udp: add prometheus support

Cargo.lock

@@ -231,6 +231,8 @@ dependencies = [
  "hex",
  "libc",
  "log",
+ "metrics",
+ "metrics-exporter-prometheus",
  "mimalloc",
  "mio",
  "num-format",

aquatic_udp/Cargo.toml

@@ -18,6 +18,7 @@ name = "aquatic_udp"
 
 [features]
 cpu-pinning = ["aquatic_common/hwloc"]
+prometheus = ["metrics", "metrics-exporter-prometheus"]
 
 [dependencies]
 aquatic_common.workspace = true
@@ -35,6 +36,8 @@ hdrhistogram = "7"
 hex = "0.4"
 libc = "0.2"
 log = "0.4"
+metrics = { version = "0.20", optional = true }
+metrics-exporter-prometheus = { version = "0.11", optional = true, default-features = false, features = ["http-listener"] }
 mimalloc = { version = "0.1", default-features = false }
 mio = { version = "0.8", features = ["net", "os-poll"] }
 num-format = "0.4"

aquatic_udp/src/config.rs

@@ -153,11 +153,16 @@ pub struct StatisticsConfig {
     pub write_html_to_file: bool,
     /// Path to save HTML file to
     pub html_file_path: PathBuf,
+    /// Run a prometheus endpoint
+    pub run_prometheus_endpoint: bool,
+    /// Address to run prometheus endpoint on
+    pub prometheus_endpoint_address: SocketAddr,
 }
 
 impl StatisticsConfig {
     pub fn active(&self) -> bool {
-        (self.interval != 0) & (self.print_to_stdout | self.write_html_to_file)
+        (self.interval != 0)
+            & (self.print_to_stdout | self.write_html_to_file | self.run_prometheus_endpoint)
     }
 }
 
@@ -169,6 +174,8 @@ impl Default for StatisticsConfig {
             print_to_stdout: false,
             write_html_to_file: false,
             html_file_path: "tmp/statistics.html".into(),
+            run_prometheus_endpoint: false,
+            prometheus_endpoint_address: SocketAddr::from(([0, 0, 0, 0], 9000)),
         }
     }
 }

aquatic_udp/src/lib.rs

@@ -141,6 +141,21 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
         let state = state.clone();
         let config = config.clone();
 
