P2P / Networking — Implementation of Perigee p2p Routing.

Cargo.toml

@@ -115,6 +115,7 @@ kaspa-p2p-flows = { version = "1.1.0-rc.2", path = "protocol/flows" }
 kaspa-p2p-lib = { version = "1.1.0-rc.2", path = "protocol/p2p" }
 kaspa-p2p-mining = { version = "1.1.0-rc.2", path = "protocol/mining" }
 kaspa-perf-monitor = { version = "1.1.0-rc.2", path = "metrics/perf_monitor" }
+kaspa-perigeemanager = {version =  "1.1.0-rc.2", path = "components/perigeemanager" }
 kaspa-pow = { version = "1.1.0-rc.2", path = "consensus/pow" }
 kaspa-rpc-core = { version = "1.1.0-rc.2", path = "rpc/core" }
 kaspa-rpc-macros = { version = "1.1.0-rc.2", path = "rpc/macros" }

components/addressmanager/src/lib.rs

@@ -290,8 +290,12 @@ impl AddressManager {
         self.address_store.iterate_addresses()
     }
 
-    pub fn iterate_prioritized_random_addresses(&self, exceptions: HashSet<NetAddress>) -> impl ExactSizeIterator<Item = NetAddress> {
-        self.address_store.iterate_prioritized_random_addresses(exceptions)
+    pub fn iterate_prioritized_random_addresses(
+        &self,
+        priorities: Vec<NetAddress>,
+        exceptions: HashSet<NetAddress>,
+    ) -> impl ExactSizeIterator<Item = NetAddress> {
+        self.address_store.iterate_prioritized_random_addresses(priorities, exceptions)
     }
 
     pub fn ban(&mut self, ip: IpAddress) {
@@ -325,6 +329,18 @@ impl AddressManager {
     pub fn get_all_banned_addresses(&self) -> Vec<IpAddress> {
         self.banned_address_store.iterator().map(|x| IpAddress::from(x.unwrap().0)).collect_vec()
     }
+
+    pub fn set_new_perigee_addresses(&mut self, addresses: Vec<NetAddress>) {
+        self.address_store.set_new_perigee_addresses(addresses);
+    }
+
+    pub fn get_perigee_addresses(&mut self) -> Vec<NetAddress> {
+        self.address_store.get_perigee_addresses().unwrap()
+    }
+
+    pub fn reset_perigee_data(&mut self) {
+        self.address_store.reset_perigee_data();
+    }
 }
 
 mod address_store_with_cache {
@@ -346,7 +362,7 @@ mod address_store_with_cache {
 
     use crate::{
         stores::{
-            address_store::{AddressesStore, DbAddressesStore, Entry},
+            address_store::{AddressesStore, AddressesStoreReader, DbAddressesStore, Entry},
             AddressKey,
         },
         NetAddress, MAX_ADDRESSES, MAX_CONNECTION_FAILED_COUNT,
@@ -369,6 +385,18 @@ mod address_store_with_cache {
             Self { db_store, addresses }
         }
 
+        pub fn set_new_perigee_addresses(&mut self, addresses: Vec<NetAddress>) {
+            self.db_store.set_new_perigee_addresses(addresses).unwrap();
+        }
+
+        pub fn get_perigee_addresses(&mut self) -> Option<Vec<NetAddress>> {
+            self.db_store.get_perigee_addresses().ok()
+        }
+
+        pub fn reset_perigee_data(&mut self) {
+            self.db_store.reset_perigee_data().unwrap();
+        }
+
         pub fn has(&mut self, address: NetAddress) -> bool {
             self.addresses.contains_key(&address.into())
         }
@@ -427,9 +455,11 @@ mod address_store_with_cache {
         ///```
         pub fn iterate_prioritized_random_addresses(
             &self,
+            priorities: Vec<NetAddress>,
             exceptions: HashSet<NetAddress>,
         ) -> impl ExactSizeIterator<Item = NetAddress> {
-            let exceptions: HashSet<AddressKey> = exceptions.into_iter().map(|addr| addr.into()).collect();
+            let exceptions: HashSet<AddressKey> =
+                exceptions.into_iter().chain(priorities.iter().cloned()).map(|addr| addr.into()).collect();
             let mut prefix_counter: HashMap<PrefixBucket, usize> = HashMap::new();
             let (mut weights, filtered_addresses): (Vec<f64>, Vec<NetAddress>) = self
                 .addresses
@@ -447,7 +477,7 @@ mod address_store_with_cache {
                 *weights.get_mut(i).unwrap() /= *prefix_counter.get(&address.prefix_bucket()).unwrap() as f64;
             }
 
-            RandomWeightedIterator::new(weights, filtered_addresses)
+            RandomWeightedIteratorWithPriorities::new(weights, filtered_addresses, priorities)
         }
 
         pub fn remove_by_ip(&mut self, ip: IpAddr) {
@@ -461,30 +491,34 @@ mod address_store_with_cache {
         Store::new(db)
     }
 
-    pub struct RandomWeightedIterator {
+    pub struct RandomWeightedIteratorWithPriorities {
         weighted_index: Option<WeightedIndex<f64>>,
         remaining: usize,
         addresses: Vec<NetAddress>,
+        priorities: Vec<NetAddress>,
     }
 
-    impl RandomWeightedIterator {
-        pub fn new(weights: Vec<f64>, addresses: Vec<NetAddress>) -> Self {
+    impl RandomWeightedIteratorWithPriorities {
+        pub fn new(weights: Vec<f64>, addresses: Vec<NetAddress>, priorities: Vec<NetAddress>) -> Self {
             assert_eq!(weights.len(), addresses.len());
-            let remaining = weights.iter().filter(|&&w| w > 0.0).count();
+            let remaining = weights.iter().filter(|&&w| w > 0.0).count() + priorities.len();
             let weighted_index = match WeightedIndex::new(weights) {
                 Ok(index) => Some(index),
                 Err(WeightedError::NoItem) => None,
                 Err(e) => panic!("{e}"),
             };
-            Self { weighted_index, remaining, addresses }
+            Self { weighted_index, remaining, addresses, priorities: priorities.into_iter().rev().collect() }
         }
     }
 
-    impl Iterator for RandomWeightedIterator {
+    impl Iterator for RandomWeightedIteratorWithPriorities {
         type Item = NetAddress;
 
         fn next(&mut self) -> Option<Self::Item> {
-            if let Some(weighted_index) = self.weighted_index.as_mut() {
+            if !self.priorities.is_empty() {
+                self.remaining -= 1;
+                Some(self.priorities.pop().unwrap())
+            } else if let Some(weighted_index) = self.weighted_index.as_mut() {
                 let i = weighted_index.sample(&mut rand::thread_rng());
                 // Zero the selected address entry
                 match weighted_index.update_weights(&[(i, &0f64)]) {
@@ -507,7 +541,7 @@ mod address_store_with_cache {
         }
     }
 
-    impl ExactSizeIterator for RandomWeightedIterator {}
+    impl ExactSizeIterator for RandomWeightedIteratorWithPriorities {}
 
     #[cfg(test)]
     mod tests {
@@ -526,11 +560,11 @@ mod address_store_with_cache {
         #[test]
         fn test_weighted_iterator() {
             let address = NetAddress::new(IpAddr::V6(Ipv6Addr::LOCALHOST).into(), 1);
-            let iter = RandomWeightedIterator::new(vec![0.2, 0.3, 0.0], vec![address, address, address]);
+            let iter = RandomWeightedIteratorWithPriorities::new(vec![0.2, 0.3, 0.0], vec![address, address, address], vec![]);
             assert_eq!(iter.len(), 2);
             assert_eq!(iter.count(), 2);
 
-            let iter = RandomWeightedIterator::new(vec![], vec![]);
+            let iter = RandomWeightedIteratorWithPriorities::new(vec![], vec![], vec![]);
             assert_eq!(iter.len(), 0);
             assert_eq!(iter.count(), 0);
         }
@@ -603,7 +637,7 @@ mod address_store_with_cache {
                 // The weight sampled expected uniform distribution
                 let prioritized_address_distribution = am
                     .lock()
-                    .iterate_prioritized_random_addresses(HashSet::new())
+                    .iterate_prioritized_random_addresses(vec![], HashSet::new())
                     .take(num_of_buckets)
                     .map(|addr| addr.prefix_bucket().as_u64() as f64)
                     .collect_vec();
@@ -624,5 +658,32 @@ mod address_store_with_cache {
             );
             assert!(adjusted_p <= significance);
         }
+
+        #[test]
+        fn test_perigee_data() {
+            let db = create_temp_db!(ConnBuilder::default().with_files_limit(10));
+            let config = Config::new(SIMNET_PARAMS);
+            let (am, _) = AddressManager::new(Arc::new(config), db.1, Arc::new(TickService::default()));
+
+            let mut am_guard = am.lock();
+
+            let perigee_addresses: Vec<NetAddress> = (0..10)
+                .map(|i| {
+                    NetAddress::new(
+                        IpAddress::from_str(&format!("{}.{}.{}.{}", i, i, i, i)).unwrap(),
+                        SIMNET_PARAMS.default_p2p_port(),
+                    )
+                })
+                .collect();
+            am_guard.set_new_perigee_addresses(perigee_addresses.clone());
+
+            let fetched_perigee_addresses = am_guard.get_perigee_addresses();
+            assert_eq!(fetched_perigee_addresses, perigee_addresses);
+
+            am_guard.reset_perigee_data();
+
+            let fetched_perigee_addresses_after_reset = am_guard.get_perigee_addresses();
+            assert!(fetched_perigee_addresses_after_reset.is_empty());
+        }
     }
 }