udp: don't incorrectly remove peers in client counting

This commit is contained in:
Joakim Frostegård 2023-06-10 00:05:23 +02:00
parent e6b0585372
commit 8323d963c9

View file

@ -83,7 +83,9 @@ pub fn run_statistics_worker(
"6".into(), "6".into(),
); );
let mut peers: IndexMap<PeerId, (PeerClient, CompactString)> = IndexMap::default(); // Store a count to enable not removing peers from the count completely
// just because they were removed from one torrent
let mut peers: IndexMap<PeerId, (usize, PeerClient, CompactString)> = IndexMap::default();
loop { loop {
let start_time = Instant::now(); let start_time = Instant::now();
@ -94,15 +96,21 @@ pub fn run_statistics_worker(
StatisticsMessage::Ipv6PeerHistogram(h) => ipv6_collector.add_histogram(&config, h), StatisticsMessage::Ipv6PeerHistogram(h) => ipv6_collector.add_histogram(&config, h),
StatisticsMessage::PeerAdded(peer_id) => { StatisticsMessage::PeerAdded(peer_id) => {
if process_peer_client_data { if process_peer_client_data {
let peer_client = peer_id.client(); peers
let prefix = peer_id.first_8_bytes_hex(); .entry(peer_id)
.or_insert_with(|| (0, peer_id.client(), peer_id.first_8_bytes_hex()))
peers.insert(peer_id, (peer_client, prefix)); .0 += 1;
} }
} }
StatisticsMessage::PeerRemoved(peer_id) => { StatisticsMessage::PeerRemoved(peer_id) => {
if process_peer_client_data { if process_peer_client_data {
peers.remove(&peer_id); if let Some((count, _, _)) = peers.get_mut(&peer_id) {
*count -= 1;
if *count == 0 {
peers.remove(&peer_id);
}
}
} }
} }
} }
@ -123,7 +131,8 @@ pub fn run_statistics_worker(
#[cfg(feature = "prometheus")] #[cfg(feature = "prometheus")]
let mut prefixes: IndexMap<CompactString, usize> = IndexMap::default(); let mut prefixes: IndexMap<CompactString, usize> = IndexMap::default();
for (peer_client, prefix) in peers.values() { // Only count peer_ids once, even if they are in multiple torrents
for (_, peer_client, prefix) in peers.values() {
*clients.entry(peer_client.to_owned()).or_insert(0) += 1; *clients.entry(peer_client.to_owned()).or_insert(0) += 1;
#[cfg(feature = "prometheus")] #[cfg(feature = "prometheus")]