From 226870df42749203a7d846bec15e8194fe1affb4 Mon Sep 17 00:00:00 2001
From: Rostislav Rumenov <rostislav.rumenov@dfinity.org>
Date: Thu, 12 Dec 2024 14:31:19 +0100
Subject: [PATCH] fix: have single place where we remove the peer from the peer
 map. (#2988)

Since we want to keep the peer map in sync with the active connection
joinmap, ideally we want to mutate both at the same place.

---------

Co-authored-by: Daniel Sharifi <40335219+DSharifi@users.noreply.github.com>
---
 .../quic_transport/src/connection_manager.rs  | 34 +++++++------------
 rs/p2p/quic_transport/src/metrics.rs          |  2 +-
 2 files changed, 13 insertions(+), 23 deletions(-)

diff --git a/rs/p2p/quic_transport/src/connection_manager.rs b/rs/p2p/quic_transport/src/connection_manager.rs
index bf0031ca976..b76aad31a10 100644
--- a/rs/p2p/quic_transport/src/connection_manager.rs
+++ b/rs/p2p/quic_transport/src/connection_manager.rs
@@ -348,8 +348,9 @@ impl ConnectionManager {
                 },
                 Some(active_result) = self.active_connections.join_next() => {
                     match active_result {
-                        Ok((_, peer_id)) => {
+                        Ok(((), peer_id)) => {
                             self.peer_map.write().unwrap().remove(&peer_id);
+                            self.metrics.peers_removed_total.inc();
                             self.connect_queue.insert(peer_id, Duration::ZERO);
                             self.metrics.peer_map_size.dec();
                             self.metrics.closed_request_handlers_total.inc();
@@ -398,21 +399,14 @@ impl ConnectionManager {
         let subnet_nodes = SomeOrAllNodes::Some(subnet_node_set);
 
         // Set new server config to only accept connections from the current set.
-        match self
-            .tls_config
+        let rustls_server_config = self.tls_config
             .server_config(subnet_nodes, self.topology.latest_registry_version())
-        {
-            Ok(rustls_server_config) => {
-                let quic_server_config = QuicServerConfig::try_from(rustls_server_config).expect("Conversion from RustTls config to Quinn config must succeed as long as this library and quinn use the same RustTls versions.");
-                let mut server_config =
-                    quinn::ServerConfig::with_crypto(Arc::new(quic_server_config));
-                server_config.transport_config(self.transport_config.clone());
-                self.endpoint.set_server_config(Some(server_config));
-            }
-            Err(e) => {
-                error!(self.log, "Failed to get certificate from crypto {:?}", e)
-            }
-        }
+            .expect("The rustls server config must be locally available, otherwise transport can't run.");
+
+        let quic_server_config = QuicServerConfig::try_from(rustls_server_config).expect("Conversion from RustTls config to Quinn config must succeed as long as this library and quinn use the same RustTls versions.");
+        let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(quic_server_config));
+        server_config.transport_config(self.transport_config.clone());
+        self.endpoint.set_server_config(Some(server_config));
 
         // Connect/Disconnect from peers according to new topology
         for (peer_id, _) in self.topology.iter() {
@@ -423,23 +417,19 @@ impl ConnectionManager {
 
         // Remove peer connections that are not part of subnet anymore.
         // Also remove peer connections that have closed connections.
-        let mut peer_map = self.peer_map.write().unwrap();
-        peer_map.retain(|peer_id, conn_handle| {
+        let peer_map = self.peer_map.read().unwrap();
+        peer_map.iter().for_each(|(peer_id, conn_handle)| {
             let peer_left_topology = !self.topology.is_member(peer_id);
             let node_left_topology = !self.topology.is_member(&self.node_id);
             // If peer is not member anymore or this node not part of subnet close connection.
             let should_close_connection = peer_left_topology || node_left_topology;
-
             if should_close_connection {
-                self.metrics.peers_removed_total.inc();
                 conn_handle
                     .conn()
                     .close(VarInt::from_u32(0), b"node not part of subnet anymore");
-                false
-            } else {
-                true
             }
         });
+
         self.metrics.peer_map_size.set(peer_map.len() as i64);
     }
 
diff --git a/rs/p2p/quic_transport/src/metrics.rs b/rs/p2p/quic_transport/src/metrics.rs
index 6678de167d0..9c3097001be 100644
--- a/rs/p2p/quic_transport/src/metrics.rs
+++ b/rs/p2p/quic_transport/src/metrics.rs
@@ -86,7 +86,7 @@ impl QuicTransportMetrics {
             ),
             peers_removed_total: metrics_registry.int_counter(
                 "quic_transport_peers_removed_total",
-                "Peers removed because they are not part of topology anymore.",
+                "Peers removed from the peer map.",
             ),
             inbound_connection_total: metrics_registry.int_counter(
                 "quic_transport_inbound_connection_total",