fix(swarm): eliminating protocol cloning when nothing is happening

swarm/src/connection.rs

@@ -32,7 +32,7 @@ pub use supported_protocols::SupportedProtocols;
 use crate::handler::{
     AddressChange, ConnectionEvent, ConnectionHandler, DialUpgradeError, FullyNegotiatedInbound,
     FullyNegotiatedOutbound, ListenUpgradeError, ProtocolSupport, ProtocolsAdded, ProtocolsChange,
-    UpgradeInfoSend,
+    ProtocolsRemoved, UpgradeInfoSend,
 };
 use crate::stream::ActiveStreamCounter;
 use crate::upgrade::{InboundUpgradeSend, OutboundUpgradeSend};
@@ -52,7 +52,7 @@ use libp2p_core::upgrade;
 use libp2p_core::upgrade::{NegotiationError, ProtocolError};
 use libp2p_core::Endpoint;
 use libp2p_identity::PeerId;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 use std::future::Future;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -153,8 +153,11 @@ where
         SubstreamRequested<THandler::OutboundOpenInfo, THandler::OutboundProtocol>,
     >,
 
-    local_supported_protocols: HashSet<StreamProtocol>,
+    local_supported_protocols:
+        HashMap<AsStrHashEq<<THandler::InboundProtocol as UpgradeInfoSend>::Info>, bool>,
     remote_supported_protocols: HashSet<StreamProtocol>,
+    temp_protocols: Vec<StreamProtocol>,
+
     idle_timeout: Duration,
     stream_counter: ActiveStreamCounter,
 }
