apollo_integration_tests: await abort handles when shutting down nodes in integration test (#14078)

Signed-off-by: Dori Medini <dori@starkware.co>

Author: dorimedini-starkware

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_positive_flow.rs

@@ -51,7 +51,7 @@ async fn main() {
         .await;
 
     info!("Shutting down nodes.");
-    integration_test_manager.shutdown_nodes(node_indices);
+    integration_test_manager.shutdown_nodes(node_indices).await;
 
     info!("Positive flow integration test completed successfully!");
 }

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_restart_flow.rs

@@ -57,7 +57,7 @@ async fn main() {
         info!("Awaiting transactions while all nodes are up");
         integration_test_manager.poll_all_running_nodes_received_more_txs(TIMEOUT).await;
 
-        integration_test_manager.shutdown_nodes([RESTART_NODE].into());
+        integration_test_manager.shutdown_nodes([RESTART_NODE].into()).await;
         info!("Awaiting transactions while node {RESTART_NODE} is down");
         integration_test_manager.poll_all_running_nodes_received_more_txs(TIMEOUT).await;
 
@@ -77,7 +77,7 @@ async fn main() {
 
         // Shutdown a second node to test that the restarted node has joined consensus (the network
         // can't reach consensus without the restarted node if the second node is down).
-        integration_test_manager.shutdown_nodes([SHUTDOWN_NODE].into());
+        integration_test_manager.shutdown_nodes([SHUTDOWN_NODE].into()).await;
         // Shutting down a node that's already down results in an error so we remove it from the set
         // here.
         node_indices.remove(&SHUTDOWN_NODE);
@@ -98,6 +98,6 @@ async fn main() {
 
     integration_test_manager.verify_block_hash_across_all_running_nodes(None).await;
 
-    integration_test_manager.shutdown_nodes(node_indices);
+    integration_test_manager.shutdown_nodes(node_indices).await;
     info!("Restart flow integration test completed successfully!");
 }

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_restart_service_multiple_nodes_flow.rs

@@ -80,7 +80,7 @@ async fn main() {
         .await;
 
     info!("Shutting down nodes.");
-    integration_test_manager.shutdown_nodes(node_indices);
+    integration_test_manager.shutdown_nodes(node_indices).await;
 
     info!("Restart service multiple nodes flow integration test completed successfully!");
 }

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_restart_service_single_node_flow.rs

@@ -89,6 +89,6 @@ async fn main() {
 
     integration_test_manager.verify_block_hash_across_all_running_nodes(None).await;
 
-    integration_test_manager.shutdown_nodes(node_indices);
+    integration_test_manager.shutdown_nodes(node_indices).await;
     info!("Restart service single node flow integration test completed successfully!");
 }

crates/apollo_integration_tests/src/bin/sequencer_node_end_to_end_integration_tests/integration_test_revert_flow.rs

@@ -69,7 +69,7 @@ async fn main() {
         .await;
 
     info!("Shutting down nodes.");
-    integration_test_manager.shutdown_nodes(node_indices.clone());
+    integration_test_manager.shutdown_nodes(node_indices.clone()).await;
 
     let expected_block_number_after_revert = REVERT_UP_TO_AND_INCLUDING.prev().unwrap_or_default();
     info!(
@@ -98,7 +98,7 @@ async fn main() {
         .await;
 
     info!("All nodes reverted to block {expected_block_number_after_revert}. Shutting down nodes.");
-    integration_test_manager.shutdown_nodes(node_indices.clone());
+    integration_test_manager.shutdown_nodes(node_indices.clone()).await;
 
     // Restore the tx generator state.
     *integration_test_manager.tx_generator_mut() = tx_generator_snapshot;
@@ -129,7 +129,7 @@ async fn main() {
         ))
         .await;
 
-    integration_test_manager.shutdown_nodes(node_indices);
+    integration_test_manager.shutdown_nodes(node_indices).await;
 
     info!("Revert flow integration test completed successfully!");
 }

crates/apollo_integration_tests/src/integration_test_manager.rs

@@ -660,23 +660,48 @@ impl IntegrationTestManager {
         });
     }
 
-    pub fn shutdown_nodes(&mut self, nodes_to_shutdown: HashSet<usize>) {
-        nodes_to_shutdown.into_iter().for_each(|index| {
+    pub async fn shutdown_nodes(&mut self, nodes_to_shutdown: HashSet<usize>) {
+        let mut all_handles: Vec<AbortOnDropHandle<()>> = Vec::new();
+
+        for index in nodes_to_shutdown {
             let running_node = self
                 .running_nodes
                 .remove(&index)
                 .unwrap_or_else(|| panic!("Node {index} is not in the running map."));
             // TODO(Tsabary): should this function call `shutdown_node_services` per node as well?
-            running_node.executable_handles.values().for_each(|handle| {
+            let RunningNode { node_setup, executable_handles } = running_node;
+            let handles: Vec<AbortOnDropHandle<()>> = executable_handles.into_values().collect();
+            for handle in &handles {
                 assert!(!handle.is_finished(), "Node {index} should still be running.");
                 handle.abort();
-            });
+            }
             assert!(
-                self.idle_nodes.insert(index, running_node.node_setup).is_none(),
+                self.idle_nodes.insert(index, node_setup).is_none(),
                 "Node {index} is already in the idle map."
             );
-            info!("Node {} has been shut down.", index);
-        });
+            info!("Node {index} has been aborted and should shut down soon.");
+            all_handles.extend(handles);
+        }
+
+        // Wait for all tasks to be cancelled. When each task is cancelled its Child process handle
+        // is dropped, which triggers kill_on_drop and sends SIGKILL to the child process.
+        let results = tokio::time::timeout(Duration::from_secs(10), join_all(all_handles))
+            .await
+            .expect("Node tasks did not cancel within timeout.");
+
+        // Assert that all tasks were cancelled. Any shutdown for any other reason should be
+        // considered a bug.
+        for result in results {
+            match result {
+                Err(e) if e.is_cancelled() => {}
+                Ok(()) => panic!("A node task exited before shutdown — node may have crashed."),
+                Err(e) => panic!("A node task failed during shutdown: {e:?}"),
+            }
+        }
+        info!("All nodes have been shut down.");
+
+        // Brief pause to let the OS release ports after SIGKILL before new processes bind them.
+        sleep(Duration::from_millis(100)).await;
     }
 
     pub async fn send_deploy_and_invoke_txs_and_verify(&mut self) {