From 0fdf6ad429f940d65a7a83b43a1842516a9a2b60 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Wed, 21 Aug 2024 11:31:36 +0100
Subject: [PATCH] Counting peers on drop

---
 crates/librqbit/src/torrent_state/live/peer/mod.rs  | 6 ++++++
 crates/librqbit/src/torrent_state/live/peers/mod.rs | 8 ++++++++
 2 files changed, 14 insertions(+)

diff --git a/crates/librqbit/src/torrent_state/live/peer/mod.rs b/crates/librqbit/src/torrent_state/live/peer/mod.rs
index 0c6b0ca..9efd6b0 100644
--- a/crates/librqbit/src/torrent_state/live/peer/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/peer/mod.rs
@@ -91,6 +91,12 @@ impl PeerStateNoMut {
         self.set(Default::default(), counters)
     }
 
+    pub fn destroy(self, counters: &[&AggregatePeerStatsAtomic]) {
+        for counter in counters {
+            counter.dec(&self.0);
+        }
+    }
+
     pub fn set(&mut self, new: PeerState, counters: &[&AggregatePeerStatsAtomic]) -> PeerState {
         for counter in counters {
             counter.incdec(&self.0, &new);
diff --git a/crates/librqbit/src/torrent_state/live/peers/mod.rs b/crates/librqbit/src/torrent_state/live/peers/mod.rs
index fc4d6ea..a66c69b 100644
--- a/crates/librqbit/src/torrent_state/live/peers/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/peers/mod.rs
@@ -25,6 +25,14 @@ pub(crate) struct PeerStates {
     pub states: DashMap<PeerHandle, Peer>,
 }
 
+impl Drop for PeerStates {
+    fn drop(&mut self) {
+        for (_, p) in std::mem::take(&mut self.states).into_iter() {
+            p.state.destroy(&[&self.session_stats.peers]);
+        }
+    }
+}
+
 impl PeerStates {
     pub fn stats(&self) -> AggregatePeerStats {
         AggregatePeerStats::from(&self.stats)