From 08b28c9e1b8a8e35df706b664e288ea15d80a987 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Wed, 7 Jun 2023 12:58:41 +0200
Subject: [PATCH] udp statistics: improve peer client reporting

- fix prometheus peer id prefix reporting
- don't report peer id prefix in html output (current method
  was incorrect and output would become huge)
---
 .../src/workers/statistics/collector.rs       |  1 -
 aquatic_udp/src/workers/statistics/mod.rs     | 60 ++++++++++---------
 aquatic_udp/templates/statistics.html         |  2 -
 3 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/aquatic_udp/src/workers/statistics/collector.rs b/aquatic_udp/src/workers/statistics/collector.rs
index af25639..3e3ed97 100644
--- a/aquatic_udp/src/workers/statistics/collector.rs
+++ b/aquatic_udp/src/workers/statistics/collector.rs
@@ -2,7 +2,6 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Instant;
 
-use crossbeam_channel::Receiver;
 use hdrhistogram::Histogram;
 use num_format::{Locale, ToFormattedString};
 use serde::Serialize;
diff --git a/aquatic_udp/src/workers/statistics/mod.rs b/aquatic_udp/src/workers/statistics/mod.rs
index 8a1d5ad..aa453e1 100644
--- a/aquatic_udp/src/workers/statistics/mod.rs
+++ b/aquatic_udp/src/workers/statistics/mod.rs
@@ -37,7 +37,7 @@ struct TemplateData {
     ipv6: CollectedStatistics,
     last_updated: String,
     peer_update_interval: String,
-    peer_clients: Vec<(CompactString, CompactString, usize)>,
+    peer_clients: Vec<(CompactString, usize)>,
 }
 
 pub fn run_statistics_worker(
@@ -71,7 +71,7 @@ pub fn run_statistics_worker(
         "6".into(),
     );
 
-    let mut peer_clients: IndexMap<PeerClient, (usize, CompactString)> = IndexMap::default();
+    let mut peer_clients: IndexMap<PeerClient, usize> = IndexMap::default();
 
     loop {
         let start_time = Instant::now();
@@ -81,15 +81,35 @@ pub fn run_statistics_worker(
                 StatisticsMessage::Ipv4PeerHistogram(h) => ipv4_collector.add_histogram(&config, h),
                 StatisticsMessage::Ipv6PeerHistogram(h) => ipv6_collector.add_histogram(&config, h),
                 StatisticsMessage::PeerAdded(peer_id) => {
-                    peer_clients
-                        .entry(peer_id.client())
-                        .or_insert((0, peer_id.first_8_bytes_hex()))
-                        .0 += 1;
+                    let client = peer_id.client();
+                    let first_8_bytes_hex = peer_id.first_8_bytes_hex();
+
+                    #[cfg(feature = "prometheus")]
+                    if config.statistics.run_prometheus_endpoint {
+                        ::metrics::increment_gauge!(
+                            "aquatic_peer_clients",
+                            1.0,
+                            "client" => client.to_string(),
+                            "peer_id_prefix_hex" => first_8_bytes_hex.to_string(),
+                        );
+                    }
+
+                    *peer_clients.entry(client).or_insert(0) += 1;
                 }
                 StatisticsMessage::PeerRemoved(peer_id) => {
                     let client = peer_id.client();
 
-                    if let Some((count, _)) = peer_clients.get_mut(&client) {
+                    #[cfg(feature = "prometheus")]
+                    if config.statistics.run_prometheus_endpoint {
+                        ::metrics::decrement_gauge!(
+                            "aquatic_peer_clients",
+                            1.0,
+                            "client" => client.to_string(),
+                            "peer_id_prefix_hex" => peer_id.first_8_bytes_hex().to_string(),
+                        );
+                    }
+
+                    if let Some(count) = peer_clients.get_mut(&client) {
                         if *count == 1 {
                             drop(count);
 
@@ -102,18 +122,6 @@ pub fn run_statistics_worker(
             }
         }
 
-        #[cfg(feature = "prometheus")]
-        if config.statistics.run_prometheus_endpoint && config.statistics.extended {
-            for (peer_client, (count, first_8_bytes)) in peer_clients.iter() {
-                ::metrics::gauge!(
-                    "aquatic_peer_clients",
-                    *count as f64,
-                    "client" => peer_client.to_string(),
-                    "peer_id_prefix_hex" => first_8_bytes.to_string(),
-                );
-            }
-        }
-
         let statistics_ipv4 = ipv4_collector.collect_from_shared(
             #[cfg(feature = "prometheus")]
             &config,
@@ -146,19 +154,13 @@ pub fn run_statistics_worker(
             let mut peer_clients = if config.statistics.extended {
                 peer_clients
                     .iter()
-                    .map(|(peer_client, (count, first_8_bytes))| {
-                        (
-                            peer_client.to_compact_string(),
-                            first_8_bytes.to_owned(),
-                            *count,
-                        )
-                    })
+                    .map(|(peer_client, count)| (peer_client.to_compact_string(), *count))
                     .collect()
             } else {
                 Vec::new()
             };
 
-            peer_clients.sort_unstable_by(|a, b| b.2.cmp(&a.2));
+            peer_clients.sort_unstable_by(|a, b| b.1.cmp(&a.1));
 
             let template_data = TemplateData {
                 stylesheet: STYLESHEET_CONTENTS.to_string(),
@@ -183,6 +185,10 @@ pub fn run_statistics_worker(
             Duration::from_secs(config.statistics.interval).checked_sub(start_time.elapsed())
         {
             ::std::thread::sleep(time_remaining);
+        } else {
+            ::log::warn!(
+                "statistics interval not long enough to process all data, output may be misleading"
+            );
         }
     }
 }
diff --git a/aquatic_udp/templates/statistics.html b/aquatic_udp/templates/statistics.html
index a27ba81..77463f6 100644
--- a/aquatic_udp/templates/statistics.html
+++ b/aquatic_udp/templates/statistics.html
@@ -260,7 +260,6 @@
         <thead>
             <tr>
                 <th>Client</th>
-                <th>Peer ID prefix (hex)</th>
                 <th>Count</th>
             </tr>
         </thead>
@@ -268,7 +267,6 @@
             {{ for value in peer_clients }}
             <tr>
                 <td>{ value.0 }</td>
-                <td>{ value.1 }</td>
                 <td>{ value.2 }</td>
             </tr>
             {{ endfor }}