feat(kad): allow to explicitly set Mode::{Client,Server}

protocols/kad/CHANGELOG.md

@@ -1,3 +1,10 @@
+## 0.44.2 - unreleased
+
+- feat(kad): allow to explicitly set Mode::{Client,Server}
+  See [PR 4132]
+
+[PR 4132]: https://github.com/libp2p/rust-libp2p/pull/4132
+
 ## 0.44.1
 
 - Expose `KBucketDistance`.

protocols/kad/src/behaviour.rs

@@ -52,7 +52,7 @@ use smallvec::SmallVec;
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::fmt;
 use std::num::NonZeroUsize;
-use std::task::{Context, Poll};
+use std::task::{Context, Poll, Waker};
 use std::time::Duration;
 use std::vec;
 use thiserror::Error;
@@ -114,6 +114,8 @@ pub struct Kademlia<TStore> {
     local_peer_id: PeerId,
 
     mode: Mode,
+    auto_mode: bool,
+    waker: Option<Waker>,
 
     /// The record storage.
     store: TStore,
@@ -456,6 +458,8 @@ where
             local_peer_id: id,
             connections: Default::default(),
             mode: Mode::Client,
+            auto_mode: true,
+            waker: None,
         }
     }
 
@@ -990,6 +994,81 @@ where
         id
     }
 
+    /// Either set Kademlia [`Mode`] explicitly via `Some(_)` or enable automatic configuration of the Kademlia [`Mode`] based on the external addresses available via `None`.
+    pub fn set_mode(&mut self, mode: Option<Mode>) {
+        match mode {
+            Some(mode) => {
+                self.mode = mode;
+                self.auto_mode = false;
+                if !self.connections.is_empty() {
+                    self.reconfigure_mode();
+                }
+            }
+            None => {
+                self.auto_mode = true;
+                self.determine_mode();
+            }
+        }
+
+        if let Some(waker) = self.waker.take() {
+            waker.wake();
+        }
+    }
+
+    fn reconfigure_mode(&mut self) {
+        self.queued_events
+            .extend(
+                self.connections
+                    .iter()
+                    .map(|(conn_id, peer_id)| ToSwarm::NotifyHandler {
+                        peer_id: *peer_id,
+                        handler: NotifyHandler::One(*conn_id),
+                        event: KademliaHandlerIn::ReconfigureMode {
+                            new_mode: self.mode,
+                        },
+                    }),
+            );
+    }
+
+    fn determine_mode(&mut self) {
+        self.mode = match (self.external_addresses.as_slice(), self.mode) {
+            ([], Mode::Server) => {
+                log::debug!("Switching to client-mode because we no longer have any confirmed external addresses");
+
+                Mode::Client
+            }
+            ([], Mode::Client) => {
+                // Previously client-mode, now also client-mode because no external addresses.
+
+                Mode::Client
+            }
+            (confirmed_external_addresses, Mode::Client) => {
+                if log::log_enabled!(log::Level::Debug) {
+                    let confirmed_external_addresses =
+                        to_comma_separated_list(confirmed_external_addresses);
+
+                    log::debug!("Switching to server-mode assuming that one of [{confirmed_external_addresses}] is externally reachable");
+                }
+
+                Mode::Server
+            }
+            (confirmed_external_addresses, Mode::Server) => {
+                debug_assert!(
+                    !confirmed_external_addresses.is_empty(),
+                    "Previous match arm handled empty list"
+                );
+
+                // Previously, server-mode, now also server-mode because > 1 external address. Don't log anything to avoid spam.
+
+                Mode::Server
+            }
+        };
+
+        if !self.connections.is_empty() {
+            self.reconfigure_mode();
+        }
+    }
+
     /// Processes discovered peers from a successful request in an iterative `Query`.
     fn discovered<'a, I>(&'a mut self, query_id: &QueryId, source: &PeerId, peers: I)
     where
@@ -2424,6 +2503,8 @@ where
                 }
             }
 
+            self.waker = Some(cx.waker().clone());
+
             // No immediate event was produced as a result of a finished query.
             // If no new events have been queued either, signal `NotReady` to
             // be polled again later.
@@ -2437,40 +2518,7 @@ where
         self.listen_addresses.on_swarm_event(&event);
         let external_addresses_changed = self.external_addresses.on_swarm_event(&event);
 
