wip/dnm: Prioritize authority records from DHT

src/lib.rs

@@ -505,6 +505,15 @@ impl Litep2p {
         self.transport_manager.add_known_address(peer, address)
     }
 
+    /// Add one ore more known addresses for peer with priority.
+    pub fn add_priority_address(
+        &mut self,
+        peer: PeerId,
+        address: impl Iterator<Item = Multiaddr>,
+    ) -> usize {
+        self.transport_manager.add_priority_address(peer, address)
+    }
+
     /// Poll next event.
     ///
     /// This function must be called in order for litep2p to make progress.

src/transport/manager/address.rs

@@ -24,11 +24,14 @@ use ip_network::IpNetwork;
 use multiaddr::{Multiaddr, Protocol};
 use multihash::Multihash;
 
-use std::collections::{hash_map::Entry, HashMap};
+use std::collections::{hash_map::Entry, HashMap, VecDeque};
 
 /// Maximum number of addresses tracked for a peer.
 const MAX_ADDRESSES: usize = 64;
 
+/// Maximum number of priority addresses returned by [`AddressStore::addresses`].
+const MAX_PRIORITY_ADDRESS: usize = 4;
+
 /// Scores for address records.
 pub mod scores {
     /// Score indicating that the connection was successfully established.
@@ -173,6 +176,13 @@ impl Ord for AddressRecord {
 pub struct AddressStore {
     /// Addresses available.
     pub addresses: HashMap<Multiaddr, AddressRecord>,
+
+    /// Stateless priority addresses that also ignore the capacity limit of the store.
+    ///
+    /// They represent the authority discovery records from the DHT and should always be
+    /// returned by the store.
+    pub priority_addresses: VecDeque<Multiaddr>,
+
     /// Maximum capacity of the address store.
     max_capacity: usize,
 }
@@ -222,6 +232,7 @@ impl AddressStore {
     pub fn new() -> Self {
         Self {
             addresses: HashMap::with_capacity(MAX_ADDRESSES),
+            priority_addresses: VecDeque::with_capacity(MAX_PRIORITY_ADDRESS),
             max_capacity: MAX_ADDRESSES,
         }
     }
@@ -239,6 +250,25 @@ impl AddressStore {
         self.addresses.is_empty()
     }
 
+    /// Insert the priority addresses into the store.
+    pub fn insert_with_priority(&mut self, records: Vec<Multiaddr>) {
+        for record in records {
+            if self.priority_addresses.contains(&record) {
+                continue;
+            }
+
+            if self.priority_addresses.len() >= MAX_PRIORITY_ADDRESS {
+                if let Some(evicted) = self.priority_addresses.pop_front() {
+                    if let Some(rec) = AddressRecord::from_multiaddr(evicted) {
+                        self.insert(rec);
+                    }
+                }
+            }
+
+            self.priority_addresses.push_back(record);
+        }
+    }
+
     /// Insert the address record into [`AddressStore`] with the provided score.
     ///
     /// If the address is not in the store, it will be inserted with a bonus for public addresses.
@@ -294,9 +324,22 @@ impl AddressStore {
 
     /// Return the available addresses sorted by score.
     pub fn addresses(&self, limit: usize) -> Vec<Multiaddr> {
-        let mut records = self.addresses.values().cloned().collect::<Vec<_>>();
-        records.sort_by_key(|rhs| std::cmp::Reverse(rhs.score));
-        records.into_iter().take(limit).map(|record| record.address).collect()
+        let mut result =
+            Vec::with_capacity(self.priority_addresses.len() + self.addresses.len().min(limit));
+
+        result.extend(self.priority_addresses.iter().cloned());
+
+        let mut candidates: Vec<_> = self
+            .addresses
+            .values()
+            .filter(|record| !self.priority_addresses.contains(record.address()))
+            .collect();
+
+        candidates.sort_unstable_by_key(|rec| std::cmp::Reverse(rec.score));
+
+        result.extend(candidates.into_iter().take(limit).map(|record| record.address().clone()));
+
+        result
     }
 }
 
@@ -617,6 +660,7 @@ mod tests {
         let mut store = AddressStore {
             addresses: HashMap::new(),
             max_capacity: 2,
+            priority_addresses: VecDeque::new(),
         };
 
         let mut rng = rand::thread_rng();

src/transport/manager/handle.rs

@@ -227,16 +227,11 @@ impl TransportManagerHandle {
         false
     }
 
-    /// Add one or more known addresses for peer.
-    ///
-    /// If peer doesn't exist, it will be added to known peers.
-    ///
-    /// Returns the number of added addresses after non-supported transports were filtered out.
-    pub fn add_known_address(
-        &mut self,
+    fn get_valid_addresses(
+        &self,
         peer: &PeerId,
         addresses: impl Iterator<Item = Multiaddr>,
-    ) -> usize {
+    ) -> HashSet<Multiaddr> {
         let mut peer_addresses = HashSet::new();
 
         for address in addresses {
@@ -270,6 +265,20 @@ impl TransportManagerHandle {
             }
         }
 
+        peer_addresses
+    }
+
+    /// Add one or more known addresses for peer.
+    ///
+    /// If peer doesn't exist, it will be added to known peers.
+    ///
+    /// Returns the number of added addresses after non-supported transports were filtered out.
+    pub fn add_known_address(
+        &mut self,
+        peer: &PeerId,
+        addresses: impl Iterator<Item = Multiaddr>,
+    ) -> usize {
+        let peer_addresses = self.get_valid_addresses(peer, addresses);
         let num_added = peer_addresses.len();
 
         tracing::trace!(
@@ -287,7 +296,32 @@ impl TransportManagerHandle {
         entry
             .addresses
             .extend(peer_addresses.into_iter().filter_map(AddressRecord::from_multiaddr));
+        num_added
+    }
+
+    /// Similar to `[Self::add_known_address]` but adds the addresses with priority, meaning they
+    /// will be returned by the store before the other addresses.
+    pub fn add_priority_address(
+        &mut self,
+        peer: &PeerId,
+        addresses: impl Iterator<Item = Multiaddr>,
+    ) -> usize {
+        let peer_addresses = self.get_valid_addresses(peer, addresses);
+        let num_added = peer_addresses.len();
 
+        tracing::trace!(
+            target: LOG_TARGET,
+            ?peer,
+            ?peer_addresses,
+            "add priority addresses",
+        );
+
+        let mut peers = self.peers.write();
+        let entry = peers.entry(*peer).or_default();
+
+        // All addresses should be valid at this point, since the peer ID was either added or
+        // double checked.
+        entry.addresses.insert_with_priority(peer_addresses.into_iter().collect());
         num_added
     }
 

src/transport/manager/mod.rs

@@ -499,6 +499,15 @@ impl TransportManager {
         self.transport_manager_handle.add_known_address(&peer, address)
     }
 
+    /// Add one or more known addresses for `peer` with priority.
+    pub fn add_priority_address(
+        &mut self,
+        peer: PeerId,
+        address: impl Iterator<Item = Multiaddr>,
+    ) -> usize {
+        self.transport_manager_handle.add_priority_address(&peer, address)
+    }
+
     /// Return multiple addresses to dial on supported protocols.
     fn supported_transports_addresses(
         addresses: &[Multiaddr],