Code to serialize peer store

Cargo.lock

@@ -26,6 +26,21 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "android-tzdata"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
+
+[[package]]
+name = "android_system_properties"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "anstream"
 version = "0.6.4"
@@ -275,6 +290,21 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "chrono"
+version = "0.4.31"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f2c685bad3eb3d45a01354cedb7d5faa66194d1d58ba6e267a8de788f79db38"
+dependencies = [
+ "android-tzdata",
+ "iana-time-zone",
+ "js-sys",
+ "num-traits",
+ "serde",
+ "wasm-bindgen",
+ "windows-targets",
+]
+
 [[package]]
 name = "clap"
 version = "4.4.8"
@@ -462,6 +492,7 @@ dependencies = [
  "lock_api",
  "once_cell",
  "parking_lot_core",
+ "serde",
 ]
 
 [[package]]
@@ -886,6 +917,29 @@ dependencies = [
  "tokio-native-tls",
 ]
 
+[[package]]
+name = "iana-time-zone"
+version = "0.1.58"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8326b86b6cff230b97d0d312a6c40a60726df3332e721f72a1b035f451663b20"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "wasm-bindgen",
+ "windows-core",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "idna"
 version = "0.4.0"
@@ -1090,6 +1144,7 @@ version = "3.2.0"
 dependencies = [
  "anyhow",
  "backoff",
+ "chrono",
  "dashmap",
  "directories",
  "futures",
@@ -2515,6 +2570,15 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
+[[package]]
+name = "windows-core"
+version = "0.51.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1f8cf84f35d2db49a46868f947758c7a1138116f7fac3bc844f43ade1292e64"
+dependencies = [
+ "windows-targets",
+]
+
 [[package]]
 name = "windows-sys"
 version = "0.48.0"

crates/dht/Cargo.toml

@@ -32,10 +32,11 @@ futures = "0.3"
 rand = "0.8"
 indexmap = "2"
 directories = "5"
-dashmap = "5.5.3"
+dashmap = {version = "5.5.3", features = ["serde"]}
 
 clone_to_owned = {path="../clone_to_owned", package="librqbit-clone-to-owned", version = "2.2.1"}
 librqbit-core = {path="../librqbit_core", version = "3.1.0"}
+chrono = {version = "0.4.31", features = ["serde"]}
 
 [dev-dependencies]
-tracing-subscriber = "0.3"
+tracing-subscriber = "0.3"

crates/dht/src/dht.rs

@@ -490,7 +490,7 @@ pub struct DhtState {
     // This is to send raw messages
     worker_sender: UnboundedSender<WorkerSendRequest>,
 
-    peer_store: PeerStore,
+    pub(crate) peer_store: PeerStore,
 }
 
 impl DhtState {
@@ -499,6 +499,7 @@ impl DhtState {
         sender: UnboundedSender<WorkerSendRequest>,
         routing_table: Option<RoutingTable>,
         listen_addr: SocketAddr,
+        peer_store: PeerStore,
     ) -> Self {
         let routing_table = routing_table.unwrap_or_else(|| RoutingTable::new(id, None));
         Self {
@@ -509,7 +510,7 @@ impl DhtState {
             worker_sender: sender,
             listen_addr,
             rate_limiter: make_rate_limiter(),
-            peer_store: PeerStore::new(id),
+            peer_store,
         }
     }
 
@@ -1056,6 +1057,7 @@ pub struct DhtConfig {
     pub bootstrap_addrs: Option<Vec<String>>,
     pub routing_table: Option<RoutingTable>,
     pub listen_addr: Option<SocketAddr>,
+    pub(crate) peer_store: Option<PeerStore>,
 }
 
 impl DhtState {
@@ -1089,6 +1091,7 @@ impl DhtState {
             in_tx,
             config.routing_table,
             listen_addr,
+            config.peer_store.unwrap_or_else(|| PeerStore::new(peer_id)),
         ));
 
         spawn(error_span!("dht"), {

crates/dht/src/peer_store.rs

@@ -2,37 +2,108 @@ use std::{
     collections::VecDeque,
     net::{SocketAddr, SocketAddrV4},
     str::FromStr,
-    sync::atomic::AtomicU64,
+    sync::atomic::AtomicU32,
-    time::Instant,
 };
 
 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_tokens: u32,
-    max_remembered_peers: usize,
+    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
+        if !self.tokens.read().iter().any(|t| {
-            .tokens
+            t.token[..] == announce.token[..]
-            .read()
+                && t.addr == std::net::SocketAddr::V4(addr)
-            .iter()
+                && t.node_id == announce.id
-            .any(|t| t.token[..] == announce.token[..] && t.addr == std::net::SocketAddr::V4(addr))
+        }) {
-        {
             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)
+        use dashmap::mapref::entry::Entry;
-            .or_default()
+        let peers_entry = self.peers.entry(announce.info_hash);
-            .push(StoredPeer {
+        let peers_len = self.peers_len.load(std::sync::atomic::Ordering::SeqCst);
-                addr,
+        match peers_entry {
-                time: Instant::now(),
+            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

crates/dht/src/persistence.rs

@@ -11,6 +11,7 @@ use std::time::Duration;
 use anyhow::Context;
 use tracing::{debug, error, error_span, info, trace, warn};
 
+use crate::peer_store::PeerStore;
 use crate::routing_table::RoutingTable;
 use crate::{Dht, DhtConfig, DhtState};
 
@@ -21,9 +22,10 @@ pub struct PersistentDhtConfig {
 }
 
 #[derive(Serialize, Deserialize)]
-struct DhtSerialize<Table> {
+struct DhtSerialize<Table, PeerStore> {
     addr: SocketAddr,
     table: Table,
+    peer_store: Option<PeerStore>,
 }
 
 pub struct PersistentDht {
@@ -40,9 +42,16 @@ fn dump_dht(dht: &Dht, filename: &Path, tempfile_name: &Path) -> anyhow::Result<
     let mut file = BufWriter::new(file);
 
     let addr = dht.listen_addr();
-    match dht
+    match dht.with_routing_table(|r| {
-        .with_routing_table(|r| serde_json::to_writer(&mut file, &DhtSerialize { addr, table: r }))
+        serde_json::to_writer(
-    {
+            &mut file,
+            &DhtSerialize {
+                addr,
+                table: r,
+                peer_store: Some(&dht.peer_store),
+            },
+        )
+    }) {
         Ok(_) => {
             trace!("dumped DHT to {:?}", &tempfile_name);
         }
@@ -79,7 +88,7 @@ impl PersistentDht {
         let de = match OpenOptions::new().read(true).open(&config_filename) {
             Ok(dht_json) => {
                 let reader = BufReader::new(dht_json);
-                match serde_json::from_reader::<_, DhtSerialize<RoutingTable>>(reader) {
+                match serde_json::from_reader::<_, DhtSerialize<RoutingTable, PeerStore>>(reader) {
                     Ok(r) => {
                         info!("loaded DHT routing table from {:?}", &config_filename);
                         Some(r)
@@ -98,14 +107,15 @@ impl PersistentDht {
                 _ => return Err(e).with_context(|| format!("error reading {config_filename:?}")),
             },
         };
-        let (listen_addr, routing_table) = de
+        let (listen_addr, routing_table, peer_store) = de
-            .map(|de| (Some(de.addr), Some(de.table)))
+            .map(|de| (Some(de.addr), Some(de.table), de.peer_store))
-            .unwrap_or((None, None));
+            .unwrap_or((None, None, None));
         let peer_id = routing_table.as_ref().map(|r| r.id());
         let dht_config = DhtConfig {
             peer_id,
             routing_table,
             listen_addr,
+            peer_store,
             ..Default::default()
         };
         let dht = DhtState::with_config(dht_config).await?;