split up librqbit torrent state into smaller files

crates/librqbit/src/torrent_state/live/peers/mod.rs

@@ -0,0 +1,114 @@
+use std::net::SocketAddr;
+
+use anyhow::Context;
+use backoff::backoff::Backoff;
+use dashmap::DashMap;
+
+use crate::{
+    torrent_state::utils::{atomic_inc, TimedExistence},
+    type_aliases::{PeerHandle, BF},
+};
+
+use self::stats::{atomic::AggregatePeerStatsAtomic, snapshot::AggregatePeerStats};
+
+use super::peer::{LivePeerState, Peer, PeerRx, PeerState, PeerTx};
+
+pub mod stats;
+
+#[derive(Default)]
+pub struct PeerStates {
+    pub stats: AggregatePeerStatsAtomic,
+    pub states: DashMap<PeerHandle, Peer>,
+}
+
+impl PeerStates {
+    pub fn stats(&self) -> AggregatePeerStats {
+        AggregatePeerStats::from(&self.stats)
+    }
+
+    pub fn add_if_not_seen(&self, addr: SocketAddr) -> Option<PeerHandle> {
+        use dashmap::mapref::entry::Entry;
+        match self.states.entry(addr) {
+            Entry::Occupied(_) => None,
+            Entry::Vacant(vac) => {
+                vac.insert(Default::default());
+                atomic_inc(&self.stats.queued);
+                atomic_inc(&self.stats.seen);
+                Some(addr)
+            }
+        }
+    }
+    pub fn with_peer<R>(&self, addr: PeerHandle, f: impl FnOnce(&Peer) -> R) -> Option<R> {
+        self.states.get(&addr).map(|e| f(e.value()))
+    }
+
+    pub fn with_peer_mut<R>(
+        &self,
+        addr: PeerHandle,
+        reason: &'static str,
+        f: impl FnOnce(&mut Peer) -> R,
+    ) -> Option<R> {
+        use crate::torrent_state::utils::timeit;
+        timeit(reason, || self.states.get_mut(&addr))
+            .map(|e| f(TimedExistence::new(e, reason).value_mut()))
+    }
+    pub fn with_live<R>(&self, addr: PeerHandle, f: impl FnOnce(&LivePeerState) -> R) -> Option<R> {
+        self.states
+            .get(&addr)
+            .and_then(|e| match &e.value().state.get() {
+                PeerState::Live(l) => Some(f(l)),
+                _ => None,
+            })
+    }
+    pub fn with_live_mut<R>(
+        &self,
+        addr: PeerHandle,
+        reason: &'static str,
+        f: impl FnOnce(&mut LivePeerState) -> R,
+    ) -> Option<R> {
+        self.with_peer_mut(addr, reason, |peer| peer.state.get_live_mut().map(f))
+            .flatten()
+    }
+
+    pub fn drop_peer(&self, handle: PeerHandle) -> Option<Peer> {
+        let p = self.states.remove(&handle).map(|r| r.1)?;
+        self.stats.dec(p.state.get());
+        Some(p)
+    }
+
+    pub fn mark_peer_interested(&self, handle: PeerHandle, is_interested: bool) -> Option<bool> {
+        self.with_live_mut(handle, "mark_peer_interested", |live| {
+            let prev = live.peer_interested;
+            live.peer_interested = is_interested;
+            prev
+        })
+    }
+    pub fn update_bitfield_from_vec(&self, handle: PeerHandle, bitfield: Vec<u8>) -> Option<()> {
+        self.with_live_mut(handle, "update_bitfield_from_vec", |live| {
+            live.bitfield = BF::from_vec(bitfield);
+        })
+    }
+    pub fn mark_peer_connecting(&self, h: PeerHandle) -> anyhow::Result<(PeerRx, PeerTx)> {
+        let rx = self
+            .with_peer_mut(h, "mark_peer_connecting", |peer| {
+                peer.state
+                    .queued_to_connecting(&self.stats)
+                    .context("invalid peer state")
+            })
+            .context("peer not found in states")??;
+        Ok(rx)
+    }
+
+    pub fn reset_peer_backoff(&self, handle: PeerHandle) {
+        self.with_peer_mut(handle, "reset_peer_backoff", |p| {
+            p.stats.backoff.reset();
+        });
+    }
+
+    pub fn mark_peer_not_needed(&self, handle: PeerHandle) -> Option<PeerState> {
+        let prev = self.with_peer_mut(handle, "mark_peer_not_needed", |peer| {
+            peer.state.to_not_needed(&self.stats)
+        })?;
+        Some(prev)
+    }
+}

crates/librqbit/src/torrent_state/live/peers/stats/atomic.rs

@@ -0,0 +1,43 @@
+use std::sync::atomic::AtomicU32;
+
+use serde::Serialize;
+
+use crate::torrent_state::{
+    live::peer::PeerState,
+    utils::{atomic_dec, atomic_inc},
+};
+
+#[derive(Debug, Default, Serialize)]
+pub struct AggregatePeerStatsAtomic {
+    pub queued: AtomicU32,
+    pub connecting: AtomicU32,
+    pub live: AtomicU32,
+    pub seen: AtomicU32,
+    pub dead: AtomicU32,
+    pub not_needed: AtomicU32,
+}
+
+impl AggregatePeerStatsAtomic {
+    pub fn counter(&self, state: &PeerState) -> &AtomicU32 {
+        match state {
+            PeerState::Connecting(_) => &self.connecting,
+            PeerState::Live(_) => &self.live,
+            PeerState::Queued => &self.queued,
+            PeerState::Dead => &self.dead,
+            PeerState::NotNeeded => &self.not_needed,
+        }
+    }
+
+    pub fn inc(&self, state: &PeerState) {
+        atomic_inc(self.counter(state));
+    }
+
+    pub fn dec(&self, state: &PeerState) {
+        atomic_dec(self.counter(state));
+    }
+
+    pub fn incdec(&self, old: &PeerState, new: &PeerState) {
+        self.dec(old);
+        self.inc(new);
+    }
+}

crates/librqbit/src/torrent_state/live/peers/stats/mod.rs

@@ -0,0 +1,2 @@
+pub mod atomic;
+pub mod snapshot;

crates/librqbit/src/torrent_state/live/peers/stats/snapshot.rs

@@ -0,0 +1,29 @@
+use std::sync::atomic::Ordering;
+
+use serde::Serialize;
+
+use super::atomic::AggregatePeerStatsAtomic;
+
+#[derive(Debug, Default, Serialize, PartialEq, Eq)]
+pub struct AggregatePeerStats {
+    pub queued: usize,
+    pub connecting: usize,
+    pub live: usize,
+    pub seen: usize,
+    pub dead: usize,
+    pub not_needed: usize,
+}
+
+impl<'a> From<&'a AggregatePeerStatsAtomic> for AggregatePeerStats {
+    fn from(s: &'a AggregatePeerStatsAtomic) -> Self {
+        let ordering = Ordering::Relaxed;
+        Self {
+            queued: s.queued.load(ordering) as usize,
+            connecting: s.connecting.load(ordering) as usize,
+            live: s.live.load(ordering) as usize,
+            seen: s.seen.load(ordering) as usize,
+            dead: s.dead.load(ordering) as usize,
+            not_needed: s.not_needed.load(ordering) as usize,
+        }
+    }
+}