track if there was a successfull bootstrap and the query id of the au…

…tomated bootstrap
libp2p · Nov 30, 2023 · f2aaad2 · f2aaad2
1 parent f72a8a1
commit f2aaad2
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 21 deletions.
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
@@ -126,8 +126,11 @@ pub struct Behaviour<TStore> {
     /// The timer used to poll [`Behaviour::bootstrap`].
     bootstrap_timer: Option<Delay>,
 
-    /// Tracks if the latest bootstrap was successfull.
-    initial_bootstrap: Option<InitialBootstrap>,
+    /// Tracks if the there was a successfull bootstrap.
+    any_bootstrap_successful: bool,
+
+    /// Tracks the `QueryId` of the automated bootstrap.
+    current_automated_bootstrap: Option<QueryId>,
 }
 
 /// The configurable strategies for the insertion of peers
@@ -477,7 +480,8 @@ where
             no_events_waker: None,
             bootstrap_interval: config.bootstrap_interval,
             bootstrap_timer,
-            initial_bootstrap: None,
+            any_bootstrap_successful: false,
+            current_automated_bootstrap: None,
         }
     }
 
@@ -579,7 +583,9 @@ where
                 };
                 match entry.insert(addresses.clone(), status) {
                     kbucket::InsertResult::Inserted => {
-                        self.handle_bootstrap();
+                        if !self.any_bootstrap_successful {
+                            self.maybe_bootstrap();
+                        }
                         self.queued_events.push_back(ToSwarm::GenerateEvent(
                             Event::RoutingUpdated {
                                 peer: *peer,
@@ -915,8 +921,7 @@ where
             Err(NoKnownPeers())
         } else {
             let inner = QueryInner::new(info);
-            let query_id = self.queries.add_iter_closest(local_key, peers, inner);
-            Ok(query_id)
+            Ok(self.queries.add_iter_closest(local_key, peers, inner))
         }
     }
 
@@ -1310,7 +1315,9 @@ where
                         let addresses = Addresses::new(a);
                         match entry.insert(addresses.clone(), new_status) {
                             kbucket::InsertResult::Inserted => {
-                                self.handle_bootstrap();
+                                if !self.any_bootstrap_successful {
+                                    self.maybe_bootstrap();
+                                }
                                 let event = Event::RoutingUpdated {
                                     peer,
                                     is_new_peer: true,
@@ -2074,11 +2081,14 @@ where
         }
     }
 
-    fn handle_bootstrap(&mut self) {
-        if self.initial_bootstrap.is_none() {
+    fn maybe_bootstrap(&mut self) {
+        if self.current_automated_bootstrap.is_none() {
             match self.bootstrap() {
-                Ok(query_id) => self.initial_bootstrap = Some(InitialBootstrap::Active(query_id)),
-                Err(_) => self.initial_bootstrap = None,
+                Ok(query_id) => {
+                    self.current_automated_bootstrap = Some(query_id);
+                    self.any_bootstrap_successful = true;
+                }
+                Err(_) => self.current_automated_bootstrap = None,
             }
         }
     }
@@ -2509,7 +2519,9 @@ where
             if let Poll::Ready(()) = Pin::new(&mut bootstrap_timer).poll(cx) {
                 if let Some(interval) = self.bootstrap_interval {
                     bootstrap_timer.reset(interval);
-                    self.handle_bootstrap();
+                    if !self.any_bootstrap_successful {
+                        self.maybe_bootstrap();
+                    }
                 }
             }
             self.bootstrap_timer = Some(bootstrap_timer);
@@ -3396,9 +3408,3 @@ where
         .collect::<Vec<_>>()
         .join(", ")
 }
-
-#[derive(PartialEq)]
-pub enum InitialBootstrap {
-    Active(QueryId),
-    Completed,
-}
diff --git a/protocols/kad/src/lib.rs b/protocols/kad/src/lib.rs
@@ -58,9 +58,9 @@ pub use behaviour::{
     AddProviderContext, AddProviderError, AddProviderOk, AddProviderPhase, AddProviderResult,
     BootstrapError, BootstrapOk, BootstrapResult, GetClosestPeersError, GetClosestPeersOk,
     GetClosestPeersResult, GetProvidersError, GetProvidersOk, GetProvidersResult, GetRecordError,
-    GetRecordOk, GetRecordResult, InboundRequest, InitialBootstrap, Mode, NoKnownPeers, PeerRecord,
-    PutRecordContext, PutRecordError, PutRecordOk, PutRecordPhase, PutRecordResult, QueryInfo,
-    QueryMut, QueryRef, QueryResult, QueryStats, RoutingUpdate,
+    GetRecordOk, GetRecordResult, InboundRequest, Mode, NoKnownPeers, PeerRecord, PutRecordContext,
+    PutRecordError, PutRecordOk, PutRecordPhase, PutRecordResult, QueryInfo, QueryMut, QueryRef,
+    QueryResult, QueryStats, RoutingUpdate,
 };
 pub use behaviour::{
     Behaviour, BucketInserts, Caching, Config, Event, ProgressStep, Quorum, StoreInserts,