Code to serialize peer store

crates/dht/src/peer_store.rs

@@ -2,37 +2,108 @@ use std::{
     collections::VecDeque,
     net::{SocketAddr, SocketAddrV4},
     str::FromStr,
-    sync::atomic::AtomicU64,
-    time::Instant,
+    sync::atomic::AtomicU32,
 };
 
 use bencode::ByteString;
+use chrono::{DateTime, Utc};
 use librqbit_core::id20::Id20;
 use parking_lot::RwLock;
 use rand::RngCore;
+use serde::{
+    ser::{SerializeMap, SerializeStruct},
+    Deserialize, Serialize,
+};
 use tracing::trace;
 
-use crate::bprotocol::{AnnouncePeer, CompactPeerInfo, Response};
+use crate::bprotocol::{AnnouncePeer, CompactPeerInfo};
 
+#[derive(Serialize, Deserialize)]
 struct StoredToken {
     token: [u8; 4],
+    #[serde(serialize_with = "crate::utils::serialize_id20")]
     node_id: Id20,
     addr: SocketAddr,
 }
 
+#[derive(Serialize, Deserialize)]
 struct StoredPeer {
     addr: SocketAddrV4,
-    time: Instant,
+    time: DateTime<Utc>,
 }
 
 pub struct PeerStore {
     self_id: Id20,
-    max_remembered_tokens: usize,
-    max_remembered_peers: usize,
+    max_remembered_tokens: u32,
+    max_remembered_peers: u32,
     max_distance: Id20,
     tokens: RwLock<VecDeque<StoredToken>>,
     peers: dashmap::DashMap<Id20, Vec<StoredPeer>>,
-    peers_len: AtomicU64,
+    peers_len: AtomicU32,
+}
+
+impl Serialize for PeerStore {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        struct SerializePeers<'a> {
+            peers: &'a dashmap::DashMap<Id20, Vec<StoredPeer>>,
+        }
+
+        impl<'a> Serialize for SerializePeers<'a> {
+            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+            where
+                S: serde::Serializer,
+            {
+                let mut m = serializer.serialize_map(None)?;
+                for entry in self.peers.iter() {
+                    m.serialize_entry(&entry.key().as_string(), &entry.value())?;
+                }
+                m.end()
+            }
+        }
+
+        let mut s = serializer.serialize_struct("PeerStore", 7)?;
+        s.serialize_field("self_id", &self.self_id.as_string())?;
+        s.serialize_field("max_remembered_tokens", &self.max_remembered_tokens)?;
+        s.serialize_field("max_remembered_peers", &self.max_remembered_peers)?;
+        s.serialize_field("max_distance", &self.max_distance.as_string())?;
+        s.serialize_field("tokens", &*self.tokens.read())?;
+        s.serialize_field("peers", &SerializePeers { peers: &self.peers })?;
+        s.serialize_field(
+            "peers_len",
+            &self.peers_len.load(std::sync::atomic::Ordering::SeqCst),
+        )?;
+        s.end()
+    }
+}
+
+impl<'de> Deserialize<'de> for PeerStore {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        #[derive(Deserialize)]
+        struct Tmp {
+            self_id: Id20,
+            max_remembered_tokens: u32,
+            max_remembered_peers: u32,
+            max_distance: Id20,
+            tokens: VecDeque<StoredToken>,
+            peers: dashmap::DashMap<Id20, Vec<StoredPeer>>,
+        }
+
+        Tmp::deserialize(deserializer).map(|tmp| Self {
+            self_id: tmp.self_id,
+            max_remembered_tokens: tmp.max_remembered_tokens,
+            max_remembered_peers: tmp.max_remembered_peers,
+            max_distance: tmp.max_distance,
+            tokens: RwLock::new(tmp.tokens),
+            peers_len: AtomicU32::new(tmp.peers.iter().map(|e| e.value().len() as u32).sum()),
+            peers: tmp.peers,
+        })
+    }
 }
 
 impl PeerStore {
@@ -44,7 +115,7 @@ impl PeerStore {
             max_distance: Id20::from_str("00000fffffffffffffffffffffffffffffffffff").unwrap(),
             tokens: RwLock::new(VecDeque::new()),
             peers: dashmap::DashMap::new(),
-            peers_len: AtomicU64::new(0),
+            peers_len: AtomicU32::new(0),
         }
     }
 
@@ -54,10 +125,10 @@ impl PeerStore {
         let mut tokens = self.tokens.write();
         tokens.push_back(StoredToken {
             token,
-            node_id,
             addr,
+            node_id,
         });
-        if tokens.len() > self.max_remembered_tokens {
+        if tokens.len() > self.max_remembered_tokens as usize {
             tokens.pop_front();
         }
         token
@@ -75,36 +146,54 @@ impl PeerStore {
                 return false;
             }
         };
-        if self.peers_len.load(std::sync::atomic::Ordering::SeqCst)
-            >= self.max_remembered_peers as u64
-        {
-            trace!("peer store: out of capacity");
-            return false;
-        }
 
         if announce.info_hash.distance(&self.self_id) > self.max_distance {
             trace!("peer store: info_hash too far to store");
             return false;
         }
-        if !self
-            .tokens
-            .read()
-            .iter()
-            .any(|t| t.token[..] == announce.token[..] && t.addr == std::net::SocketAddr::V4(addr))
-        {
+        if !self.tokens.read().iter().any(|t| {
+            t.token[..] == announce.token[..]
+                && t.addr == std::net::SocketAddr::V4(addr)
+                && t.node_id == announce.id
+        }) {
             trace!("peer store: can't find this token / addr combination");
             return false;
         }
+
         if announce.implied_port == 0 {
             addr.set_port(announce.port);
         }
-        self.peers
-            .entry(announce.info_hash)
-            .or_default()
-            .push(StoredPeer {
-                addr,
-                time: Instant::now(),
-            });
+
+        use dashmap::mapref::entry::Entry;
+        let peers_entry = self.peers.entry(announce.info_hash);
+        let peers_len = self.peers_len.load(std::sync::atomic::Ordering::SeqCst);
+        match peers_entry {
+            Entry::Occupied(mut occ) => {
+                if let Some(s) = occ.get_mut().iter_mut().find(|s| s.addr == addr) {
+                    s.time = Utc::now();
+                    return true;
+                }
+                if peers_len >= self.max_remembered_peers {
+                    trace!("peer store: out of capacity");
+                    return false;
+                }
+                occ.get_mut().push(StoredPeer {
+                    addr,
+                    time: Utc::now(),
+                });
+            }
+            Entry::Vacant(vac) => {
+                if peers_len >= self.max_remembered_peers {
+                    trace!("peer store: out of capacity");
+                    return false;
+                }
+                vac.insert(vec![StoredPeer {
+                    addr,
+                    time: Utc::now(),
+                }]);
+            }
+        }
+
         self.peers_len
             .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
         true