fix(connect): add shutdown handle for initial join task

Author: sanity

crates/core/src/node/p2p_impl.rs

@@ -1,6 +1,7 @@
 use std::{collections::HashSet, convert::Infallible, sync::Arc, time::Duration};
 
 use futures::{future::BoxFuture, FutureExt};
+use tokio::task::JoinHandle;
 use tracing::Instrument;
 
 use super::{
@@ -44,6 +45,7 @@ pub(crate) struct NodeP2P {
     should_try_connect: bool,
     client_events_task: BoxFuture<'static, anyhow::Error>,
     contract_executor_task: BoxFuture<'static, anyhow::Error>,
+    initial_join_task: Option<JoinHandle<()>>,
 }
 
 impl NodeP2P {
@@ -181,10 +183,14 @@ impl NodeP2P {
 
         Ok(())
     }
-    pub(super) async fn run_node(self) -> anyhow::Result<Infallible> {
+    pub(super) async fn run_node(mut self) -> anyhow::Result<Infallible> {
         if self.should_try_connect {
-            connect::initial_join_procedure(self.op_manager.clone(), &self.conn_manager.gateways)
-                .await?;
+            let join_handle = connect::initial_join_procedure(
+                self.op_manager.clone(),
+                &self.conn_manager.gateways,
+            )
+            .await?;
+            self.initial_join_task = Some(join_handle);
 
             // After connecting to gateways, aggressively try to reach min_connections
             // This is important for fast startup and avoiding on-demand connection delays
@@ -199,7 +205,8 @@ impl NodeP2P {
             self.node_controller,
         );
 
-        tokio::select!(
+        let join_task = self.initial_join_task.take();
+        let result = tokio::select!(
             r = f => {
                let Err(e) = r;
                tracing::error!("Network event listener exited: {}", e);
@@ -213,7 +220,13 @@ impl NodeP2P {
                 tracing::error!("Contract executor task exited: {:?}", e);
                 Err(e)
             }
-        )
+        );
+
+        if let Some(handle) = join_task {
+            handle.abort();
+        }
+
+        result
     }
 
     pub(crate) async fn build<CH, const CLIENTS: usize, ER>(
@@ -343,6 +356,7 @@ impl NodeP2P {
             location: config.location,
             client_events_task,
             contract_executor_task,
+            initial_join_task: None,
         })
     }
 }

crates/core/src/node/testing_impl.rs

@@ -780,7 +780,8 @@ where
     NB: NetworkBridge + NetworkBridgeExt,
     UsrEv: ClientEventsProxy + Send + 'static,
 {
-    connect::initial_join_procedure(config.op_manager.clone(), &config.gateways).await?;
+    let join_task =
+        connect::initial_join_procedure(config.op_manager.clone(), &config.gateways).await?;
     let (client_responses, _cli_response_sender) = contract::client_responses_channel();
     let span = {
         config
@@ -811,9 +812,13 @@ where
         .parent_span
         .clone()
         .unwrap_or_else(|| tracing::info_span!("event_listener", peer = %config.peer_key));
-    run_event_listener(node_controller_rx, config)
+    let result = run_event_listener(node_controller_rx, config)
         .instrument(parent_span)
-        .await
+        .await;
+
+    join_task.abort();
+    let _ = join_task.await;
+    result
 }
 
 /// Starts listening to incoming events. Will attempt to join the ring if any gateways have been provided.

crates/core/src/operations/connect.rs

@@ -12,7 +12,7 @@ use std::time::{Duration, Instant};
 use futures::{stream::FuturesUnordered, StreamExt};
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
-use tokio::task;
+use tokio::task::{self, JoinHandle};
 
 use crate::client_events::HostResult;
 use crate::dev_tool::Location;
@@ -297,6 +297,9 @@ impl RelayContext for RelayEnv<'_> {
     }
 
     fn courtesy_hint(&self, _acceptor: &PeerKeyLocation, _joiner: &PeerKeyLocation) -> bool {
+        // Courtesy slots still piggyback on regular connections. Flag the first acceptance so the
+        // joiner can prioritise it, and keep the logic simple until dedicated courtesy tracking
+        // is wired in (see courtesy-connection-budget branch).
         self.op_manager.ring.open_connections() == 0
     }
 }
@@ -836,15 +839,15 @@ pub(crate) async fn join_ring_request(
 pub(crate) async fn initial_join_procedure(
     op_manager: Arc<OpManager>,
     gateways: &[PeerKeyLocation],
-) -> Result<(), OpError> {
+) -> Result<JoinHandle<()>, OpError> {
     let number_of_parallel_connections = {
         let max_potential_conns_per_gw = op_manager.ring.max_hops_to_live;
         let needed_to_cover_max =
             op_manager.ring.connection_manager.max_connections / max_potential_conns_per_gw;
         gateways.iter().take(needed_to_cover_max).count().max(2)
     };
     let gateways = gateways.to_vec();
-    task::spawn(async move {
+    let handle = task::spawn(async move {
         if gateways.is_empty() {
             tracing::warn!("No gateways available, aborting join procedure");
             return;
@@ -940,7 +943,7 @@ pub(crate) async fn initial_join_procedure(
             tokio::time::sleep(Duration::from_secs(wait_time)).await;
         }
     });
-    Ok(())
+    Ok(handle)
 }
 
 #[cfg(test)]