change the way of killing cancelled tasks (#874)

<!-- Reference any GitHub issues resolved by this PR -->

Closes #

## Introduced changes

<!-- A brief description of the changes -->

-
-

## Breaking changes

<!-- List of all breaking changes, if applicable -->

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [x] Updated relevant documentation
- [x] Added relevant tests
- [x] Performed self-review of the code
- [x] Added changes to `CHANGELOG.md`

(cherry picked from commit 0947313)

Author: tkumor3

crates/forge/src/lib.rs

@@ -12,7 +12,7 @@ use futures::StreamExt;
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use running::blocking_run_from_test;
 use serde::Deserialize;
-use tokio::sync::mpsc::channel;
+use tokio::sync::mpsc::{channel, Sender};
 
 use std::sync::Arc;
 use test_case_summary::TestCaseSummary;
@@ -366,7 +366,7 @@ pub async fn run(
         .await?;
 
         fuzzing_happened |= was_fuzzed;
-        summaries.push(summary.clone());
+        summaries.push(summary);
     }
 
     pretty_printing::print_test_summary(&summaries);
@@ -397,6 +397,17 @@ async fn run_tests_from_crate(
     let mut tasks = FuturesUnordered::new();
     let test_cases = &tests.test_cases;
 
+    // Initiate two channels to manage the `--exit-first` flag.
+    // Owing to `cheatnet` fork's utilization of its own Tokio runtime for RPC requests,
+    // test execution must occur within a `tokio::spawn_blocking`.
+    // As `spawn_blocking` can't be prematurely cancelled (refer: https://dtantsur.github.io/rust-openstack/tokio/task/fn.spawn_blocking.html),
+    // a channel is used to signal the task that test processing is no longer necessary.
+    let (send, mut rec) = channel(1);
+
+    // The second channel serves as a hold point to ensure all tasks complete
+    // their shutdown procedures before moving forward (more info: https://tokio.rs/tokio/topics/shutdown)
+    let (send_shut_down, mut rec_shut_down) = channel(1);
+
     for case in test_cases.iter() {
         let case_name = case.name.as_str();
 
@@ -414,6 +425,8 @@ async fn run_tests_from_crate(
             runner_config.clone(),
             runner_params.clone(),
             cancellation_tokens.clone(),
+            &send,
+            &send_shut_down,
         ));
     }
 
@@ -431,6 +444,12 @@ async fn run_tests_from_crate(
         results.push(result);
     }
 
+    rec.close();
+
+    // Waiting for things to finish shutting down
+    drop(send_shut_down);
+    let _ = rec_shut_down.recv().await;
+
     Ok((
         TestCrateSummary {
             test_case_summaries: results,
@@ -441,13 +460,16 @@ async fn run_tests_from_crate(
     ))
 }
 
+#[allow(clippy::too_many_arguments)]
 fn choose_test_strategy_and_run(
     args: Vec<ConcreteTypeId>,
     case: Arc<TestCase>,
     runner: Arc<SierraCasmRunner>,
     runner_config: Arc<RunnerConfig>,
     runner_params: Arc<RunnerParams>,
     cancellation_tokens: Arc<CancellationTokens>,
+    send: &Sender<()>,
+    send_shut_down: &Sender<()>,
 ) -> JoinHandle<Result<TestCaseSummary>> {
     if args.is_empty() {
         run_single_test(
@@ -456,6 +478,8 @@ fn choose_test_strategy_and_run(
             runner_config,
             runner_params,
             cancellation_tokens,
+            send.clone(),
+            send_shut_down.clone(),
         )
     } else {
         run_with_fuzzing(
@@ -465,6 +489,7 @@ fn choose_test_strategy_and_run(
             runner_config,
             runner_params,
             cancellation_tokens,
+            send_shut_down.clone(),
         )
     }
 }
@@ -475,6 +500,8 @@ fn run_single_test(
     runner_config: Arc<RunnerConfig>,
     runner_params: Arc<RunnerParams>,
     cancellation_tokens: Arc<CancellationTokens>,
+    send: Sender<()>,
+    send_shut_down: Sender<()>,
 ) -> JoinHandle<Result<TestCaseSummary>> {
     let exit_first = runner_config.exit_first;
     tokio::task::spawn(async move {
@@ -489,8 +516,8 @@ fn run_single_test(
                 // one of a test returns Err
                 Ok(TestCaseSummary::Interrupted {  })
             },
-            result = blocking_run_from_test(vec![], case.clone(),runner,  runner_config.clone(), runner_params.clone() , None) => {
-                match result {
+            result = blocking_run_from_test(vec![], case.clone(),runner,  runner_config.clone(), runner_params.clone(), send.clone(), send_shut_down.clone() ) => {
+                match result? {
                     Ok(result) => {
                         if exit_first {
                             if let TestCaseSummary::Failed { .. } = &result {
@@ -516,10 +543,11 @@ fn run_with_fuzzing(
     runner_config: Arc<RunnerConfig>,
     runner_params: Arc<RunnerParams>,
     cancellation_tokens: Arc<CancellationTokens>,
+    send_shut_down: Sender<()>,
 ) -> JoinHandle<Result<TestCaseSummary>> {
     tokio::task::spawn(async move {
         let token = CancellationToken::new();
-
+        let (send, mut rec) = channel(1);
         let args = args
             .iter()
             .map(|arg| {
@@ -540,9 +568,6 @@ fn run_with_fuzzing(
         let mut fuzzer = RandomFuzzer::create(fuzzer_seed, fuzzer_runs, &args)?;
 
         let mut tasks = FuturesUnordered::new();
-        // Pattern in order to waiting for things to finish shutting down
-        // https://tokio.rs/tokio/topics/shutdown
-        let (send, mut recv) = channel(1);
 
         for _ in 1..=fuzzer_runs {
             let args = fuzzer.next_args();
@@ -556,14 +581,12 @@ fn run_with_fuzzing(
                 cancellation_tokens.clone(),
                 token.clone(),
                 send.clone(),
+                send_shut_down.clone(),
             ));
         }
 
         let mut results = vec![];
 
-        // Graceful Shutdown Pattern
-        drop(send);
-
         while let Some(task) = tasks.next().await {
             let result = task??;
             results.push(result.clone());
@@ -576,7 +599,7 @@ fn run_with_fuzzing(
             }
         }
 
-        let _ = recv.recv().await;
+        rec.close();
 
         let runs = u32::try_from(
             results
@@ -627,7 +650,8 @@ fn run_fuzzing_subtest(
     runner_params: Arc<RunnerParams>,
     cancellation_tokens: Arc<CancellationTokens>,
     cancellation_fuzzing_token: CancellationToken,
-    send: tokio::sync::mpsc::Sender<()>,
+    send: Sender<()>,
+    send_shut_down: Sender<()>,
 ) -> JoinHandle<Result<TestCaseSummary>> {
     let c = case.clone();
     task::spawn(async move {
@@ -654,9 +678,10 @@ fn run_fuzzing_subtest(
                 runner,
                 runner_config.clone(),
                 runner_params.clone(),
-                Some(send),
+                send.clone(),
+                send_shut_down.clone()
             ) => {
-                match result {
+                match result? {
                     Ok(result) => {
                         if let TestCaseSummary::Failed { .. } = &result {
                             if runner_config.exit_first {

crates/forge/src/main.rs

@@ -123,6 +123,7 @@ fn main_execution() -> Result<bool> {
         .max_blocking_threads(cores_approx)
         .enable_all()
         .build()?;
+
     let all_failed_tests = rt.block_on({
         rt.spawn(async move {
             let mut all_failed_tests = vec![];

crates/forge/src/running.rs

@@ -35,6 +35,7 @@ use starknet_api::patricia_key;
 use starknet_api::transaction::Calldata;
 use test_collector::{ForkConfig, TestCase};
 use tokio::sync::mpsc::Sender;
+use tokio::task::JoinHandle;
 
 use crate::scarb::ForkTarget;
 use crate::test_case_summary::TestCaseSummary;
@@ -70,25 +71,31 @@ fn build_hints_dict<'b>(
     (hints_dict, string_to_hint)
 }
 
-pub(crate) async fn blocking_run_from_test(
+pub(crate) fn blocking_run_from_test(
     args: Vec<Felt252>,
     case: Arc<TestCase>,
     runner: Arc<SierraCasmRunner>,
     runner_config: Arc<RunnerConfig>,
     runner_params: Arc<RunnerParams>,
-    sender: Option<Sender<()>>,
-) -> Result<TestCaseSummary> {
+    send: Sender<()>,
+    send_shut_down: Sender<()>,
+) -> JoinHandle<Result<TestCaseSummary>> {
     tokio::task::spawn_blocking(move || {
+        // Due to the inability of spawn_blocking to be abruptly cancelled,
+        // a channel is used to receive information indicating
+        // that the execution of the task is no longer necessary.
+        if send.is_closed() {
+            return Err(anyhow::anyhow!("stop spawn_blocking"));
+        }
         run_test_case(
             args,
             &case,
             &runner,
             &runner_config,
             &runner_params,
-            &sender,
+            &send_shut_down,
         )
     })
-    .await?
 }
 
 fn build_context() -> EntryPointExecutionContext {
@@ -146,7 +153,7 @@ pub(crate) fn run_test_case(
     runner: &SierraCasmRunner,
     runner_config: &Arc<RunnerConfig>,
     runner_params: &Arc<RunnerParams>,
-    _sender: &Option<Sender<()>>,
+    _send_shut_down: &Sender<()>,
 ) -> Result<TestCaseSummary> {
     let available_gas = if let Some(available_gas) = &case.available_gas {
         Some(*available_gas)

crates/forge/src/test_crate_summary.rs

@@ -2,7 +2,7 @@ use crate::test_case_summary::TestCaseSummary;
 use crate::{RunnerStatus, TestCrateType};
 
 /// Summary of the test run in the file
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq)]
 pub struct TestCrateSummary {
     /// Summaries of each test case in the file
     pub test_case_summaries: Vec<TestCaseSummary>,

crates/forge/tests/integration/common/running_tests.rs

@@ -12,6 +12,7 @@ use tokio::runtime::Runtime;
 
 pub fn run_test_case(test: &TestCase) -> Vec<TestCrateSummary> {
     let rt = Runtime::new().expect("Could not instantiate Runtime");
+
     rt.block_on(run(
         &test.path().unwrap(),
         &String::from("src"),

crates/forge/tests/integration/setup_fork.rs

@@ -94,6 +94,7 @@ fn fork_aliased_decorator() {
     ).as_str());
 
     let rt = Runtime::new().expect("Could not instantiate Runtime");
+
     let result = rt
         .block_on(run(
             &test.path().unwrap(),