From a29538d5121142541fb1a27b389ad73e42e075de Mon Sep 17 00:00:00 2001
From: Emilia Hane <elsaemiliaevahane@gmail.com>
Date: Mon, 22 Jan 2024 21:09:42 +0100
Subject: [PATCH] Allow optimistic mirror of kbuckets with read lock only

---
 src/discv5.rs               | 28 ++++++++++++++++++++++++++++
 src/handler/request_call.rs |  2 +-
 src/kbucket.rs              |  2 +-
 src/kbucket/entry.rs        |  4 +---
 4 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/discv5.rs b/src/discv5.rs
index 0acf1587a..29d11da8b 100644
--- a/src/discv5.rs
+++ b/src/discv5.rs
@@ -453,6 +453,17 @@ impl<P: ProtocolIdentity> Discv5<P> {
             .collect()
     }
 
+    /// Takes only a read lock on the kbuckets. This means pending entries aren't added. Otherwise
+    /// same as [`Self::table_entries_id`]. Returns an iterator over all ENR node IDs of nodes
+    /// currently contained in the routing table.
+    pub fn table_entries_id_rlock(&self) -> Vec<NodeId> {
+        self.kbuckets
+            .read()
+            .iter_ref()
+            .map(|entry| *entry.node.key.preimage())
+            .collect()
+    }
+
     /// Returns an iterator over all the ENR's of nodes currently contained in the routing table.
     pub fn table_entries_enr(&self) -> Vec<Enr> {
         self.kbuckets
@@ -477,6 +488,23 @@ impl<P: ProtocolIdentity> Discv5<P> {
             .collect()
     }
 
+    /// Takes only a read lock on the kbuckets. This means pending entries aren't added. Otherwise
+    /// same as [`Self::table_entries`]. Returns an iterator over all ENR node IDs of nodes
+    /// currently contained in the routing table.
+    pub fn table_entries_rlock(&self) -> Vec<(NodeId, Enr, NodeStatus)> {
+        self.kbuckets
+            .read()
+            .iter_ref()
+            .map(|entry| {
+                (
+                    *entry.node.key.preimage(),
+                    entry.node.value.clone(),
+                    entry.status,
+                )
+            })
+            .collect()
+    }
+
     /// Requests the ENR of a node corresponding to multiaddr or multi-addr string.
     ///
     /// Only `ed25519` and `secp256k1` key types are currently supported.
diff --git a/src/handler/request_call.rs b/src/handler/request_call.rs
index b91c7643e..42506aa76 100644
--- a/src/handler/request_call.rs
+++ b/src/handler/request_call.rs
@@ -1,4 +1,4 @@
-pub use crate::node_info::{NodeAddress, NodeContact};
+pub use crate::node_info::NodeContact;
 use crate::{
     packet::Packet,
     rpc::{Request, RequestBody},
diff --git a/src/kbucket.rs b/src/kbucket.rs
index 67d137c6a..1489a7bdc 100644
--- a/src/kbucket.rs
+++ b/src/kbucket.rs
@@ -533,7 +533,7 @@ where
 
     /// Returns an iterator over all the entries in the routing table.
     /// Does not add pending node to kbucket to get an iterator which
-    /// takes a reference instead of a mutable reference.
+    /// takes a mutable reference instead of a reference.
     pub fn iter_ref(&self) -> impl Iterator<Item = EntryRefView<'_, TNodeId, TVal>> {
         self.buckets.iter().flat_map(move |table| {
             table.iter().map(move |n| EntryRefView {
diff --git a/src/kbucket/entry.rs b/src/kbucket/entry.rs
index 1e0304048..97be95de1 100644
--- a/src/kbucket/entry.rs
+++ b/src/kbucket/entry.rs
@@ -25,9 +25,7 @@
 //! representing the nodes participating in the Kademlia DHT.
 
 pub use super::{
-    bucket::{
-        AppliedPending, ConnectionState, InsertResult, Node, NodeStatus, MAX_NODES_PER_BUCKET,
-    },
+    bucket::{AppliedPending, ConnectionState, InsertResult, Node, NodeStatus},
     key::*,
     ConnectionDirection,
 };