networking: Fix get_closest_peers in batches.

autonomys · Dec 20, 2023 · a3a19e1 · a3a19e1
1 parent 9d6a432
commit a3a19e1
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/crates/subspace-networking/src/node.rs b/crates/subspace-networking/src/node.rs
@@ -400,15 +400,21 @@ impl Node {
         &self,
         key: Multihash,
     ) -> Result<impl Stream<Item = PeerId>, GetClosestPeersError> {
-        self.get_closest_peers_internal(key).await
+        self.get_closest_peers_internal(key, true).await
     }
 
     /// Get closest peers by multihash key using Kademlia DHT.
     async fn get_closest_peers_internal(
         &self,
         key: Multihash,
+        acquire_permit: bool,
     ) -> Result<impl Stream<Item = PeerId>, GetClosestPeersError> {
-        let permit = self.shared.rate_limiter.acquire_kademlia_permit().await;
+        let permit = if acquire_permit {
+            Some(self.shared.rate_limiter.acquire_kademlia_permit().await)
+        } else {
+            None
+        };
+
         trace!(?key, "Starting 'GetClosestPeers' request.");
 
         let (result_sender, result_receiver) = mpsc::unbounded();
@@ -602,7 +608,7 @@ impl NodeRequestsBatchHandle {
         &mut self,
         key: Multihash,
     ) -> Result<impl Stream<Item = PeerId>, GetClosestPeersError> {
-        self.node.get_closest_peers_internal(key).await
+        self.node.get_closest_peers_internal(key, false).await
     }
     /// Sends the generic request to the peer and awaits the result.
     pub async fn send_generic_request<Request>(

diff --git a/crates/subspace-networking/src/node_runner.rs b/crates/subspace-networking/src/node_runner.rs
@@ -60,7 +60,7 @@ enum QueryResultSender {
     ClosestPeers {
         sender: mpsc::UnboundedSender<PeerId>,
         // Just holding onto permit while data structure is not dropped
-        _permit: RateLimiterPermit,
+        _permit: Option<RateLimiterPermit>,
     },
     Providers {
         sender: mpsc::UnboundedSender<PeerId>,

diff --git a/crates/subspace-networking/src/shared.rs b/crates/subspace-networking/src/shared.rs
@@ -79,7 +79,7 @@ pub(crate) enum Command {
     GetClosestPeers {
         key: Multihash,
         result_sender: mpsc::UnboundedSender<PeerId>,
-        permit: RateLimiterPermit,
+        permit: Option<RateLimiterPermit>,
     },
     GenericRequest {
         peer_id: PeerId,