split up librqbit torrent state into smaller files

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

@@ -0,0 +1,199 @@
+pub mod stats;
+
+use std::collections::HashSet;
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::Context;
+use backoff::{ExponentialBackoff, ExponentialBackoffBuilder};
+use librqbit_core::id20::Id20;
+use librqbit_core::lengths::{ChunkInfo, ValidPieceIndex};
+use serde::Serialize;
+use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
+
+use crate::peer_connection::WriterRequest;
+use crate::type_aliases::BF;
+
+use super::peers::stats::atomic::AggregatePeerStatsAtomic;
+
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct InflightRequest {
+    pub piece: ValidPieceIndex,
+    pub chunk: u32,
+}
+
+impl From<&ChunkInfo> for InflightRequest {
+    fn from(c: &ChunkInfo) -> Self {
+        Self {
+            piece: c.piece_index,
+            chunk: c.chunk_index,
+        }
+    }
+}
+
+// TODO: Arc can be removed probably, as UnboundedSender should be clone + it can be downgraded to weak.
+pub type PeerRx = UnboundedReceiver<WriterRequest>;
+pub type PeerTx = UnboundedSender<WriterRequest>;
+
+pub trait SendMany {
+    fn send_many(&self, requests: impl IntoIterator<Item = WriterRequest>) -> anyhow::Result<()>;
+}
+
+impl SendMany for PeerTx {
+    fn send_many(&self, requests: impl IntoIterator<Item = WriterRequest>) -> anyhow::Result<()> {
+        requests
+            .into_iter()
+            .try_for_each(|r| self.send(r))
+            .context("peer dropped")
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct Peer {
+    pub state: PeerStateNoMut,
+    pub stats: stats::atomic::PeerStats,
+}
+
+#[derive(Debug, Default)]
+pub enum PeerState {
+    #[default]
+    // Will be tried to be connected as soon as possible.
+    Queued,
+    Connecting(PeerTx),
+    Live(LivePeerState),
+    // There was an error, and it's waiting for exponential backoff.
+    Dead,
+    // We don't need to do anything with the peer any longer.
+    // The peer has the full torrent, and we have the full torrent, so no need
+    // to keep talking to it.
+    NotNeeded,
+}
+
+impl std::fmt::Display for PeerState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}
+
+impl PeerState {
+    pub fn name(&self) -> &'static str {
+        match self {
+            PeerState::Queued => "queued",
+            PeerState::Connecting(_) => "connecting",
+            PeerState::Live(_) => "live",
+            PeerState::Dead => "dead",
+            PeerState::NotNeeded => "not needed",
+        }
+    }
+
+    pub fn take_live_no_counters(self) -> Option<LivePeerState> {
+        match self {
+            PeerState::Live(l) => Some(l),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct PeerStateNoMut(PeerState);
+
+impl PeerStateNoMut {
+    pub fn get(&self) -> &PeerState {
+        &self.0
+    }
+
+    pub fn take(&mut self, counters: &AggregatePeerStatsAtomic) -> PeerState {
+        self.set(Default::default(), counters)
+    }
+
+    pub fn set(&mut self, new: PeerState, counters: &AggregatePeerStatsAtomic) -> PeerState {
+        counters.incdec(&self.0, &new);
+        std::mem::replace(&mut self.0, new)
+    }
+
+    pub fn get_live(&self) -> Option<&LivePeerState> {
+        match &self.0 {
+            PeerState::Live(l) => Some(l),
+            _ => None,
+        }
+    }
+
+    pub fn get_live_mut(&mut self) -> Option<&mut LivePeerState> {
+        match &mut self.0 {
+            PeerState::Live(l) => Some(l),
+            _ => None,
+        }
+    }
+
+    pub fn queued_to_connecting(
+        &mut self,
+        counters: &AggregatePeerStatsAtomic,
+    ) -> Option<(PeerRx, PeerTx)> {
+        if let PeerState::Queued = &self.0 {
+            let (tx, rx) = unbounded_channel();
+            let tx_2 = tx.clone();
+            self.set(PeerState::Connecting(tx), counters);
+            Some((rx, tx_2))
+        } else {
+            None
+        }
+    }
+    pub fn connecting_to_live(
+        &mut self,
+        peer_id: Id20,
+        counters: &AggregatePeerStatsAtomic,
+    ) -> Option<&mut LivePeerState> {
+        if let PeerState::Connecting(_) = &self.0 {
+            let tx = match self.take(counters) {
+                PeerState::Connecting(tx) => tx,
+                _ => unreachable!(),
+            };
+            self.set(PeerState::Live(LivePeerState::new(peer_id, tx)), counters);
+            self.get_live_mut()
+        } else {
+            None
+        }
+    }
+
+    pub fn to_dead(&mut self, counters: &AggregatePeerStatsAtomic) -> PeerState {
+        self.set(PeerState::Dead, counters)
+    }
+
+    pub fn to_not_needed(&mut self, counters: &AggregatePeerStatsAtomic) -> PeerState {
+        self.set(PeerState::NotNeeded, counters)
+    }
+}
+
+#[derive(Debug)]
+pub struct LivePeerState {
+    pub peer_id: Id20,
+    pub peer_interested: bool,
+
+    // This is used to track the pieces the peer has.
+    pub bitfield: BF,
+
+    // When the peer sends us data this is used to track if we asked for it.
+    pub inflight_requests: HashSet<InflightRequest>,
+
+    // The main channel to send requests to peer.
+    pub tx: PeerTx,
+}
+
+impl LivePeerState {
+    pub fn new(peer_id: Id20, tx: PeerTx) -> Self {
+        LivePeerState {
+            peer_id,
+            peer_interested: false,
+            bitfield: BF::new(),
+            inflight_requests: Default::default(),
+            tx,
+        }
+    }
+
+    pub fn has_full_torrent(&self, total_pieces: usize) -> bool {
+        self.bitfield
+            .get(0..total_pieces)
+            .map_or(false, |s| s.all())
+    }
+}

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

@@ -0,0 +1,41 @@
+use std::{
+    sync::{
+        atomic::{AtomicU32, AtomicU64},
+        Arc,
+    },
+    time::Duration,
+};
+
+use backoff::{ExponentialBackoff, ExponentialBackoffBuilder};
+
+#[derive(Default, Debug)]
+pub struct PeerCounters {
+    pub fetched_bytes: AtomicU64,
+    pub total_time_connecting_ms: AtomicU64,
+    pub connection_attempts: AtomicU32,
+    pub connections: AtomicU32,
+    pub errors: AtomicU32,
+    pub fetched_chunks: AtomicU32,
+    pub downloaded_and_checked_pieces: AtomicU32,
+    pub downloaded_and_checked_bytes: AtomicU64,
+}
+
+#[derive(Debug)]
+pub struct PeerStats {
+    pub counters: Arc<PeerCounters>,
+    pub backoff: ExponentialBackoff,
+}
+
+impl Default for PeerStats {
+    fn default() -> Self {
+        Self {
+            counters: Arc::new(Default::default()),
+            backoff: ExponentialBackoffBuilder::new()
+                .with_initial_interval(Duration::from_secs(10))
+                .with_multiplier(6.)
+                .with_max_interval(Duration::from_secs(3600))
+                .with_max_elapsed_time(Some(Duration::from_secs(86400)))
+                .build(),
+        }
+    }
+}

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

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

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

@@ -0,0 +1,74 @@
+use std::{collections::HashMap, sync::atomic::Ordering};
+
+use serde::{Deserialize, Serialize};
+
+use crate::torrent_state::live::peer::{Peer, PeerState};
+
+#[derive(Serialize, Deserialize)]
+pub struct PeerCounters {
+    pub fetched_bytes: u64,
+    pub total_time_connecting_ms: u64,
+    pub connection_attempts: u32,
+    pub connections: u32,
+    pub errors: u32,
+    pub fetched_chunks: u32,
+    pub downloaded_and_checked_pieces: u32,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct PeerStats {
+    pub counters: PeerCounters,
+    pub state: &'static str,
+}
+
+impl From<&super::atomic::PeerCounters> for PeerCounters {
+    fn from(counters: &super::atomic::PeerCounters) -> Self {
+        Self {
+            fetched_bytes: counters.fetched_bytes.load(Ordering::Relaxed),
+            total_time_connecting_ms: counters.total_time_connecting_ms.load(Ordering::Relaxed),
+            connection_attempts: counters.connection_attempts.load(Ordering::Relaxed),
+            connections: counters.connections.load(Ordering::Relaxed),
+            errors: counters.errors.load(Ordering::Relaxed),
+            fetched_chunks: counters.fetched_chunks.load(Ordering::Relaxed),
+            downloaded_and_checked_pieces: counters
+                .downloaded_and_checked_pieces
+                .load(Ordering::Relaxed),
+        }
+    }
+}
+
+impl From<&Peer> for PeerStats {
+    fn from(peer: &Peer) -> Self {
+        Self {
+            counters: peer.stats.counters.as_ref().into(),
+            state: peer.state.get().name(),
+        }
+    }
+}
+
+#[derive(Serialize)]
+pub struct PeerStatsSnapshot {
+    pub peers: HashMap<String, PeerStats>,
+}
+
+#[derive(Clone, Copy, Default, Deserialize)]
+pub enum PeerStatsFilterState {
+    All,
+    #[default]
+    Live,
+}
+
+impl PeerStatsFilterState {
+    pub fn matches(&self, s: &PeerState) -> bool {
+        match (self, s) {
+            (Self::All, _) => true,
+            (Self::Live, PeerState::Live(_)) => true,
+            _ => false,
+        }
+    }
+}
+
+#[derive(Default, Deserialize)]
+pub struct PeerStatsFilter {
+    pub state: PeerStatsFilterState,
+}