From 0204b6fcc2dea04d85f667fb0e321d9a0e0c7c52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Wed, 14 Jun 2023 11:36:08 +0200
Subject: [PATCH] ws: support serving metrics on peer id prefixes

---
 aquatic_ws/src/config.rs         |  9 ++++++-
 aquatic_ws/src/workers/socket.rs | 42 ++++++++++++++++++++++++++------
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/aquatic_ws/src/config.rs b/aquatic_ws/src/config.rs
index 36bb828..25651e1 100644
--- a/aquatic_ws/src/config.rs
+++ b/aquatic_ws/src/config.rs
@@ -156,10 +156,16 @@ pub struct MetricsConfig {
     pub prometheus_endpoint_address: SocketAddr,
     /// Update metrics for torrent count this often (seconds)
     pub torrent_count_update_interval: u64,
-    /// Collect data on peer clients.
+    /// Serve information on peer clients
     ///
     /// Expect a certain CPU hit
     pub peer_clients: bool,
+    /// Serve information on all peer id prefixes
+    ///
+    /// Requires `peer_clients` to be activated.
+    ///
+    /// Expect a certain CPU hit
+    pub peer_id_prefixes: bool,
 }
 
 #[cfg(feature = "metrics")]
@@ -170,6 +176,7 @@ impl Default for MetricsConfig {
             prometheus_endpoint_address: SocketAddr::from(([0, 0, 0, 0], 9000)),
             torrent_count_update_interval: 10,
             peer_clients: false,
+            peer_id_prefixes: false,
         }
     }
 }
diff --git a/aquatic_ws/src/workers/socket.rs b/aquatic_ws/src/workers/socket.rs
index 8559612..09c1884 100644
--- a/aquatic_ws/src/workers/socket.rs
+++ b/aquatic_ws/src/workers/socket.rs
@@ -52,7 +52,7 @@ struct ConnectionReference {
     valid_until: ValidUntil,
     announced_info_hashes: HashMap<InfoHash, PeerId>,
     ip_version: IpVersion,
-    opt_peer_client: Option<PeerClient>,
+    opt_peer_client: Option<(PeerClient, String)>,
 }
 
 pub async fn run_socket_worker(
@@ -226,12 +226,20 @@ pub async fn run_socket_worker(
                         }
 
                         #[cfg(feature = "prometheus")]
-                        if let Some(peer_client) = reference.opt_peer_client {
+                        if let Some((peer_client, prefix)) = reference.opt_peer_client {
                             ::metrics::decrement_gauge!(
                                 "aquatic_peer_clients",
                                 1.0,
                                 "client" => peer_client.to_string(),
                             );
+
+                            if config.metrics.peer_id_prefixes {
+                                ::metrics::decrement_gauge!(
+                                    "aquatic_peer_id_prefixes",
+                                    1.0,
+                                    "prefix_hex" => prefix.to_string(),
+                                );
+                            }
                         }
                     }
                 }), tq_regular)
@@ -259,15 +267,24 @@ async fn clean_connections(
     connection_slab.borrow_mut().retain(|_, reference| {
         if reference.valid_until.valid(now) {
             #[cfg(feature = "prometheus")]
-            if let Some(peer_client) = &reference.opt_peer_client {
+            if let Some((peer_client, prefix)) = &reference.opt_peer_client {
                 // As long as connection is still alive, increment peer client
-                // gauge by zero to prevent it from being removed due to
+                // gauges by zero to prevent them from being removed due to
                 // idleness
+
                 ::metrics::increment_gauge!(
                     "aquatic_peer_clients",
                     0.0,
                     "client" => peer_client.to_string(),
                 );
+
+                if config.metrics.peer_id_prefixes {
+                    ::metrics::increment_gauge!(
+                        "aquatic_peer_id_prefixes",
+                        0.0,
+                        "prefix_hex" => prefix.to_string(),
+                    );
+                }
             }
 
             true
@@ -629,9 +646,10 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
                                     && self.config.metrics.peer_clients
                                     && connection_reference.opt_peer_client.is_none()
                                 {
-                                    let client =
-                                        aquatic_peer_id::PeerId(announce_request.peer_id.0)
-                                            .client();
+                                    let peer_id =
+                                        aquatic_peer_id::PeerId(announce_request.peer_id.0);
+                                    let client = peer_id.client();
+                                    let prefix = peer_id.first_8_bytes_hex().to_string();
 
                                     ::metrics::increment_gauge!(
                                         "aquatic_peer_clients",
@@ -639,7 +657,15 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
                                         "client" => client.to_string(),
                                     );
 
-                                    connection_reference.opt_peer_client = Some(client);
+                                    if self.config.metrics.peer_id_prefixes {
+                                        ::metrics::increment_gauge!(
+                                            "aquatic_peer_id_prefixes",
+                                            1.0,
+                                            "prefix_hex" => prefix.to_string(),
+                                        );
+                                    }
+
+                                    connection_reference.opt_peer_client = Some((client, prefix));
                                 };
 
                                 entry.insert(announce_request.peer_id);