feat(swarm): improve PeerAddresses configurability

protocols/identify/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 - Add `hide_listen_addrs` option to prevent leaking (local) listen addresses.
   See [PR 5507](https://github.com/libp2p/rust-libp2p/pull/5507).
+- Add `with_cache_config` to fully configure the identify cache using `PeerAddressesConfig`.
+  Deprecating `with_cache_size` in favor of `with_cache_config`.
+  See [PR 5574](https://github.com/libp2p/rust-libp2p/pull/5574).
 
 ## 0.45.0
 

protocols/identify/src/behaviour.rs

@@ -28,8 +28,8 @@ use libp2p_identity::PublicKey;
 use libp2p_swarm::behaviour::{ConnectionClosed, ConnectionEstablished, DialFailure, FromSwarm};
 use libp2p_swarm::{
     ConnectionDenied, DialError, ExternalAddresses, ListenAddresses, NetworkBehaviour,
-    NotifyHandler, PeerAddresses, StreamUpgradeError, THandlerInEvent, ToSwarm,
-    _address_translation,
+    NotifyHandler, PeerAddresses, PeerAddressesConfig, StreamUpgradeError, THandlerInEvent,
+    ToSwarm, _address_translation,
 };
 use libp2p_swarm::{ConnectionId, THandler, THandlerOutEvent};
 
@@ -144,10 +144,12 @@ pub struct Config {
 
     /// How many entries of discovered peers to keep before we discard
     /// the least-recently used one.
-    ///
-    /// Disabled by default.
+    #[deprecated(since = "0.45.1", note = "Use `Config::cache_config` instead.")]
     pub cache_size: usize,
 
+    /// Configuration for the LRU cache of discovered peers.
+    pub cache_config: Option<PeerAddressesConfig>,
+
     /// Whether to include our listen addresses in our responses. If enabled,
     /// we will effectively only share our external addresses.
     ///
@@ -159,13 +161,15 @@ impl Config {
     /// Creates a new configuration for the identify [`Behaviour`] that
     /// advertises the given protocol version and public key.
     pub fn new(protocol_version: String, local_public_key: PublicKey) -> Self {
+        #[allow(deprecated)]
         Self {
             protocol_version,
             agent_version: format!("rust-libp2p/{}", env!("CARGO_PKG_VERSION")),
             local_public_key,
             interval: Duration::from_secs(5 * 60),
             push_listen_addr_updates: false,
             cache_size: 100,
+            cache_config: None, // TODO: when removing `cache_size`, replace `None` with `Some(Default::default())`
             hide_listen_addrs: false,
         }
     }
@@ -191,7 +195,24 @@ impl Config {
         self
     }
 
+    /// Configures the LRU cache responsible for caching addresses of discovered peers.
+    ///
+    /// If set to [`None`], caching is disabled.
+    pub fn with_cache_config(mut self, cache_config: Option<PeerAddressesConfig>) -> Self {
+        #[allow(deprecated)]
+        {
+            // set cache_size to 0 to ensure if user call `with_cache_config(None)`, caching do get disabled.
+            self.cache_size = 0;
+        }
+        self.cache_config = cache_config;
+        self
+    }
+
     /// Configures the size of the LRU cache, caching addresses of discovered peers.
+    ///
+    /// If `cache_config` is set, then `cache_size` is ignored.
+    #[deprecated(since = "0.45.1", note = "Use `Config::with_cache_config` instead.")]
+    #[allow(deprecated)]
     pub fn with_cache_size(mut self, cache_size: usize) -> Self {
         self.cache_size = cache_size;
         self
@@ -206,11 +227,15 @@ impl Config {
 
 impl Behaviour {
     /// Creates a new identify [`Behaviour`].
-    pub fn new(config: Config) -> Self {
-        let discovered_peers = match NonZeroUsize::new(config.cache_size) {
-            None => PeerCache::disabled(),
-            Some(size) => PeerCache::enabled(size),
-        };
+    pub fn new(mut config: Config) -> Self {
+        #[allow(deprecated)]
+        // If `cache_config` value is not set but `cache_size` is, we build `cache_config` from `cache_size`.
+        if config.cache_config.is_none() {
+            config.cache_config = NonZeroUsize::new(config.cache_size)
+                .map(|cache_size| PeerAddressesConfig::default().with_number_of_peers(cache_size));
+        }
+
+        let discovered_peers = PeerCache::new(config.cache_config.clone());
 
         Self {
             config,
@@ -615,12 +640,8 @@ fn multiaddr_matches_peer_id(addr: &Multiaddr, peer_id: &PeerId) -> bool {
 struct PeerCache(Option<PeerAddresses>);
 
 impl PeerCache {
-    fn disabled() -> Self {
-        Self(None)
-    }
-
-    fn enabled(size: NonZeroUsize) -> Self {
-        Self(Some(PeerAddresses::new(size)))
+    fn new(cache_config: Option<PeerAddressesConfig>) -> Self {
+        Self(cache_config.map(PeerAddresses::new))
     }
 
     fn get(&mut self, peer: &PeerId) -> Vec<Multiaddr> {

protocols/identify/tests/smoke.rs

@@ -1,9 +1,10 @@
 use futures::StreamExt;
 use libp2p_identify as identify;
-use libp2p_swarm::{Swarm, SwarmEvent};
+use libp2p_swarm::{PeerAddressesConfig, Swarm, SwarmEvent};
 use libp2p_swarm_test::SwarmExt;
 use std::collections::HashSet;
 use std::iter;
+use std::num::NonZeroUsize;
 use std::time::{Duration, Instant};
 use tracing_subscriber::EnvFilter;
 
@@ -161,7 +162,10 @@ async fn emits_unique_listen_addresses() {
             identify::Config::new("a".to_string(), identity.public())
                 .with_agent_version("b".to_string())
                 .with_interval(Duration::from_secs(1))
-                .with_cache_size(10),
+                .with_cache_config(Some(
+                    PeerAddressesConfig::default()
+                        .with_number_of_peers(NonZeroUsize::new(10).expect("10 != 0")),
+                )),
         )
     });
     let mut swarm2 = Swarm::new_ephemeral(|identity| {
@@ -233,7 +237,10 @@ async fn hides_listen_addresses() {
             identify::Config::new("a".to_string(), identity.public())
                 .with_agent_version("b".to_string())
                 .with_interval(Duration::from_secs(1))
-                .with_cache_size(10),
+                .with_cache_config(Some(
+                    PeerAddressesConfig::default()
+                        .with_number_of_peers(NonZeroUsize::new(10).expect("10 != 0")),
+                )),
         )
     });
     let mut swarm2 = Swarm::new_ephemeral(|identity| {

swarm/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 - Don't report `NewExternalAddrCandidate` for confirmed external addresses.
   See [PR 5582](https://github.com/libp2p/rust-libp2p/pull/5582).
+- Add `PeerAddressesConfig` and the possibility to configure the number of addresses cached per peer.
+  See [PR 5574](https://github.com/libp2p/rust-libp2p/pull/5574).
 
 ## 0.45.1
 

swarm/src/behaviour.rs

@@ -26,7 +26,7 @@ pub mod toggle;
 
 pub use external_addresses::ExternalAddresses;
 pub use listen_addresses::ListenAddresses;
-pub use peer_addresses::PeerAddresses;
+pub use peer_addresses::{PeerAddresses, PeerAddressesConfig};
 
 use crate::connection::ConnectionId;
 use crate::dial_opts::DialOpts;

swarm/src/behaviour/peer_addresses.rs

@@ -8,16 +8,56 @@ use lru::LruCache;
 
 use std::num::NonZeroUsize;
 
+#[derive(Debug, Clone)]
+/// Configuration of a [`PeerAddresses`] instance.
+pub struct PeerAddressesConfig {
+    /// Capacity of the [`PeerAddresses`] cache.
+    number_of_peers: NonZeroUsize,
+
+    /// Maximum number of cached addresses per peer.
+    number_of_addresses_per_peer: NonZeroUsize,
+}
+
+impl PeerAddressesConfig {
+    /// Configure the capacity of the [`PeerAddresses`] cache.
+    pub fn with_number_of_peers(mut self, number_of_peers: NonZeroUsize) -> Self {
+        self.number_of_peers = number_of_peers;
+        self
+    }
+
+    /// Configure the maximum number of cached addresses per peer.
+    pub fn with_number_of_addresses_by_peer(
+        mut self,
+        number_of_addresses_per_peer: NonZeroUsize,
+    ) -> Self {
+        self.number_of_addresses_per_peer = number_of_addresses_per_peer;
+        self
+    }
+}
+
+impl Default for PeerAddressesConfig {
+    fn default() -> Self {
+        Self {
+            number_of_peers: NonZeroUsize::new(100).expect("100 != 0"),
+            number_of_addresses_per_peer: NonZeroUsize::new(10).expect("10 != 0"),
+        }
+    }
+}
+
 /// Struct for tracking peers' external addresses of the [`Swarm`](crate::Swarm).
 #[derive(Debug)]
-pub struct PeerAddresses(LruCache<PeerId, LruCache<Multiaddr, ()>>);
+pub struct PeerAddresses {
+    config: PeerAddressesConfig,
+    inner: LruCache<PeerId, LruCache<Multiaddr, ()>>,
+}
 
 impl PeerAddresses {
     /// Creates a [`PeerAddresses`] cache with capacity for the given number of peers.
     ///
-    /// For each peer, we will at most store 10 addresses.
-    pub fn new(number_of_peers: NonZeroUsize) -> Self {
-        Self(LruCache::new(number_of_peers))
+    /// For each peer, we will at most store `config.number_of_addresses_by_peer` addresses.
+    pub fn new(config: PeerAddressesConfig) -> Self {
+        let inner = LruCache::new(config.number_of_peers);
+        Self { config, inner }
     }
 
     /// Feed a [`FromSwarm`] event to this struct.
@@ -50,12 +90,12 @@ impl PeerAddresses {
     pub fn add(&mut self, peer: PeerId, address: Multiaddr) -> bool {
         match prepare_addr(&peer, &address) {
             Ok(address) => {
-                if let Some(cached) = self.0.get_mut(&peer) {
+                if let Some(cached) = self.inner.get_mut(&peer) {
                     cached.put(address, ()).is_none()
                 } else {
-                    let mut set = LruCache::new(NonZeroUsize::new(10).expect("10 > 0"));
+                    let mut set = LruCache::new(self.config.number_of_addresses_per_peer);
                     set.put(address, ());
-                    self.0.put(peer, set);
+                    self.inner.put(peer, set);
 
                     true
                 }
@@ -66,7 +106,7 @@ impl PeerAddresses {
 
     /// Returns peer's external addresses.
     pub fn get(&mut self, peer: &PeerId) -> impl Iterator<Item = Multiaddr> + '_ {
-        self.0
+        self.inner
             .get(peer)
             .into_iter()
             .flat_map(|c| c.iter().map(|(m, ())| m))
@@ -76,7 +116,7 @@ impl PeerAddresses {
     /// Removes address from peer addresses cache.
     /// Returns true if the address was removed.
     pub fn remove(&mut self, peer: &PeerId, address: &Multiaddr) -> bool {
-        match self.0.get_mut(peer) {
+        match self.inner.get_mut(peer) {
             Some(addrs) => match prepare_addr(peer, address) {
                 Ok(address) => addrs.pop(&address).is_some(),
                 Err(_) => false,
@@ -92,7 +132,7 @@ fn prepare_addr(peer: &PeerId, addr: &Multiaddr) -> Result<Multiaddr, Multiaddr>
 
 impl Default for PeerAddresses {
     fn default() -> Self {
-        Self(LruCache::new(NonZeroUsize::new(100).unwrap()))
+        Self::new(Default::default())
     }
 }
 

swarm/src/lib.rs

@@ -111,7 +111,8 @@ pub use behaviour::{
     AddressChange, CloseConnection, ConnectionClosed, DialFailure, ExpiredListenAddr,
     ExternalAddrExpired, ExternalAddresses, FromSwarm, ListenAddresses, ListenFailure,
     ListenerClosed, ListenerError, NetworkBehaviour, NewExternalAddrCandidate,
-    NewExternalAddrOfPeer, NewListenAddr, NotifyHandler, PeerAddresses, ToSwarm,
+    NewExternalAddrOfPeer, NewListenAddr, NotifyHandler, PeerAddresses, PeerAddressesConfig,
+    ToSwarm,
 };
 pub use connection::pool::ConnectionCounters;
 pub use connection::{ConnectionError, ConnectionId, SupportedProtocols};