From 70c59834ba7d73c0491b2b734801c43f3478acd5 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Mon, 26 Feb 2024 09:47:46 +0000
Subject: [PATCH] Simplify the trait for tracker comms

---
 crates/librqbit/src/session.rs       |  4 +--
 crates/librqbit/src/tracker_comms.rs | 52 ++++++++++++++--------------
 2 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs
index d7bf12c..46c9694 100644
--- a/crates/librqbit/src/session.rs
+++ b/crates/librqbit/src/session.rs
@@ -41,7 +41,7 @@ use crate::{
     torrent_state::{
         ManagedTorrentBuilder, ManagedTorrentHandle, ManagedTorrentState, TorrentStateLive,
     },
-    tracker_comms::{self, TrackerComms},
+    tracker_comms::TrackerComms,
     type_aliases::PeerStream,
 };
 
@@ -1093,7 +1093,7 @@ impl Session {
             self.peer_id,
             trackers,
             // TODO: report actual bytes, not zeroes.
-            Box::new(tracker_comms::TorrentStatsForTrackerDummy {}),
+            Box::new(()),
             force_tracker_interval,
             cancel,
             announce_port,
diff --git a/crates/librqbit/src/tracker_comms.rs b/crates/librqbit/src/tracker_comms.rs
index e5b83ef..d1a7813 100644
--- a/crates/librqbit/src/tracker_comms.rs
+++ b/crates/librqbit/src/tracker_comms.rs
@@ -20,21 +20,24 @@ use librqbit_core::hash_id::Id20;
 pub struct TrackerComms {
     info_hash: Id20,
     peer_id: Id20,
-    stats: Box<dyn TorrentStatsForTracker>,
+    stats: Box<dyn TorrentStatsProvider>,
     force_tracker_interval: Option<Duration>,
     cancellation_token: CancellationToken,
     tx: Sender,
     tcp_listen_port: Option<u16>,
 }
 
-pub trait TorrentStatsForTracker: Send + Sync {
-    fn get_uploaded_bytes(&self) -> u64;
-    fn get_downloaded_bytes(&self) -> u64;
-    fn get_total_bytes(&self) -> u64;
+#[derive(Default)]
+pub struct TrackerCommsStats {
+    pub uploaded_bytes: u64,
+    pub downloaded_bytes: u64,
+    pub total_bytes: u64,
+}
 
-    fn get_left_to_download_bytes(&self) -> u64 {
-        let total = self.get_total_bytes();
-        let down = self.get_downloaded_bytes();
+impl TrackerCommsStats {
+    pub fn get_left_to_download_bytes(&self) -> u64 {
+        let total = self.total_bytes;
+        let down = self.downloaded_bytes;
         if total >= down {
             return total - down;
         }
@@ -42,18 +45,13 @@ pub trait TorrentStatsForTracker: Send + Sync {
     }
 }
 
-pub struct TorrentStatsForTrackerDummy {}
-impl TorrentStatsForTracker for TorrentStatsForTrackerDummy {
-    fn get_uploaded_bytes(&self) -> u64 {
-        0
-    }
+pub trait TorrentStatsProvider: Send + Sync {
+    fn get(&self) -> TrackerCommsStats;
+}
 
-    fn get_downloaded_bytes(&self) -> u64 {
-        0
-    }
-
-    fn get_total_bytes(&self) -> u64 {
-        0
+impl TorrentStatsProvider for () {
+    fn get(&self) -> TrackerCommsStats {
+        Default::default()
     }
 }
 
@@ -64,7 +62,7 @@ impl TrackerComms {
         info_hash: Id20,
         peer_id: Id20,
         trackers: Vec<String>,
-        stats: Box<dyn TorrentStatsForTracker>,
+        stats: Box<dyn TorrentStatsProvider>,
         force_interval: Option<Duration>,
         cancellation_token: CancellationToken,
         tcp_listen_port: Option<u16>,
@@ -131,13 +129,14 @@ impl TrackerComms {
     ) -> anyhow::Result<()> {
         let mut event = Some(tracker_comms_http::TrackerRequestEvent::Started);
         loop {
+            let stats = self.stats.get();
             let request = tracker_comms_http::TrackerRequest {
                 info_hash: self.info_hash,
                 peer_id: self.peer_id,
                 port: 6778,
-                uploaded: self.stats.get_uploaded_bytes(),
-                downloaded: self.stats.get_downloaded_bytes(),
-                left: self.stats.get_left_to_download_bytes(),
+                uploaded: stats.uploaded_bytes,
+                downloaded: stats.downloaded_bytes,
+                left: stats.get_left_to_download_bytes(),
                 compact: true,
                 no_peer_id: false,
                 event,
@@ -212,12 +211,13 @@ impl TrackerComms {
                 tokio::time::sleep(i).await;
             }
 
+            let stats = self.stats.get();
             let request = AnnounceFields {
                 info_hash: self.info_hash,
                 peer_id: self.peer_id,
-                downloaded: self.stats.get_downloaded_bytes(),
-                left: self.stats.get_left_to_download_bytes(),
-                uploaded: self.stats.get_uploaded_bytes(),
+                downloaded: stats.downloaded_bytes,
+                left: stats.get_left_to_download_bytes(),
+                uploaded: stats.uploaded_bytes,
                 event: EVENT_NONE,
                 key: 0, // whatever that is?
                 port: self.tcp_listen_port.unwrap_or(0),