-        self.mode = match (self.external_addresses.as_slice(), self.mode) {
-            ([], Mode::Server) => {
-                log::debug!("Switching to client-mode because we no longer have any confirmed external addresses");
-
-                Mode::Client
-            }
-            ([], Mode::Client) => {
-                // Previously client-mode, now also client-mode because no external addresses.
-
-                Mode::Client
-            }
-            (confirmed_external_addresses, Mode::Client) => {
-                if log::log_enabled!(log::Level::Debug) {
-                    let confirmed_external_addresses =
-                        to_comma_separated_list(confirmed_external_addresses);
-
-                    log::debug!("Switching to server-mode assuming that one of [{confirmed_external_addresses}] is externally reachable");
-                }
-
-                Mode::Server
-            }
-            (confirmed_external_addresses, Mode::Server) => {
-                debug_assert!(
-                    !confirmed_external_addresses.is_empty(),
-                    "Previous match arm handled empty list"
-                );
-
-                // Previously, server-mode, now also server-mode because > 1 external address. Don't log anything to avoid spam.
-
-                Mode::Server
-            }
-        };
-
-        if external_addresses_changed && !self.connections.is_empty() {
+        if self.auto_mode && external_addresses_changed {
             let num_connections = self.connections.len();
 
             log::debug!(
@@ -2479,18 +2527,7 @@ where
                 if num_connections > 1 { "s" } else { "" }
             );
 
-            self.queued_events
-                .extend(
-                    self.connections
-                        .iter()
-                        .map(|(conn_id, peer_id)| ToSwarm::NotifyHandler {
-                            peer_id: *peer_id,
-                            handler: NotifyHandler::One(*conn_id),
-                            event: KademliaHandlerIn::ReconfigureMode {
-                                new_mode: self.mode,
-                            },
-                        }),
-                );
+            self.determine_mode();
         }
 
         match event {

protocols/kad/src/lib.rs

@@ -68,7 +68,7 @@ pub use behaviour::{
     AddProviderContext, AddProviderError, AddProviderOk, AddProviderPhase, AddProviderResult,
     BootstrapError, BootstrapOk, BootstrapResult, GetClosestPeersError, GetClosestPeersOk,
     GetClosestPeersResult, GetProvidersError, GetProvidersOk, GetProvidersResult, GetRecordError,
-    GetRecordOk, GetRecordResult, InboundRequest, NoKnownPeers, PeerRecord, PutRecordContext,
+    GetRecordOk, GetRecordResult, InboundRequest, Mode, NoKnownPeers, PeerRecord, PutRecordContext,
     PutRecordError, PutRecordOk, PutRecordPhase, PutRecordResult, QueryInfo, QueryMut, QueryRef,
     QueryResult, QueryStats, RoutingUpdate,
 };

protocols/kad/tests/client_mode.rs

@@ -1,7 +1,7 @@
 use libp2p_identify as identify;
 use libp2p_identity as identity;
 use libp2p_kad::store::MemoryStore;
-use libp2p_kad::{Kademlia, KademliaConfig, KademliaEvent};
+use libp2p_kad::{Kademlia, KademliaConfig, KademliaEvent, Mode};
 use libp2p_swarm::Swarm;
 use libp2p_swarm_test::SwarmExt;
 
@@ -111,6 +111,50 @@ async fn adding_an_external_addresses_activates_server_mode_on_existing_connecti
     }
 }
 
+#[async_std::test]
+async fn set_client_to_server_mode() {
+    let _ = env_logger::try_init();
+
+    let mut client = Swarm::new_ephemeral(MyBehaviour::new);
+    client.behaviour_mut().kad.set_mode(Some(Mode::Client));
+
+    let mut server = Swarm::new_ephemeral(MyBehaviour::new);
+
+    server.listen().await;
+    client.connect(&mut server).await;
+
+    let server_peer_id = *server.local_peer_id();
+
+    match libp2p_swarm_test::drive(&mut client, &mut server).await {
+        (
+            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(_), MyBehaviourEvent::Kad(KademliaEvent::RoutingUpdated { peer, .. })],
+            [MyBehaviourEvent::Identify(_), MyBehaviourEvent::Identify(identify::Event::Received { info, .. })],
+        ) => {
+            assert_eq!(peer, server_peer_id);
+            assert!(info
+                .protocols
+                .iter()
+                .all(|proto| libp2p_kad::PROTOCOL_NAME.ne(proto)))
+        }
+        other => panic!("Unexpected events: {other:?}"),
+    }
+
+    client.behaviour_mut().kad.set_mode(Some(Mode::Server));
+
+    match libp2p_swarm_test::drive(&mut client, &mut server).await {
+        (
+            [MyBehaviourEvent::Identify(_)],
+            [MyBehaviourEvent::Identify(identify::Event::Received { info, .. }), MyBehaviourEvent::Kad(_)],
+        ) => {
+            assert!(info
+                .protocols
+                .iter()
+                .any(|proto| libp2p_kad::PROTOCOL_NAME.eq(proto)))
+        }
+        other => panic!("Unexpected events: {other:?}"),
+    }
+}
+
 #[derive(libp2p_swarm::NetworkBehaviour)]
 #[behaviour(prelude = "libp2p_swarm::derive_prelude")]
 struct MyBehaviour {