@@ -186,10 +189,22 @@ where
         max_negotiating_inbound_streams: usize,
         idle_timeout: Duration,
     ) -> Self {
-        let initial_protocols = gather_supported_protocols(&handler);
-        if !initial_protocols.is_empty() {
+        let local_supported_protocols = handler
+            .listen_protocol()
+            .upgrade()
+            .protocol_info()
+            .map(|i| (AsStrHashEq(i), true))
+            .collect::<HashMap<_, _>>();
+
+        if !local_supported_protocols.is_empty() {
+            let temp = local_supported_protocols
+                .keys()
+                .filter_map(|i| StreamProtocol::try_from_owned(i.0.as_ref().to_owned()).ok())
+                .collect::<Vec<_>>();
             handler.on_connection_event(ConnectionEvent::LocalProtocolsChange(
-                ProtocolsChange::Added(ProtocolsAdded::from_set(&initial_protocols)),
+                ProtocolsChange::Added(ProtocolsAdded {
+                    protocols: temp.iter(),
+                }),
             ));
         }
         Connection {
@@ -201,8 +216,9 @@ where
             substream_upgrade_protocol_override,
             max_negotiating_inbound_streams,
             requested_substreams: Default::default(),
-            local_supported_protocols: initial_protocols,
+            local_supported_protocols,
             remote_supported_protocols: Default::default(),
+            temp_protocols: Default::default(),
             idle_timeout,
             stream_counter: ActiveStreamCounter::default(),
         }
@@ -250,6 +266,7 @@ where
             substream_upgrade_protocol_override,
             local_supported_protocols: supported_protocols,
             remote_supported_protocols,
+            temp_protocols,
             idle_timeout,
             stream_counter,
             ..
@@ -286,24 +303,35 @@ where
                 Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(
                     ProtocolSupport::Added(protocols),
                 )) => {
-                    if let Some(added) =
-                        ProtocolsChange::add(remote_supported_protocols, &protocols)
-                    {
-                        handler.on_connection_event(ConnectionEvent::RemoteProtocolsChange(added));
-                        remote_supported_protocols.extend(protocols);
+                    let added = protocols
+                        .into_iter()
+                        .filter(|i| remote_supported_protocols.insert(i.clone()))
+                        .collect::<Vec<_>>();
+
+                    if !added.is_empty() {
+                        handler.on_connection_event(ConnectionEvent::RemoteProtocolsChange(
+                            ProtocolsChange::Added(ProtocolsAdded {
+                                protocols: added.iter(),
+                            }),
+                        ));
                     }
 
                     continue;
                 }
                 Poll::Ready(ConnectionHandlerEvent::ReportRemoteProtocols(
                     ProtocolSupport::Removed(protocols),
                 )) => {
-                    if let Some(removed) =
-                        ProtocolsChange::remove(remote_supported_protocols, &protocols)
-                    {
-                        handler
-                            .on_connection_event(ConnectionEvent::RemoteProtocolsChange(removed));
-                        remote_supported_protocols.retain(|p| !protocols.contains(p));
+                    let removed = protocols
+                        .into_iter()
+                        .filter_map(|i| remote_supported_protocols.take(&i))
+                        .collect::<Vec<_>>();
+
+                    if !removed.is_empty() {
+                        handler.on_connection_event(ConnectionEvent::RemoteProtocolsChange(
+                            ProtocolsChange::Removed(ProtocolsRemoved {
+                                protocols: removed.iter(),
+                            }),
+                        ));
                     }
 
                     continue;
@@ -431,17 +459,18 @@ where
                 }
             }
 
-            let new_protocols = gather_supported_protocols(handler);
-            let changes = ProtocolsChange::from_full_sets(supported_protocols, &new_protocols);
-
-            if !changes.is_empty() {
-                for change in changes {
-                    handler.on_connection_event(ConnectionEvent::LocalProtocolsChange(change));
-                }
-
-                *supported_protocols = new_protocols;
-
-                continue; // Go back to the top, handler can potentially make progress again.
+            let changes = ProtocolsChange::from_full_sets(
+                supported_protocols,
+                handler.listen_protocol().upgrade().protocol_info(),
+                temp_protocols,
+            );
+            let mut has_changes = false;
+            for change in changes {
+                handler.on_connection_event(ConnectionEvent::LocalProtocolsChange(change));
+                has_changes = true;
+            }
+            if has_changes {
+                continue;
             }
 
             return Poll::Pending; // Nothing can make progress, return `Pending`.
@@ -454,15 +483,6 @@ where
     }
 }
 
-fn gather_supported_protocols(handler: &impl ConnectionHandler) -> HashSet<StreamProtocol> {
-    handler
-        .listen_protocol()
-        .upgrade()
-        .protocol_info()
-        .filter_map(|i| StreamProtocol::try_from_owned(i.as_ref().to_owned()).ok())
-        .collect()
-}
-
 fn compute_new_shutdown(
     handler_keep_alive: bool,
     current_shutdown: &Shutdown,
@@ -737,6 +757,58 @@ enum Shutdown {
     Later(Delay, Instant),
 }
 
+/// The endpoint roles associated with a pending peer-to-peer connection.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+enum PendingPoint {
+    /// The socket comes from a dialer.
+    ///
+    /// There is no single address associated with the Dialer of a pending
+    /// connection. Addresses are dialed in parallel. Only once the first dial
+    /// is successful is the address of the connection known.
+    Dialer {
+        /// Same as [`ConnectedPoint::Dialer`] `role_override`.
+        role_override: Endpoint,
+    },
+    /// The socket comes from a listener.
+    Listener {
+        /// Local connection address.
+        local_addr: Multiaddr,
+        /// Address used to send back data to the remote.
+        send_back_addr: Multiaddr,
+    },
+}
+
+impl From<ConnectedPoint> for PendingPoint {
+    fn from(endpoint: ConnectedPoint) -> Self {
+        match endpoint {
+            ConnectedPoint::Dialer { role_override, .. } => PendingPoint::Dialer { role_override },
+            ConnectedPoint::Listener {
+                local_addr,
+                send_back_addr,
+            } => PendingPoint::Listener {
+                local_addr,
+                send_back_addr,
+            },
+        }
+    }
+}
+
+pub(crate) struct AsStrHashEq<T>(pub(crate) T);
+
+impl<T: AsRef<str>> Eq for AsStrHashEq<T> {}
+
+impl<T: AsRef<str>> PartialEq for AsStrHashEq<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.as_ref() == other.0.as_ref()
+    }
+}
+
+impl<T: AsRef<str>> std::hash::Hash for AsStrHashEq<T> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.0.as_ref().hash(state)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -848,6 +920,51 @@ mod tests {
         assert_eq!(connection.handler.local_removed, vec![vec!["/foo"]]);
     }
 
