Prometheus Endpoint (#706)

worldcoin · Mar 5, 2024 · b5bd04e · b5bd04e
1 parent e6fba00
commit b5bd04e
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 54 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -5,7 +5,6 @@ use std::time::Duration;
 use ethers::types::{Address, H160};
 use semaphore::Field;
 use serde::{Deserialize, Serialize};
-use telemetry_batteries::metrics::prometheus::PrometheusExporterConfig;
 
 use crate::prover::ProverConfig;
 use crate::utils::secret::SecretUrl;
@@ -213,50 +212,20 @@ pub struct ServiceConfig {
     #[serde(default = "default::service_name")]
     pub service_name: String,
     pub datadog:      Option<DatadogConfig>,
-    #[serde(default = "default::metrics")]
-    pub metrics:      Option<MetricsConfig>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(rename_all = "snake_case")]
-pub enum MetricsConfig {
-    Prometheus(PrometheusExporterConfig),
-    Statsd(StatsdConfig),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DatadogConfig {
     pub traces_endpoint: Option<String>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct StatsdConfig {
-    pub metrics_host:        String,
-    pub metrics_port:        u16,
-    pub metrics_queue_size:  usize,
-    pub metrics_buffer_size: usize,
-    pub metrics_prefix:      String,
-}
-
 pub mod default {
     use std::time::Duration;
 
-    use telemetry_batteries::metrics::prometheus::PrometheusExporterConfig;
-
-    use super::MetricsConfig;
-
     pub fn service_name() -> String {
         "signup_sequencer".to_string()
     }
 
-    pub fn metrics() -> Option<MetricsConfig> {
-        Some(MetricsConfig::Prometheus(
-            PrometheusExporterConfig::HttpListener {
-                listen_address: "0.0.0.0:9998".parse().unwrap(),
-            },
-        ))
-    }
-
     pub fn oz_api_url() -> String {
         "https://api.defender.openzeppelin.com".to_string()
     }
@@ -425,15 +394,13 @@ mod tests {
 
         [service.datadog]
         traces_endpoint = "http://localhost:8126"
-
-        [service.metrics.prometheus.http_listener]
-        listen_address = "0.0.0.0:9998"
     "#};
 
     #[test]
     fn full_toml_round_trip() {
         let config: Config = toml::from_str(FULL_TOML).unwrap();
         let serialized = toml::to_string_pretty(&config).unwrap();
+        println!("{}", serialized);
         similar_asserts::assert_eq!(serialized.trim(), FULL_TOML.trim());
     }
 }
diff --git a/src/main.rs b/src/main.rs
@@ -10,12 +10,10 @@ use std::path::PathBuf;
 
 use clap::Parser;
 use signup_sequencer::app::App;
-use signup_sequencer::config::{Config, MetricsConfig, ServiceConfig};
+use signup_sequencer::config::{Config, ServiceConfig};
 use signup_sequencer::server;
 use signup_sequencer::shutdown::watch_shutdown_signals;
 use signup_sequencer::task_monitor::TaskMonitor;
-use telemetry_batteries::metrics::prometheus::PrometheusBattery;
-use telemetry_batteries::metrics::statsd::StatsdBattery;
 use telemetry_batteries::tracing::datadog::DatadogBattery;
 use telemetry_batteries::tracing::stdout::StdoutBattery;
 use telemetry_batteries::tracing::TracingShutdownHandle;
@@ -79,22 +77,6 @@ fn load_config(args: &Args) -> anyhow::Result<Config> {
 }
 
 fn init_telemetry(service: &ServiceConfig) -> anyhow::Result<TracingShutdownHandle> {
-    match service.metrics.clone() {
-        Some(MetricsConfig::Prometheus(prometheus)) => {
-            PrometheusBattery::init(Some(prometheus))?;
-        }
-        Some(MetricsConfig::Statsd(statsd)) => {
-            StatsdBattery::init(
-                &statsd.metrics_host,
-                statsd.metrics_port,
-                statsd.metrics_queue_size,
-                statsd.metrics_buffer_size,
-                Some(&statsd.metrics_prefix),
-            )?;
-        }
-        _ => {}
-    }
-
     if let Some(ref datadog) = service.datadog {
         Ok(DatadogBattery::init(
             datadog.traces_endpoint.as_deref(),

diff --git a/src/server.rs b/src/server.rs
@@ -5,10 +5,13 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use axum::extract::{Query, State};
+use axum::response::Response;
 use axum::routing::{get, post};
 use axum::{middleware, Json, Router};
 use error::Error;
-use hyper::StatusCode;
+use hyper::header::CONTENT_TYPE;
+use hyper::{Body, StatusCode};
+use prometheus::{Encoder, TextEncoder};
 use tracing::info;
 
 use crate::app::App;
@@ -132,6 +135,23 @@ async fn health() -> Result<(), Error> {
     Ok(())
 }
 
+async fn metrics() -> Result<Response<Body>, Error> {
+    let encoder = TextEncoder::new();
+
+    let metric_families = prometheus::gather();
+    let mut buffer = vec![];
+    encoder
+        .encode(&metric_families, &mut buffer)
+        .map_err(|e| Error::Other(e.into()))?;
+
+    let response = Response::builder()
+        .status(200)
+        .header(CONTENT_TYPE, encoder.format_type())
+        .body(Body::from(buffer))?;
+
+    Ok(response)
+}
+
 /// # Errors
 ///
 /// Will return `Err` if `options.server` URI is not http, incorrectly includes
@@ -169,6 +189,7 @@ pub async fn bind_from_listener(
         .route("/listBatchSizes", get(list_batch_sizes))
         // Health check, return 200 OK
         .route("/health", get(health))
+        .route("/metrics", get(metrics))
         .layer(middleware::from_fn(
             custom_middleware::api_metrics_layer::middleware,
         ))

diff --git a/tests/common/mod.rs b/tests/common/mod.rs
@@ -618,11 +618,34 @@ pub async fn spawn_app(config: Config) -> anyhow::Result<(JoinHandle<()>, Socket
 
     info!("Checking app health");
     check_health(&local_addr).await?;
+
+    info!("Checking metrics");
+    check_metrics(&local_addr).await?;
+
     info!("App ready");
 
     Ok((app, local_addr))
 }
 
+pub async fn check_metrics(socket_addr: &SocketAddr) -> anyhow::Result<()> {
+    let uri = format!("http://{}", socket_addr);
+    let client = Client::new();
+    let req = Request::builder()
+        .method("GET")
+        .uri(uri.to_owned() + "/metrics")
+        .body(Body::empty())
+        .expect("Failed to create metrics hyper::Body");
+    let response = client
+        .request(req)
+        .await
+        .context("Failed to execute metrics request.")?;
+    if !response.status().is_success() {
+        anyhow::bail!("Metrics endpoint failed");
+    }
+
+    Ok(())
+}
+
 pub async fn check_health(socket_addr: &SocketAddr) -> anyhow::Result<()> {
     let uri = format!("http://{}", socket_addr);
     let client = Client::new();