Add registry method for metrics and tests

Layr-Labs · Jun 10, 2024 · 8ed044c · 8ed044c
1 parent 50818c1
commit 8ed044c
Show file tree

Hide file tree

Showing 7 changed files with 196 additions and 8 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -42,6 +42,7 @@ rustdoc.all = "warn"
 
 
 [workspace.dependencies]
+metrics-exporter-prometheus = "0.12.0"
 ethers =  "2.0.14"
 ark-ff = "0.4.0"
 eyre = "0.6.12"
@@ -75,6 +76,8 @@ tokio = {version = "1.37.0" , features = ["test-util", "full","sync"] }
 futures-util = "0.3.30"
 thiserror = "1.0"
 tracing = "0.1.40"
+hyper = "0.14.25"
+
 
 #misc 
 parking_lot = "0.12"

diff --git a/crates/metrics/Cargo.toml b/crates/metrics/Cargo.toml
@@ -14,4 +14,13 @@ eigen-metrics-derive.workspace = true
 metrics.workspace = true
 
 #prometheus
-prometheus = "0.13.4"
+prometheus = "0.13.4"
+metrics-exporter-prometheus.workspace = true
+
+reqwest = "0.12.4"
+eyre.workspace = true
+
+#http
+hyper.workspace = true
+
+tokio.workspace = true
diff --git a/crates/metrics/collectors/economic/Cargo.toml b/crates/metrics/collectors/economic/Cargo.toml
@@ -10,4 +10,4 @@ repository.workspace = true
 ethers.workspace = true
 eigen-client-elcontracts.workspace = true
 eigen-client-avsregistry.workspace = true
-alloy-primitives.workspace = true
+alloy-primitives.workspace = true
diff --git a/crates/metrics/src/eigenmetrics.rs b/crates/metrics/src/eigenmetrics.rs
@@ -21,4 +21,8 @@ impl EigenMetrics {
         gauge.performance_score.set(100 as f64);
         gauge
     }
+
+    pub fn performance_score(&self) -> Gauge {
+        return self.performance_score.clone();
+    }
 }
diff --git a/crates/metrics/src/lib.rs b/crates/metrics/src/lib.rs
@@ -1,2 +1,3 @@
 pub mod eigenmetrics;
 pub mod noopmetrics;
+pub mod prometheus;
diff --git a/crates/metrics/src/prometheus.rs b/crates/metrics/src/prometheus.rs
@@ -0,0 +1,77 @@
+use crate::eigenmetrics::EigenMetrics;
+use eyre::Result;
+use hyper::{
+    body::Body,
+    service::{make_service_fn, service_fn},
+    Request, Response, Server,
+};
+use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
+use std::{convert::Infallible, net::SocketAddr};
+use tokio::runtime::Runtime;
+use tokio::time::sleep;
+use tokio::time::Duration;
+
+fn init_registry() -> PrometheusHandle {
+    let recorder = PrometheusBuilder::new().build_recorder();
+    let handle = recorder.handle();
+    let boxed_recorder = Box::new(recorder);
+    let static_recorder: &'static dyn metrics::Recorder = Box::leak(boxed_recorder);
+    metrics::set_recorder(static_recorder).expect("failed to set metrics recorder");
+    handle
+}
+
+async fn serve_metrics(addr: SocketAddr, handle: PrometheusHandle) -> eyre::Result<()> {
+    let make_svc = make_service_fn(move |_| {
+        let handle = handle.clone();
+
+        async move {
+            Ok::<_, Infallible>(service_fn(move |_: Request<Body>| {
+                let metrics = handle.render();
+                async move { Ok::<_, Infallible>(Response::new(Body::from(metrics))) }
+            }))
+        }
+    });
+
+    let server = Server::bind(&addr).serve(make_svc);
+    server.await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_prometheus_server() {
+    let socket: SocketAddr = "127.0.0.1:9091".parse().unwrap();
+    let handle = init_registry();
+
+    // Initialize EigenMetrics
+    let metrics = EigenMetrics::new();
+
+    // Run the metrics server in a background task
+    let server_handle = tokio::spawn(async move {
+        serve_metrics(socket, handle).await.unwrap();
+    });
+
+    sleep(Duration::from_secs(1)).await;
+
+    async fn get_metrics_body(client: &reqwest::Client, url: &str) -> String {
+        let resp = client.get(url).send().await.unwrap();
+        resp.text().await.unwrap()
+    }
+    let client = reqwest::Client::new();
+    let resp = client
+        .get("http://127.0.0.1:9091/metrics")
+        .send()
+        .await
+        .unwrap();
+
+    let mut body = resp.text().await.unwrap();
+    assert!(body.contains("eigen_performance_score 100"));
+
+    metrics.performance_score().set(80.0);
+
+    sleep(Duration::from_secs(1)).await;
+
+    body = get_metrics_body(&client, "http://127.0.0.1:9091/metrics").await;
+    assert!(body.contains("eigen_performance_score 80"));
+    // Shutdown the server
+    server_handle.abort();
+}