+    #[test]
+    #[ignore]
+    fn repoll_with_active_protocols() {
+        fn run_benchmark(protcol_count: usize, iters: usize) {
+            // we need longer protocol names
+            macro_rules! protocols {
+                ($($name:ident)*) => {
+                    concat!(
+                        $(
+                            "/",
+                            stringify!($name),
+                            "ffffffffffffffffffffff ",
+                        )*
+                    )
+                }
+            }
+
+            const PROTOCOLS: &str = protocols!(a b c d e f g h i j k l m n o p q r s t u v);
+            let protocols = PROTOCOLS.split(' ').collect::<Vec<_>>();
+
+            let mut connection = Connection::new(
+                StreamMuxerBox::new(PendingStreamMuxer),
+                ConfigurableProtocolConnectionHandler::default(),
+                None,
+                0,
+                Duration::ZERO,
+            );
+
+            connection.handler.listen_on(&protocols[..protcol_count]);
+
+            let now = Instant::now();
+            for _ in 0..iters {
+                let _ = connection.poll_noop_waker();
+            }
+            let elapsed = now.elapsed();
+            println!("{protcol_count} {elapsed:?}");
+        }
+
+        let iters = 3_000_000;
+        run_benchmark(2, iters);
+        run_benchmark(4, iters);
+        run_benchmark(10, iters);
+        run_benchmark(20, iters);
+    }
+
     #[test]
     fn only_propagtes_actual_changes_to_remote_protocols_to_handler() {
         let mut connection = Connection::new(
@@ -1315,39 +1432,3 @@ mod tests {
         }
     }
 }
-
-/// The endpoint roles associated with a pending peer-to-peer connection.
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-enum PendingPoint {
-    /// The socket comes from a dialer.
-    ///
-    /// There is no single address associated with the Dialer of a pending
-    /// connection. Addresses are dialed in parallel. Only once the first dial
-    /// is successful is the address of the connection known.
-    Dialer {
-        /// Same as [`ConnectedPoint::Dialer`] `role_override`.
-        role_override: Endpoint,
-    },
-    /// The socket comes from a listener.
-    Listener {
-        /// Local connection address.
-        local_addr: Multiaddr,
-        /// Address used to send back data to the remote.
-        send_back_addr: Multiaddr,
-    },
-}
-
-impl From<ConnectedPoint> for PendingPoint {
-    fn from(endpoint: ConnectedPoint) -> Self {
-        match endpoint {
-            ConnectedPoint::Dialer { role_override, .. } => PendingPoint::Dialer { role_override },
-            ConnectedPoint::Listener {
-                local_addr,
-                send_back_addr,
-            } => PendingPoint::Listener {
-                local_addr,
-                send_back_addr,
-            },
-        }
-    }
-}

swarm/src/connection/supported_protocols.rs

@@ -40,7 +40,6 @@ impl SupportedProtocols {
 mod tests {
     use super::*;
     use crate::handler::{ProtocolsAdded, ProtocolsRemoved};
-    use once_cell::sync::Lazy;
 
     #[test]
     fn protocols_change_added_returns_correct_changed_value() {
@@ -70,19 +69,24 @@ mod tests {
     }
 
     fn add_foo() -> ProtocolsChange<'static> {
-        ProtocolsChange::Added(ProtocolsAdded::from_set(&FOO_PROTOCOLS))
+        ProtocolsChange::Added(ProtocolsAdded {
+            protocols: FOO_PROTOCOLS.iter(),
+        })
     }
 
     fn add_foo_bar() -> ProtocolsChange<'static> {
-        ProtocolsChange::Added(ProtocolsAdded::from_set(&FOO_BAR_PROTOCOLS))
+        ProtocolsChange::Added(ProtocolsAdded {
+            protocols: FOO_BAR_PROTOCOLS.iter(),
+        })
     }
 
     fn remove_foo() -> ProtocolsChange<'static> {
-        ProtocolsChange::Removed(ProtocolsRemoved::from_set(&FOO_PROTOCOLS))
+        ProtocolsChange::Removed(ProtocolsRemoved {
+            protocols: FOO_PROTOCOLS.iter(),
+        })
     }
 
-    static FOO_PROTOCOLS: Lazy<HashSet<StreamProtocol>> =
-        Lazy::new(|| HashSet::from([StreamProtocol::new("/foo")]));
-    static FOO_BAR_PROTOCOLS: Lazy<HashSet<StreamProtocol>> =
-        Lazy::new(|| HashSet::from([StreamProtocol::new("/foo"), StreamProtocol::new("/bar")]));
+    static FOO_PROTOCOLS: &[StreamProtocol] = &[StreamProtocol::new("/foo")];
+    static FOO_BAR_PROTOCOLS: &[StreamProtocol] =
+        &[StreamProtocol::new("/foo"), StreamProtocol::new("/bar")];
 }