chore: improve the creation of the p2p handlers and managers (#3467)

Author: rumenov

rs/p2p/consensus_manager/src/lib.rs

@@ -28,6 +28,22 @@ mod sender;
 type StartConsensusManagerFn =
     Box<dyn FnOnce(Arc<dyn Transport>, watch::Receiver<SubnetTopology>) -> Vec<Shutdown>>;
 
+pub struct AbortableBroadcastChannelManager(Vec<StartConsensusManagerFn>);
+
+impl AbortableBroadcastChannelManager {
+    pub fn start(
+        self,
+        transport: Arc<dyn Transport>,
+        topology_watcher: watch::Receiver<SubnetTopology>,
+    ) -> Vec<Shutdown> {
+        let mut ret = vec![];
+        for client in self.0 {
+            ret.append(&mut client(transport.clone(), topology_watcher.clone()));
+        }
+        ret
+    }
+}
+
 /// Same order of magnitude as the number of active artifacts.
 const MAX_OUTBOUND_CHANNEL_SIZE: usize = 100_000;
 
@@ -113,20 +129,11 @@ impl AbortableBroadcastChannelBuilder {
         (outbound_tx, inbound_rx)
     }
 
-    pub fn router(&mut self) -> Router {
-        self.router.take().unwrap_or_default()
-    }
-
-    pub fn run(
-        self,
-        transport: Arc<dyn Transport>,
-        topology_watcher: watch::Receiver<SubnetTopology>,
-    ) -> Vec<Shutdown> {
-        let mut ret = vec![];
-        for client in self.clients {
-            ret.append(&mut client(transport.clone(), topology_watcher.clone()));
-        }
-        ret
+    pub fn build(self) -> (Router, AbortableBroadcastChannelManager) {
+        (
+            self.router.unwrap(),
+            AbortableBroadcastChannelManager(self.clients),
+        )
     }
 }
 

rs/p2p/consensus_manager/tests/test.rs

@@ -345,17 +345,18 @@ fn load_test(
     for i in 0..num_peers {
         let node = node_test_id(i);
         let processor = TestConsensus::new(log.clone(), node, 256 * (i as usize + 1), i % 2 == 0);
-        let (jh, mut cm) =
+        let (jh, cm) =
             start_consensus_manager(no_op_logger(), rt.handle().clone(), processor.clone());
         jhs.push(jh);
-        nodes.push((node, cm.router()));
-        cms.push((node, cm));
+        let (r, m) = cm.build();
+        nodes.push((node, r));
+        cms.push((node, m));
         node_advert_map.insert(node, processor);
     }
     let (nodes, topology_watcher) = fully_connected_localhost_subnet(rt.handle(), log, id, nodes);
     for ((node1, transport), (node2, cm)) in nodes.into_iter().zip(cms.into_iter()) {
         assert!(node1 == node2);
-        cm.run(transport, topology_watcher.clone());
+        cm.start(transport, topology_watcher.clone());
     }
 
     rt.block_on(async move {

rs/p2p/state_sync_manager/src/lib.rs

@@ -47,25 +47,23 @@ const ADVERT_BROADCAST_INTERVAL: Duration = Duration::from_secs(5);
 const ADVERT_BROADCAST_TIMEOUT: Duration =
     ADVERT_BROADCAST_INTERVAL.saturating_sub(Duration::from_secs(2));
 
-pub fn build_axum_router<T: 'static>(
-    state_sync: Arc<dyn StateSyncClient<Message = T>>,
-    log: ReplicaLogger,
+pub fn build_state_sync_manager<T: Send + 'static>(
+    log: &ReplicaLogger,
     metrics_registry: &MetricsRegistry,
-) -> (
-    Router,
-    tokio::sync::mpsc::Receiver<(StateSyncArtifactId, NodeId)>,
-) {
+    rt_handle: &tokio::runtime::Handle,
+    state_sync: Arc<dyn StateSyncClient<Message = T>>,
+) -> (Router, StateSyncManager<T>) {
     let metrics = StateSyncManagerHandlerMetrics::new(metrics_registry);
     let shared_chunk_state = Arc::new(StateSyncChunkHandler::new(
         log.clone(),
-        state_sync,
+        state_sync.clone(),
         metrics.clone(),
     ));
 
-    let (tx, rx) = tokio::sync::mpsc::channel(20);
-    let advert_handler_state = Arc::new(StateSyncAdvertHandler::new(log, tx));
+    let (advert_sender, advert_receiver) = tokio::sync::mpsc::channel(20);
+    let advert_handler_state = Arc::new(StateSyncAdvertHandler::new(log.clone(), advert_sender));
 
-    let app = Router::new()
+    let router = Router::new()
         .route(STATE_SYNC_CHUNK_PATH, any(state_sync_chunk_handler))
         .with_state(shared_chunk_state)
         .route(
@@ -74,45 +72,37 @@ pub fn build_axum_router<T: 'static>(
         )
         .with_state(advert_handler_state);
 
-    (app, rx)
-}
-
-pub fn start_state_sync_manager<T: Send + 'static>(
-    log: &ReplicaLogger,
-    metrics: &MetricsRegistry,
-    rt: &Handle,
-    transport: Arc<dyn Transport>,
-    state_sync: Arc<dyn StateSyncClient<Message = T>>,
-    advert_receiver: tokio::sync::mpsc::Receiver<(StateSyncArtifactId, NodeId)>,
-) -> Shutdown {
-    let state_sync_manager_metrics = StateSyncManagerMetrics::new(metrics);
+    let state_sync_manager_metrics = StateSyncManagerMetrics::new(metrics_registry);
     let manager = StateSyncManager {
         log: log.clone(),
-        rt: rt.clone(),
+        rt: rt_handle.clone(),
         metrics: state_sync_manager_metrics,
-        transport,
         state_sync,
         advert_receiver,
         ongoing_state_sync: None,
     };
-    Shutdown::spawn_on_with_cancellation(
-        |cancellation: CancellationToken| manager.run(cancellation),
-        rt,
-    )
+    (router, manager)
 }
 
-struct StateSyncManager<T> {
+pub struct StateSyncManager<T> {
     log: ReplicaLogger,
     rt: Handle,
     metrics: StateSyncManagerMetrics,
-    transport: Arc<dyn Transport>,
     state_sync: Arc<dyn StateSyncClient<Message = T>>,
     advert_receiver: tokio::sync::mpsc::Receiver<(StateSyncArtifactId, NodeId)>,
     ongoing_state_sync: Option<OngoingStateSyncHandle>,
 }
 
 impl<T: 'static + Send> StateSyncManager<T> {
-    async fn run(mut self, cancellation: CancellationToken) {
+    pub fn start(self, transport: Arc<dyn Transport>) -> Shutdown {
+        let rt_handle = self.rt.clone();
+        Shutdown::spawn_on_with_cancellation(
+            |cancellation: CancellationToken| self.run(cancellation, transport),
+            &rt_handle,
+        )
+    }
+
+    async fn run(mut self, cancellation: CancellationToken, transport: Arc<dyn Transport>) {
         let mut interval = tokio::time::interval(ADVERT_BROADCAST_INTERVAL);
         interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
         let mut advertise_task = JoinSet::new();
@@ -127,15 +117,15 @@ impl<T: 'static + Send> StateSyncManager<T> {
                         Self::send_state_adverts(
                             self.rt.clone(),
                             self.state_sync.clone(),
-                            self.transport.clone(),
+                            transport.clone(),
                             self.metrics.clone(),
                             cancellation.clone(),
                         ),
                         &self.rt
                     );
                 },
                 Some((advert, peer_id)) = self.advert_receiver.recv() =>{
-                    self.handle_advert(advert, peer_id).await;
+                    self.handle_advert(advert, peer_id, transport.clone()).await;
                 }
                 Some(_) = advertise_task.join_next() => {}
             }
@@ -146,7 +136,12 @@ impl<T: 'static + Send> StateSyncManager<T> {
         }
     }
 
