Support Prometheus target metrics (#560)

* Add support for Prometheus metrics

* changelog

* clippies?

* Support explicit metrics selection

* clippy: allow long function

Author: GeorgeHahn

CHANGELOG.md

@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Added
 - Added target metrics support for Go expvars
+- Added target metrics support for Prometheus
 
 ### Fixed
 - Fixed throttle behavior for generators that run very quickly

src/target_metrics.rs

@@ -9,12 +9,15 @@ use serde::Deserialize;
 use crate::signals::Shutdown;
 
 pub mod expvar;
+pub mod prometheus;
 
 #[derive(Debug, Clone, Copy)]
 /// Errors produced by [`Server`]
 pub enum Error {
-    /// See [`crate::target_metrics::expvar::Expvar`] for details.
+    /// See [`crate::target_metrics::expvar::Error`] for details.
     Expvar(expvar::Error),
+    /// See [`crate::target_metrics::prometheus::Error`] for details.
+    Prometheus(prometheus::Error),
 }
 
 #[derive(Debug, Deserialize, PartialEq, Eq)]
@@ -23,13 +26,17 @@ pub enum Error {
 pub enum Config {
     /// See [`crate::target_metrics::expvar::Config`] for details.
     Expvar(expvar::Config),
+    /// See [`crate::target_metrics::prometheus::Config`] for details.
+    Prometheus(prometheus::Config),
 }
 
 /// The `target_metrics` server.
 #[derive(Debug)]
 pub enum Server {
     /// See [`crate::target_metrics::expvar::Expvar`] for details.
     Expvar(expvar::Expvar),
+    /// See [`crate::target_metrics::prometheus::Prometheus`] for details.
+    Prometheus(prometheus::Prometheus),
 }
 
 impl Server {
@@ -42,6 +49,9 @@ impl Server {
     pub fn new(config: Config, shutdown: Shutdown) -> Self {
         match config {
             Config::Expvar(conf) => Self::Expvar(expvar::Expvar::new(conf, shutdown)),
+            Config::Prometheus(conf) => {
+                Self::Prometheus(prometheus::Prometheus::new(conf, shutdown))
+            }
         }
     }
 
@@ -61,6 +71,7 @@ impl Server {
     pub async fn run(self) -> Result<(), Error> {
         match self {
             Server::Expvar(inner) => inner.run().await.map_err(Error::Expvar),
+            Server::Prometheus(inner) => inner.run().await.map_err(Error::Prometheus),
         }
     }
 }

src/target_metrics/expvar.rs

@@ -44,7 +44,6 @@ impl Expvar {
     /// using Go's expvar format.
     ///
     pub(crate) fn new(config: Config, shutdown: Shutdown) -> Self {
-        info!("expvar created");
         Self { config, shutdown }
     }
 
@@ -67,7 +66,7 @@ impl Expvar {
             loop {
                 tokio::time::sleep(Duration::from_secs(1)).await;
 
-                let Ok(resp) = client.get(&self.config.uri).send().await else {
+                let Ok(resp) = client.get(&self.config.uri).timeout(Duration::from_secs(1)).send().await else {
                     info!("failed to get expvar uri");
                     continue;
                 };

src/target_metrics/prometheus.rs

@@ -0,0 +1,215 @@
+//! Prometheus target metrics fetcher
+//!
+//! This module scrapes Prometheus/OpenMetrics formatted metrics from the target
+//! software.
+//!
+
+use std::{collections::HashMap, str::FromStr, time::Duration};
+
+use metrics::{counter, gauge};
+use serde::Deserialize;
+use tracing::{error, info, trace, warn};
+
+use crate::signals::Shutdown;
+
+#[derive(Debug, Clone, Copy, thiserror::Error)]
+/// Errors produced by [`Prometheus`]
+pub enum Error {
+    /// Prometheus scraper shut down unexpectedly
+    #[error("Unexpected shutdown")]
+    EarlyShutdown,
+}
+
+#[derive(Debug, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+/// Configuration for collecting Prometheus based target metrics
+pub struct Config {
+    /// URI to scrape
+    uri: String,
+    /// Metric names to scrape. Leave unset to scrape all metrics.
+    metrics: Option<Vec<String>>,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum MetricType {
+    Gauge,
+    Counter,
+    Histogram,
+    Summary,
+}
+
+#[derive(Debug)]
+enum MetricTypeParseError {
+    UnknownType,
+}
+
+impl FromStr for MetricType {
+    type Err = MetricTypeParseError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "counter" => Ok(Self::Counter),
+            "gauge" => Ok(Self::Gauge),
+            "histogram" => Ok(Self::Histogram),
+            "summary" => Ok(Self::Summary),
+            _ => Err(MetricTypeParseError::UnknownType),
+        }
+    }
+}
+
+/// The `Prometheus` target metrics implementation.
+#[derive(Debug)]
+pub struct Prometheus {
+    config: Config,
+    shutdown: Shutdown,
+}
+
+impl Prometheus {
+    /// Create a new [`Prometheus`] instance
+    ///
+    /// This is responsible for scraping metrics from the target process in the
+    /// Prometheus format.
+    ///
+    pub(crate) fn new(config: Config, shutdown: Shutdown) -> Self {
+        Self { config, shutdown }
+    }
+
+    /// Run this [`Server`] to completion
+    ///
+    /// Scrape metrics from the target at 1Hz.
+    ///
+    /// # Errors
+    ///
+    /// None are known.
+    ///
+    /// # Panics
+    ///
+    /// None are known.
+    #[allow(clippy::cast_sign_loss)]
+    #[allow(clippy::cast_possible_truncation)]
+    #[allow(clippy::too_many_lines)]
+    pub(crate) async fn run(mut self) -> Result<(), Error> {
+        info!("Prometheus target metrics scraper running");
+        let client = reqwest::Client::new();
+
+        let server = async move {
+            loop {
+                tokio::time::sleep(Duration::from_secs(1)).await;
+
+                let Ok(resp) = client.get(&self.config.uri).timeout(Duration::from_secs(1)).send().await else {
+                    info!("failed to get Prometheus uri");
+                    continue;
+                };
+
+                let Ok(text) = resp.text().await else {
+                    info!("failed to read Prometheus response");
+                    continue;
+                };
+
+                // remember the type for each metric across lines
+                let mut typemap = HashMap::new();
+
+                // this deserves a real parser, but this will do for now.
+                // Format doc: https://github.com/prometheus/docs/blob/main/content/docs/instrumenting/exposition_formats.md
+                for line in text.lines() {
+                    if line.starts_with("# HELP") {
+                        continue;
+                    }
+
+                    if line.starts_with("# TYPE") {
+                        let mut parts = line.split_ascii_whitespace().skip(2);
+                        let name = parts.next().unwrap();
+                        let metric_type = parts.next().unwrap();
+                        let metric_type: MetricType = metric_type.parse().unwrap();
+                        // summary and histogram metrics additionally report names suffixed with _sum, _count, _bucket
+                        if matches!(metric_type, MetricType::Histogram | MetricType::Summary) {
+                            typemap.insert(format!("{name}_sum"), metric_type);
+                            typemap.insert(format!("{name}_count"), metric_type);
+                            typemap.insert(format!("{name}_bucket"), metric_type);
+                        }
+                        typemap.insert(name.to_owned(), metric_type);
+                        continue;
+                    }
+
+                    let mut parts = line.split_ascii_whitespace();
+                    let name_and_labels = parts.next().unwrap();
+                    let value = parts.next().unwrap();
+
+                    let (name, labels) = {
+                        if let Some((name, labels)) = name_and_labels.split_once('{') {
+                            let labels = labels.trim_end_matches('}');
+                            let labels = labels.split(',').map(|label| {
+                                let (label_name, label_value) = label.split_once('=').unwrap();
+                                let label_value = label_value.trim_matches('\"');
+                                (label_name.to_owned(), label_value.to_owned())
+                            });
+                            let labels = labels.collect::<Vec<_>>();
+                            (name, Some(labels))
+                        } else {
+                            (name_and_labels, None)
+                        }
+                    };
+
+                    let metric_type = typemap.get(name);
+                    let name = name.replace("__", ".");
+
+                    if let Some(metrics) = &self.config.metrics {
+                        if !metrics.contains(&name) {
+                            continue;
+                        }
+                    }
+
+                    match metric_type {
+                        Some(MetricType::Gauge) => {
+                            let Ok(value): Result<f64, _> = value.parse() else {
+                                let e = value.parse::<f64>().unwrap_err();
+                                warn!("{e}: {name} = {value}");
+                                continue;
+                            };
+
+                            trace!("gauge: {name} = {value}");
+                            gauge!(format!("target/{name}"), value, &labels.unwrap_or_default());
+                        }
+                        Some(MetricType::Counter) => {
+                            let Ok(value): Result<f64, _> = value.parse() else {
+                                let e = value.parse::<f64>().unwrap_err();
+                                warn!("{e}: {name} = {value}");
+                                continue;
+                            };
+
+                            let value = if value < 0.0 {
+                                warn!("Negative counter value unhandled");
+                                continue;
+                            } else {
+                                // clippy shows "error: casting `f64` to `u64` may lose the sign of the value". This is
+                                // guarded by the sign check above.
+                                if value > u64::MAX as f64 {
+                                    warn!("Counter value above maximum limit");
+                                    continue;
+                                }
+                                value as u64
+                            };
+
+                            trace!("counter: {name} = {value}");
+                            counter!(format!("target/{name}"), value, &labels.unwrap_or_default());
+                        }
+                        Some(_) | None => {
+                            trace!("unsupported metric type: {name} = {value}");
+                        }
+                    }
+                }
+            }
+        };
+
+        tokio::select! {
+            _res = server => {
+                error!("server shutdown unexpectedly");
+                 Err(Error::EarlyShutdown)
+            }
+            _ = self.shutdown.recv() => {
+                info!("shutdown signal received");
+                 Ok(())
+            }
+        }
+    }
+}