first draft

Cargo.toml

@@ -65,7 +65,7 @@ cuprate-levin               = { path = "net/levin"              ,default-feature
 cuprate-wire                = { path = "net/wire"               ,default-features = false}
 cuprate-p2p                 = { path = "p2p/p2p"                ,default-features = false}
 cuprate-p2p-core            = { path = "p2p/p2p-core"           ,default-features = false}
-cuprate-p2p-bucket          = { path = "p2p/p2p-bucket"         ,default-features = false}
+cuprate-p2p-bucket          = { path = "p2p/bucket"             ,default-features = false}
 cuprate-dandelion-tower     = { path = "p2p/dandelion-tower"    ,default-features = false}
 cuprate-async-buffer        = { path = "p2p/async-buffer"       ,default-features = false}
 cuprate-address-book        = { path = "p2p/address-book"       ,default-features = false}

p2p/address-book/Cargo.toml

@@ -10,6 +10,7 @@ authors = ["Boog900"]
 cuprate-constants   = { workspace = true }
 cuprate-pruning     = { workspace = true }
 cuprate-p2p-core    = { workspace = true, features = ["borsh"] }
+cuprate-p2p-bucket  = { workspace = true }
 
 tower = { workspace = true, features = ["util"] }
 tokio = { workspace = true, features = ["time", "fs", "rt"]}
@@ -19,7 +20,7 @@ futures = { workspace = true, features = ["std"] }
 
 thiserror = { workspace = true }
 tracing = { workspace = true, features = ["std", "attributes"] }
-indexmap = { workspace = true, features = ["std"] }
+#indexmap = { workspace = true, features = ["std"] }
 
 rand = { workspace = true, features = ["std", "std_rng"] }
 

p2p/address-book/src/book.rs

@@ -357,6 +357,13 @@ impl<Z: BorshNetworkZone> AddressBook<Z> {
         self.connected_peers.insert(internal_peer_id, peer);
         Ok(())
     }
+
+    fn add_white_peer(
+        &mut self,
+        white_peer: ZoneSpecificPeerListEntryBase<Z::Addr>
+    ) {
+        self.white_list.add_new_peer(white_peer);
+    }
 }
 
 impl<Z: BorshNetworkZone> Service<AddressBookRequest<Z>> for AddressBook<Z> {
@@ -401,6 +408,10 @@ impl<Z: BorshNetworkZone> Service<AddressBookRequest<Z>> for AddressBook<Z> {
                 self.handle_incoming_peer_list(peer_list);
                 Ok(AddressBookResponse::Ok)
             }
+            AddressBookRequest::AddWhitePeer(white_peer) =>{
+                self.add_white_peer(white_peer);
+                Ok(AddressBookResponse::Ok)
+            }
             AddressBookRequest::TakeRandomWhitePeer { height } => self
                 .take_random_white_peer(height)
                 .map(AddressBookResponse::Peer)

p2p/address-book/src/peer_list.rs

@@ -1,11 +1,11 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 
-use indexmap::IndexMap;
 use rand::prelude::*;
 
 use cuprate_constants::block::MAX_BLOCK_HEIGHT_USIZE;
 use cuprate_p2p_core::{services::ZoneSpecificPeerListEntryBase, NetZoneAddress, NetworkZone};
 use cuprate_pruning::PruningSeed;
+use cuprate_p2p_bucket::BucketMap;
 
 #[cfg(test)]
 pub(crate) mod tests;
@@ -16,7 +16,7 @@ pub(crate) mod tests;
 #[derive(Debug)]
 pub(crate) struct PeerList<Z: NetworkZone> {
     /// The peers with their peer data.
-    pub peers: IndexMap<Z::Addr, ZoneSpecificPeerListEntryBase<Z::Addr>>,
+    pub peers: BucketMap<ZoneSpecificPeerListEntryBase<Z::Addr>, 8>,
     /// An index of Pruning seed to address, so can quickly grab peers with the blocks
     /// we want.
     ///
@@ -33,7 +33,7 @@ pub(crate) struct PeerList<Z: NetworkZone> {
 impl<Z: NetworkZone> PeerList<Z> {
     /// Creates a new peer list.
     pub(crate) fn new(list: Vec<ZoneSpecificPeerListEntryBase<Z::Addr>>) -> Self {
-        let mut peers = IndexMap::with_capacity(list.len());
+        let mut peers = BucketMap::<_, 8>::new();
         let mut pruning_seeds = BTreeMap::new();
         let mut ban_ids = HashMap::with_capacity(list.len());
 
@@ -48,7 +48,7 @@ impl<Z: NetworkZone> PeerList<Z> {
                 .or_insert_with(Vec::new)
                 .push(peer.adr);
 
-            peers.insert(peer.adr, peer);
+            peers.push(peer);
         }
         Self {
             peers,
@@ -64,7 +64,7 @@ impl<Z: NetworkZone> PeerList<Z> {
 
     /// Adds a new peer to the peer list
     pub(crate) fn add_new_peer(&mut self, peer: ZoneSpecificPeerListEntryBase<Z::Addr>) {
-        if self.peers.insert(peer.adr, peer).is_none() {
+        if self.peers.push(peer).is_none() {
             #[expect(clippy::unwrap_or_default, reason = "It's more clear with this")]
             self.pruning_seeds
                 .entry(peer.pruning_seed)
@@ -94,7 +94,14 @@ impl<Z: NetworkZone> PeerList<Z> {
         // Take a random peer and see if it's in the list of must_keep_peers, if it is try again.
         // TODO: improve this
 
+        // Return early if no items
+        if self.peers.is_empty() {
+            return None
+        }
+
         for _ in 0..3 {
+
+            // In case we require a specific height
             if let Some(needed_height) = block_needed {
                 let (_, addresses_with_block) = self.pruning_seeds.iter().find(|(seed, _)| {
                     // TODO: factor in peer blockchain height?
@@ -109,18 +116,15 @@ impl<Z: NetworkZone> PeerList<Z> {
                 }
 
                 return self.remove_peer(&peer);
-            }
-            let len = self.len();
-
-            if len == 0 {
-                return None;
-            }
-
-            let n = r.gen_range(0..len);
-
-            let (&key, _) = self.peers.get_index(n).unwrap();
-            if !must_keep_peers.contains(&key) {
-                return self.remove_peer(&key);
+            } 
+
+            // Straightforward if we don't need a specific height.
+            let random_peer = self.peers.get_random(r);
+            if let Some(peer) = random_peer {
+
+                if !must_keep_peers.contains(&peer.adr) {
+                    return self.remove_peer(&peer.adr.clone())
+                }
             }
         }
 
@@ -132,26 +136,22 @@ impl<Z: NetworkZone> PeerList<Z> {
         r: &mut R,
         len: usize,
     ) -> Vec<ZoneSpecificPeerListEntryBase<Z::Addr>> {
-        let mut peers = self.peers.values().copied().choose_multiple(r, len);
-        // Order of the returned peers is not random, I am unsure of the impact of this, potentially allowing someone to make guesses about which peers
-        // were connected first.
-        // So to mitigate this shuffle the result.
-        peers.shuffle(r);
-        peers.drain(len.min(peers.len())..peers.len());
-        peers
+        (0..len)
+            .filter_map(|_| self.peers.get_random(r).copied())
+            .collect::<Vec<_>>()
     }
 
     /// Returns a mutable reference to a peer.
     pub(crate) fn get_peer_mut(
         &mut self,
-        peer: &Z::Addr,
+        peer_addr: &Z::Addr,
     ) -> Option<&mut ZoneSpecificPeerListEntryBase<Z::Addr>> {
-        self.peers.get_mut(peer)
+        self.peers.get_mut(peer_addr)
     }
 
     /// Returns true if the list contains this peer.
     pub(crate) fn contains_peer(&self, peer: &Z::Addr) -> bool {
-        self.peers.contains_key(peer)
+        self.peers.contains(peer)
     }
 
     /// Removes a peer from the pruning idx
@@ -197,7 +197,7 @@ impl<Z: NetworkZone> PeerList<Z> {
         &mut self,
         peer: &Z::Addr,
     ) -> Option<ZoneSpecificPeerListEntryBase<Z::Addr>> {
-        let peer_eb = self.peers.swap_remove(peer)?;
+        let peer_eb = self.peers.remove(peer)?;
         self.remove_peer_from_all_idxs(&peer_eb);
         Some(peer_eb)
     }
@@ -226,13 +226,13 @@ impl<Z: NetworkZone> PeerList<Z> {
         let target_removed = self.len() - new_len;
         let mut removed_count = 0;
         let mut peers_to_remove: Vec<Z::Addr> = Vec::with_capacity(target_removed);
-
+        
         for peer_adr in self.peers.keys() {
             if removed_count >= target_removed {
                 break;
             }
-            if !must_keep_peers.contains(peer_adr) {
-                peers_to_remove.push(*peer_adr);
+            if !must_keep_peers.contains(&peer_adr) {
+                peers_to_remove.push(peer_adr);
                 removed_count += 1;
             }
         }

p2p/address-book/src/peer_list/tests.rs

@@ -68,7 +68,7 @@ fn peer_list_reduce_length() {
 #[test]
 fn peer_list_reduce_length_with_peers_we_need() {
     let mut peer_list = make_fake_peer_list(0, 500);
-    let must_keep_peers = peer_list.peers.keys().copied().collect::<HashSet<_>>();
+    let must_keep_peers = peer_list.peers.keys().collect::<HashSet<_>>();
 
     let target_len = 49;
 
@@ -93,7 +93,7 @@ fn peer_list_remove_specific_peer() {
         addrs.iter().for_each(|adr| assert_ne!(adr, &peer.adr));
     }
 
-    assert!(!peers.contains_key(&peer.adr));
+    assert!(!peers.contains(&peer.adr));
 }
 
 #[test]
@@ -170,7 +170,7 @@ fn peer_list_ban_peers() {
     let ban_id = peer.adr.ban_id();
 
     assert_eq!(peer_list.ban_ids.get(&ban_id), None);
-    for (addr, _) in peer_list.peers {
+    for addr in peer_list.peers.keys() {
         assert_ne!(addr.ban_id(), ban_id);
     }
 }

p2p/address-book/src/store.rs

@@ -86,11 +86,11 @@ mod tests {
         assert_eq!(gray_list.peers.len(), gray_list_2.peers.len());
 
         for addr in white_list.peers.keys() {
-            assert!(white_list_2.contains_peer(addr));
+            assert!(white_list_2.contains_peer(&addr));
         }
 
         for addr in gray_list.peers.keys() {
-            assert!(gray_list_2.contains_peer(addr));
+            assert!(gray_list_2.contains_peer(&addr));
         }
     }
 }