From 1e525eb91e1ecd568f8956eb30007ee6bfcbab3f Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Mon, 25 Dec 2023 08:58:42 -0500
Subject: [PATCH] Move per-peer piece stats into methods

---
 crates/librqbit/src/torrent_state/live/mod.rs |  7 +----
 .../torrent_state/live/peer/stats/atomic.rs   | 26 ++++++++++++++++++-
 .../torrent_state/live/peer/stats/snapshot.rs |  2 ++
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs
index c439df2..411e817 100644
--- a/crates/librqbit/src/torrent_state/live/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/mod.rs
@@ -1474,12 +1474,7 @@ impl PeerHandler {
 
                         // Per-peer piece counters.
                         self.counters
-                            .downloaded_and_checked_pieces
-                            .fetch_add(1, Ordering::Relaxed);
-                        self.counters
-                            .downloaded_and_checked_bytes
-                            .fetch_add(piece_len, Ordering::Relaxed);
-
+                            .on_piece_downloaded(piece_len, full_piece_download_time);
                         self.state.peers.reset_peer_backoff(self.addr);
 
                         debug!("piece={} successfully downloaded and verified", index);
diff --git a/crates/librqbit/src/torrent_state/live/peer/stats/atomic.rs b/crates/librqbit/src/torrent_state/live/peer/stats/atomic.rs
index 2933d07..5ab7cda 100644
--- a/crates/librqbit/src/torrent_state/live/peer/stats/atomic.rs
+++ b/crates/librqbit/src/torrent_state/live/peer/stats/atomic.rs
@@ -1,6 +1,6 @@
 use std::{
     sync::{
-        atomic::{AtomicU32, AtomicU64},
+        atomic::{AtomicU32, AtomicU64, Ordering},
         Arc,
     },
     time::Duration,
@@ -19,6 +19,30 @@ pub(crate) struct PeerCountersAtomic {
     pub fetched_chunks: AtomicU32,
     pub downloaded_and_checked_pieces: AtomicU32,
     pub downloaded_and_checked_bytes: AtomicU64,
+    pub total_piece_download_ms: AtomicU64,
+}
+
+impl PeerCountersAtomic {
+    pub(crate) fn on_piece_downloaded(&self, piece_len: u64, elapsed: Duration) {
+        let elapsed = elapsed.as_millis() as u64;
+        self.total_piece_download_ms
+            .fetch_add(elapsed, Ordering::Release);
+        self.downloaded_and_checked_pieces
+            .fetch_add(1, Ordering::Release);
+        self.downloaded_and_checked_bytes
+            .fetch_add(piece_len, Ordering::Relaxed);
+    }
+
+    pub(crate) fn average_piece_download_time(&self) -> Option<Duration> {
+        let downloaded_pieces = self.downloaded_and_checked_pieces.load(Ordering::Acquire);
+        let total_download_time = self.total_piece_download_ms.load(Ordering::Acquire);
+        if total_download_time == 0 || downloaded_pieces == 0 {
+            return None;
+        }
+        Some(Duration::from_millis(
+            total_download_time / downloaded_pieces as u64,
+        ))
+    }
 }
 
 #[derive(Debug)]
diff --git a/crates/librqbit/src/torrent_state/live/peer/stats/snapshot.rs b/crates/librqbit/src/torrent_state/live/peer/stats/snapshot.rs
index df18007..1e0b2bd 100644
--- a/crates/librqbit/src/torrent_state/live/peer/stats/snapshot.rs
+++ b/crates/librqbit/src/torrent_state/live/peer/stats/snapshot.rs
@@ -14,6 +14,7 @@ pub struct PeerCounters {
     pub errors: u32,
     pub fetched_chunks: u32,
     pub downloaded_and_checked_pieces: u32,
+    pub total_piece_download_ms: u64,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -37,6 +38,7 @@ impl From<&super::atomic::PeerCountersAtomic> for PeerCounters {
             downloaded_and_checked_pieces: counters
                 .downloaded_and_checked_pieces
                 .load(Ordering::Relaxed),
+            total_piece_download_ms: counters.total_piece_download_ms.load(Ordering::Relaxed),
         }
     }
 }