-    async fn handle_advert(&mut self, artifact_id: StateSyncArtifactId, peer_id: NodeId) {
+    async fn handle_advert(
+        &mut self,
+        artifact_id: StateSyncArtifactId,
+        peer_id: NodeId,
+        transport: Arc<dyn Transport>,
+    ) {
         self.metrics.adverts_received_total.inc();
         // Remove ongoing state sync if finished or try to add peer if ongoing.
         if let Some(ongoing) = &mut self.ongoing_state_sync {
@@ -184,7 +179,7 @@ impl<T: 'static + Send> StateSyncManager<T> {
                 self.metrics.ongoing_state_sync_metrics.clone(),
                 Arc::new(Mutex::new(chunkable)),
                 artifact_id.clone(),
-                self.transport.clone(),
+                transport,
             );
             // Add peer that initiated this state sync to ongoing state sync.
             ongoing
@@ -330,14 +325,17 @@ mod tests {
             };
 
             let (handler_tx, handler_rx) = tokio::sync::mpsc::channel(100);
-            start_state_sync_manager(
-                &log,
-                &MetricsRegistry::default(),
-                rt.handle(),
-                Arc::new(t) as Arc<_>,
-                Arc::new(s) as Arc<_>,
-                handler_rx,
-            );
+            let metrics = StateSyncManagerMetrics::new(&MetricsRegistry::default());
+
+            let manager = StateSyncManager {
+                log: log.clone(),
+                advert_receiver: handler_rx,
+                ongoing_state_sync: None,
+                metrics,
+                state_sync: Arc::new(s) as Arc<_>,
+                rt: rt.handle().clone(),
+            };
+            let _ = manager.start(Arc::new(t) as Arc<_>);
             rt.block_on(async move {
                 handler_tx.send((id, NODE_1)).await.unwrap();
                 handler_tx.send((old_id, NODE_2)).await.unwrap();

rs/p2p/state_sync_manager/tests/common.rs

@@ -418,25 +418,18 @@ pub fn create_node(
         disconnected: Arc::new(AtomicBool::new(false)),
     });
 
-    let (router, rx) = ic_state_sync_manager::build_axum_router(
-        state_sync.clone(),
-        log.clone(),
+    let (router, manager) = ic_state_sync_manager::build_state_sync_manager(
+        &log,
         &MetricsRegistry::default(),
+        rt,
+        state_sync.clone(),
     );
     let transport = transport_router.add_peer(
         NodeId::from(PrincipalId::new_node_test_id(node_num)),
         router,
         link.0,
         link.1,
     );
-    let shutdown = ic_state_sync_manager::start_state_sync_manager(
-        &log,
-        &MetricsRegistry::default(),
-        rt,
-        Arc::new(transport),
-        state_sync.clone(),
-        rx,
-    );
-
+    let shutdown = manager.start(Arc::new(transport));
     (state_sync, shutdown)
 }

rs/p2p/test_utils/src/turmoil.rs

@@ -358,19 +358,21 @@ pub fn add_transport_to_sim<F>(
             let this_ip = turmoil::lookup(peer.to_string());
             let custom_udp = CustomUdp::new(this_ip, udp_listener);
 
-            let state_sync_rx = if let Some(ref state_sync) = state_sync_client_clone {
-                let (state_sync_router, state_sync_rx) = ic_state_sync_manager::build_axum_router(
-                    state_sync.clone(),
-                    log.clone(),
-                    &MetricsRegistry::default(),
-                );
+            let state_sync_manager = if let Some(ref state_sync) = state_sync_client_clone {
+                let (state_sync_router, state_sync_manager) =
+                    ic_state_sync_manager::build_state_sync_manager(
+                        &log,
+                        &MetricsRegistry::default(),
+                        &tokio::runtime::Handle::current(),
+                        state_sync.clone(),
+                    );
                 router = Some(router.unwrap_or_default().merge(state_sync_router));
-                Some(state_sync_rx)
+                Some(state_sync_manager)
             } else {
                 None
             };
 
-            let _artifact_processor_jh = if let Some(consensus) = consensus_manager_clone {
+            let con = if let Some(consensus) = consensus_manager_clone {
                 let bouncer_factory = Arc::new(consensus.clone().read().unwrap().clone());
                 let downloader = FetchArtifact::new(
                     log.clone(),
@@ -388,9 +390,11 @@ pub fn add_transport_to_sim<F>(
                     consensus.clone(),
                     consensus.clone().read().unwrap().clone(),
                 );
-                router = Some(router.unwrap_or_default().merge(consensus_builder.router()));
 
-                Some(artifact_processor_jh)
+                let (consensus_router, manager) = consensus_builder.build();
+                router = Some(router.unwrap_or_default().merge(consensus_router));
+
+                Some((artifact_processor_jh, manager))
             } else {
                 None
             };
@@ -407,17 +411,12 @@ pub fn add_transport_to_sim<F>(
                 router.unwrap_or_default(),
             ));
 
-            consensus_builder.run(transport.clone(), topology_watcher_clone.clone());
+            if let Some((_, con_manager)) = con {
+                con_manager.start(transport.clone(), topology_watcher_clone.clone());
+            }
 
-            if let Some(state_sync_rx) = state_sync_rx {
-                ic_state_sync_manager::start_state_sync_manager(
-                    &log,
-                    &MetricsRegistry::default(),
-                    &tokio::runtime::Handle::current(),
-                    transport.clone(),
-                    state_sync_client_clone.unwrap().clone(),
-                    state_sync_rx,
-                );
+            if let Some(state_sync_manager) = state_sync_manager {
+                state_sync_manager.start(transport.clone());
             }
 
             post_setup_future_clone(peer, transport).await;

rs/replica/setup_ic_network/src/lib.rs

@@ -131,7 +131,7 @@ pub fn setup_consensus_and_p2p(
 ) {
     let consensus_pool_cache = consensus_pool.read().unwrap().get_cache();
 
-    let (ingress_pool, ingress_sender, join_handles, mut p2p_consensus) = start_consensus(
+    let (ingress_pool, ingress_sender, join_handles, p2p_consensus) = start_consensus(
         log,
         metrics_registry,
         rt_handle,
@@ -158,18 +158,21 @@ pub fn setup_consensus_and_p2p(
         max_certified_height_tx,
     );
 
-    // StateSync receive side => handler definition
-    let (state_sync_router, state_sync_manager_rx) = ic_state_sync_manager::build_axum_router(
-        state_sync_client.clone(),
-        log.clone(),
-        metrics_registry,
-    );
+    // StateSync receive side + handler definition
+    let (state_sync_manager_router, state_sync_manager_runner) =
+        ic_state_sync_manager::build_state_sync_manager(
+            log,
+            metrics_registry,
+            rt_handle,
+            state_sync_client.clone(),
+        );
 
-    // Consensus receive side => handler definition
+    // Consensus receive side + handler definition
+    let (consensus_manager_router, consensus_manager_runner) = p2p_consensus.build();
 
     // Merge all receive side handlers => router
-    let p2p_router = state_sync_router
-        .merge(p2p_consensus.router())
+    let p2p_router = state_sync_manager_router
+        .merge(consensus_manager_router)
         .layer(TraceLayer::new_for_http());
     // Quic transport
     let (_, topology_watcher) = ic_peer_manager::start_peer_manager(
@@ -200,17 +203,8 @@ pub fn setup_consensus_and_p2p(
     ));
 
     // Start the main event loops for StateSync and Consensus
-
-    let _state_sync_manager = ic_state_sync_manager::start_state_sync_manager(
-        log,
-        metrics_registry,
-        rt_handle,
-        quic_transport.clone(),
-        state_sync_client,
-        state_sync_manager_rx,
-    );
-
-    let _cancellation_token = p2p_consensus.run(quic_transport, topology_watcher);
+    let _state_sync_manager = state_sync_manager_runner.start(quic_transport.clone());
+    let _cancellation_token = consensus_manager_runner.start(quic_transport, topology_watcher);
 
     (ingress_pool, ingress_sender, join_handles)
 }