From 07a5e560b7bfa7643d0ed7f3883692db06b4687e Mon Sep 17 00:00:00 2001
From: Charles Kerr <charles@charleskerr.com>
Date: Thu, 22 Dec 2022 17:43:36 -0600
Subject: [PATCH] fixup: peer counts (#4445)

---
 libtransmission/peer-common.h | 10 ----------
 libtransmission/peer-mgr.cc   | 30 +++++++++++++-----------------
 libtransmission/peer-msgs.cc  |  6 ++++++
 libtransmission/peer-msgs.h   | 12 ++++++++++++
 4 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/libtransmission/peer-common.h b/libtransmission/peer-common.h
index ca242a5e3..16702741e 100644
--- a/libtransmission/peer-common.h
+++ b/libtransmission/peer-common.h
@@ -10,7 +10,6 @@
 #endif
 
 #include <array>
-#include <atomic>
 #include <cstdint> // uint8_t, uint32_t, uint64_t
 
 #include "transmission.h"
@@ -214,12 +213,6 @@ public:
 
     virtual void requestBlocks(tr_block_span_t const* block_spans, size_t n_spans) = 0;
 
-    [[nodiscard]] static auto peer_count() noexcept
-    {
-        // the number of currently-connected peers
-        return n_peers_.load();
-    }
-
     struct RequestLimit
     {
         // How many blocks we could request.
@@ -261,9 +254,6 @@ public:
 
     // how many requests we made to this peer and then canceled
     tr_recentHistory<uint16_t> cancels_sent_to_peer;
-
-private:
-    static inline auto n_peers_ = std::atomic<size_t>{};
 };
 
 /***
diff --git a/libtransmission/peer-mgr.cc b/libtransmission/peer-mgr.cc
index b6b2b7bab..fe5b5109f 100644
--- a/libtransmission/peer-mgr.cc
+++ b/libtransmission/peer-mgr.cc
@@ -599,7 +599,6 @@ tr_peer::tr_peer(tr_torrent const* tor, peer_atom* atom_in)
     , atom{ atom_in }
     , blame{ tor->blockCount() }
 {
-    ++n_peers_;
 }
 
 tr_peer::~tr_peer()
@@ -613,9 +612,6 @@ tr_peer::~tr_peer()
     {
         atom->is_connected = false;
     }
-
-    [[maybe_unused]] auto const n_prev = n_peers_--;
-    TR_ASSERT(n_prev > 0U);
 }
 
 /**
@@ -2369,10 +2365,9 @@ void closeBadPeers(tr_swarm* s, time_t const now_sec)
     }
 }
 
-void enforceTorrentPeerLimit(tr_swarm* swarm)
+void enforceSwarmPeerLimit(tr_swarm* swarm, size_t max)
 {
     // do we have too many peers?
-    auto const max = swarm->tor->peerLimit();
     if (auto const n = swarm->peerCount(); n <= max)
     {
         return;
@@ -2386,25 +2381,27 @@ void enforceTorrentPeerLimit(tr_swarm* swarm)
 
 void enforceSessionPeerLimit(tr_session* session)
 {
-    // do we have too many peers?
-    auto const n_peers = tr_peer::peer_count();
+    // No need to disconnect if we are under the peer limit
     auto const max = session->peerLimit();
-    if (n_peers <= max)
+    if (tr_peerMsgs::size() <= max)
     {
         return;
     }
 
-    // make a list of all the peers
+    // Make a list of all the peers.
     auto peers = std::vector<tr_peer*>{};
-    peers.reserve(n_peers);
+    peers.reserve(tr_peerMsgs::size());
     for (auto const* const tor : session->torrents())
     {
         peers.insert(std::end(peers), std::begin(tor->swarm->peers), std::end(tor->swarm->peers));
     }
 
-    // close all but the `max` most active
-    std::partial_sort(std::begin(peers), std::begin(peers) + max, std::end(peers), ComparePeerByActivity{});
-    std::for_each(std::begin(peers) + max, std::end(peers), closePeer);
+    TR_ASSERT(tr_peerMsgs::size() == std::size(peers));
+    if (std::size(peers) > max)
+    {
+        std::partial_sort(std::begin(peers), std::begin(peers) + max, std::end(peers), ComparePeerByActivity{});
+        std::for_each(std::begin(peers) + max, std::end(peers), closePeer);
+    }
 }
 
 } // namespace disconnect_helpers
@@ -2436,7 +2433,7 @@ void tr_peerMgr::reconnectPulse()
     {
         if (tor->isRunning)
         {
-            enforceTorrentPeerLimit(tor->swarm);
+            enforceSwarmPeerLimit(tor->swarm, tor->peerLimit());
         }
     }
 
@@ -2654,8 +2651,7 @@ struct peer_candidate
     auto const max_candidates = static_cast<size_t>(session->peerLimit() * 0.95);
 
     // don't start any new handshakes if we're full up
-    auto const peer_count = tr_peer::peer_count();
-    if (max_candidates <= peer_count)
+    if (max_candidates <= tr_peerMsgs::size())
     {
         return {};
     }
diff --git a/libtransmission/peer-msgs.cc b/libtransmission/peer-msgs.cc
index e89c9c82a..0a9027c20 100644
--- a/libtransmission/peer-msgs.cc
+++ b/libtransmission/peer-msgs.cc
@@ -231,6 +231,12 @@ static void updateDesiredRequestCount(tr_peerMsgsImpl* msgs);
 #define logdbg(msgs, text) myLogMacro(msgs, TR_LOG_DEBUG, text)
 #define logtrace(msgs, text) myLogMacro(msgs, TR_LOG_TRACE, text)
 
+tr_peerMsgs::~tr_peerMsgs()
+{
+    [[maybe_unused]] auto const n_prev = n_peers_--;
+    TR_ASSERT(n_prev > 0U);
+}
+
 /**
  * Low-level communication state information about a connected peer.
  *
diff --git a/libtransmission/peer-msgs.h b/libtransmission/peer-msgs.h
index 1eaf77680..806f6529a 100644
--- a/libtransmission/peer-msgs.h
+++ b/libtransmission/peer-msgs.h
@@ -9,6 +9,7 @@
 #error only libtransmission should #include this header.
 #endif
 
+#include <atomic>
 #include <cstdint> // int8_t
 #include <cstddef> // size_t
 #include <ctime> // time_t
@@ -35,6 +36,14 @@ public:
         : tr_peer{ tor, atom_in }
         , have_{ tor->pieceCount() }
     {
+        ++n_peers_;
+    }
+
+    virtual ~tr_peerMsgs() override;
+
+    [[nodiscard]] static size_t size() noexcept
+    {
+        return n_peers_.load();
     }
 
     [[nodiscard]] virtual bool is_peer_choked() const noexcept = 0;
@@ -67,6 +76,9 @@ public:
 
 protected:
     tr_bitfield have_;
+
+private:
+    static inline auto n_peers_ = std::atomic<size_t>{};
 };
 
 tr_peerMsgs* tr_peerMsgsNew(