+        #[cfg(feature = "prometheus")]
+        if config.statistics.run_prometheus_endpoint {
+            use metrics_exporter_prometheus::PrometheusBuilder;
+
+            PrometheusBuilder::new()
+                .with_http_listener(config.statistics.prometheus_endpoint_address)
+                .install()
+                .with_context(|| {
+                    format!(
+                        "Install prometheus endpoint on {}",
+                        config.statistics.prometheus_endpoint_address
+                    )
+                })?;
+        }
+
         Builder::new()
             .name("statistics".into())
             .spawn(move || {

aquatic_udp/src/workers/statistics/collector.rs

@@ -14,15 +14,17 @@ pub struct StatisticsCollector {
     last_update: Instant,
     pending_histograms: Vec<Histogram<u64>>,
     last_complete_histogram: PeerHistogramStatistics,
+    ip_version: String,
 }
 
 impl StatisticsCollector {
-    pub fn new(shared: Arc<Statistics>) -> Self {
+    pub fn new(shared: Arc<Statistics>, ip_version: String) -> Self {
         Self {
             shared,
             last_update: Instant::now(),
             pending_histograms: Vec::new(),
             last_complete_histogram: Default::default(),
+            ip_version,
         }
     }
 
@@ -35,16 +37,28 @@ impl StatisticsCollector {
         }
     }
 
-    pub fn collect_from_shared(&mut self) -> CollectedStatistics {
+    pub fn collect_from_shared(&mut self, config: &Config) -> CollectedStatistics {
         let requests_received = Self::fetch_and_reset(&self.shared.requests_received);
         let responses_sent_connect = Self::fetch_and_reset(&self.shared.responses_sent_connect);
         let responses_sent_announce = Self::fetch_and_reset(&self.shared.responses_sent_announce);
         let responses_sent_scrape = Self::fetch_and_reset(&self.shared.responses_sent_scrape);
         let responses_sent_error = Self::fetch_and_reset(&self.shared.responses_sent_error);
+
         let bytes_received = Self::fetch_and_reset(&self.shared.bytes_received);
         let bytes_sent = Self::fetch_and_reset(&self.shared.bytes_sent);
-        let num_torrents = Self::sum_atomic_usizes(&self.shared.torrents);
+
-        let num_peers = Self::sum_atomic_usizes(&self.shared.peers);
+        let num_torrents_by_worker: Vec<usize> = self
+            .shared
+            .torrents
+            .iter()
+            .map(|n| n.load(Ordering::Relaxed))
+            .collect();
+        let num_peers_by_worker: Vec<usize> = self
+            .shared
+            .peers
+            .iter()
+            .map(|n| n.load(Ordering::Relaxed))
+            .collect();
 
         let elapsed = {
             let now = Instant::now();
@@ -56,13 +70,76 @@ impl StatisticsCollector {
             elapsed
         };
 
-        let requests_per_second = requests_received / elapsed;
+        #[cfg(feature = "prometheus")]
-        let responses_per_second_connect = responses_sent_connect / elapsed;
+        if config.statistics.run_prometheus_endpoint {
-        let responses_per_second_announce = responses_sent_announce / elapsed;
+            ::metrics::counter!(
-        let responses_per_second_scrape = responses_sent_scrape / elapsed;
+                "aquatic_requests_total",
-        let responses_per_second_error = responses_sent_error / elapsed;
+                requests_received.try_into().unwrap(),
-        let bytes_received_per_second = bytes_received / elapsed;
+                "ip_version" => self.ip_version.clone(),
-        let bytes_sent_per_second = bytes_sent / elapsed;
+            );
+            ::metrics::counter!(
+                "aquatic_responses_total",
+                responses_sent_connect.try_into().unwrap(),
+                "type" => "connect",
+                "ip_version" => self.ip_version.clone(),
+            );
+            ::metrics::counter!(
+                "aquatic_responses_total",
+                responses_sent_announce.try_into().unwrap(),
+                "type" => "announce",
+                "ip_version" => self.ip_version.clone(),
+            );
+            ::metrics::counter!(
+                "aquatic_responses_total",
+                responses_sent_scrape.try_into().unwrap(),
+                "type" => "scrape",
+                "ip_version" => self.ip_version.clone(),
+            );
+            ::metrics::counter!(
+                "aquatic_responses_total",
+                responses_sent_error.try_into().unwrap(),
+                "type" => "error",
+                "ip_version" => self.ip_version.clone(),
+            );
+            ::metrics::counter!(
+                "aquatic_rx_bytes",
+                bytes_received.try_into().unwrap(),
+                "ip_version" => self.ip_version.clone(),
+            );
+            ::metrics::counter!(
+                "aquatic_tx_bytes",
+                bytes_sent.try_into().unwrap(),
+                "ip_version" => self.ip_version.clone(),
+            );
+
+            for (worker_index, n) in num_torrents_by_worker.iter().copied().enumerate() {
+                ::metrics::gauge!(
+                    "aquatic_torrents",
+                    n as f64,
+                    "ip_version" => self.ip_version.clone(),
+                    "worker_index" => worker_index.to_string(),
+                );
+            }
+            for (worker_index, n) in num_peers_by_worker.iter().copied().enumerate() {
+                ::metrics::gauge!(
+                    "aquatic_peers",
+                    n as f64,
+                    "ip_version" => self.ip_version.clone(),
+                    "worker_index" => worker_index.to_string(),
+                );
+            }
+        }
+
+        let num_peers: usize = num_peers_by_worker.into_iter().sum();
+        let num_torrents: usize = num_torrents_by_worker.into_iter().sum();
+
+        let requests_per_second = requests_received as f64 / elapsed;
+        let responses_per_second_connect = responses_sent_connect as f64 / elapsed;
+        let responses_per_second_announce = responses_sent_announce as f64 / elapsed;
+        let responses_per_second_scrape = responses_sent_scrape as f64 / elapsed;
+        let responses_per_second_error = responses_sent_error as f64 / elapsed;
+        let bytes_received_per_second = bytes_received as f64 / elapsed;
+        let bytes_sent_per_second = bytes_sent as f64 / elapsed;
 
         let responses_per_second_total = responses_per_second_connect
             + responses_per_second_announce
@@ -89,12 +166,8 @@ impl StatisticsCollector {
         }
     }
 
-    fn sum_atomic_usizes(values: &[AtomicUsize]) -> usize {
+    fn fetch_and_reset(atomic: &AtomicUsize) -> usize {
-        values.iter().map(|n| n.load(Ordering::Relaxed)).sum()
+        atomic.fetch_and(0, Ordering::Relaxed)
-    }
-
-    fn fetch_and_reset(atomic: &AtomicUsize) -> f64 {
-        atomic.fetch_and(0, Ordering::Relaxed) as f64
     }
 }
 

aquatic_udp/src/workers/statistics/mod.rs

@@ -57,8 +57,8 @@ pub fn run_statistics_worker(
         None
     };
 
-    let mut ipv4_collector = StatisticsCollector::new(shared_state.statistics_ipv4);
+    let mut ipv4_collector = StatisticsCollector::new(shared_state.statistics_ipv4, "4".into());
-    let mut ipv6_collector = StatisticsCollector::new(shared_state.statistics_ipv6);
+    let mut ipv6_collector = StatisticsCollector::new(shared_state.statistics_ipv6, "6".into());
 
     loop {
         ::std::thread::sleep(Duration::from_secs(config.statistics.interval));
@@ -70,8 +70,8 @@ pub fn run_statistics_worker(
             }
         }
 
-        let statistics_ipv4 = ipv4_collector.collect_from_shared();
+        let statistics_ipv4 = ipv4_collector.collect_from_shared(&config);
-        let statistics_ipv6 = ipv6_collector.collect_from_shared();
+        let statistics_ipv6 = ipv6_collector.collect_from_shared(&config);
 
         if config.statistics.print_to_stdout {
             println!("General:");