Metrics exporter to Prometheus with OTLP

ipa-core/Cargo.toml

@@ -88,6 +88,7 @@ ipa-metrics = { path = "../ipa-metrics" }
 ipa-metrics-tracing = { optional = true, path = "../ipa-metrics-tracing" }
 ipa-step = { version = "*", path = "../ipa-step" }
 ipa-step-derive = { version = "*", path = "../ipa-step-derive" }
+ipa-metrics-prometheus = { path = "../ipa-metrics-prometheus" }
 
 aes = "0.8.3"
 async-trait = "0.1.79"
@@ -176,6 +177,7 @@ rustls = { version = "0.23" }
 tempfile = "3"
 ipa-metrics-tracing = { path = "../ipa-metrics-tracing" }
 ipa-metrics = { path = "../ipa-metrics", features = ["partitions"] }
+ipa-metrics-prometheus = { path = "../ipa-metrics-prometheus" }
 
 
 [lib]

ipa-core/src/app.rs

@@ -3,6 +3,7 @@
 use async_trait::async_trait;
 
 use crate::{
+    cli::LoggingHandle,
     executor::IpaRuntime,
     helpers::{
         query::{CompareStatusRequest, PrepareQuery, QueryConfig, QueryInput},
@@ -65,6 +66,7 @@
     /// the flamegraph
     mpc_transport: MpcTransportImpl,
     shard_transport: ShardTransportImpl,
+    logging_handle: LoggingHandle,
 }
 
 impl Setup {
@@ -96,11 +98,13 @@
         self,
         mpc_transport: MpcTransportImpl,
         shard_transport: ShardTransportImpl,
+        logging_handle: LoggingHandle,
     ) -> HelperApp {
         let app = Arc::new(Inner {
             query_processor: self.query_processor,
             mpc_transport,
             shard_transport,
+            logging_handle,
         });
         self.mpc_handler
             .set_handler(Arc::downgrade(&app) as Weak<dyn RequestHandler<HelperIdentity>>);
@@ -277,6 +281,11 @@
                 let query_id = ext_query_id(&req)?;
                 HelperResponse::from(qp.kill(query_id)?)
             }
+            RouteId::Metrics => {
+                let logging_handler = &self.logging_handle;
+                let metrics_handle = &logging_handler.metrics_handle;
+                HelperResponse::from(metrics_handle.scrape_metrics())
+            }
         })
     }
 }

ipa-core/src/bin/helper.rs

@@ -271,7 +271,7 @@
         Some(shard_handler),
     );
 
-    let _app = setup.connect(transport.clone(), shard_transport.clone());
+    let _app = setup.connect(transport.clone(), shard_transport.clone(), logging_handle);
 
     let listener = create_listener(args.server_socket_fd)?;
     let shard_listener = create_listener(args.shard_server_socket_fd)?;

ipa-core/src/cli/metric_collector.rs

@@ -3,13 +3,14 @@
 use ipa_metrics::{
     MetricChannelType, MetricsCollectorController, MetricsCurrentThreadContext, MetricsProducer,
 };
+use ipa_metrics_prometheus::PrometheusMetricsExporter;
 use tokio::runtime::Builder;
 
 /// Holds a reference to metrics controller and producer
 pub struct CollectorHandle {
     thread_handle: JoinHandle<()>,
     /// This will be used once we start consuming metrics
-    _controller: MetricsCollectorController,
+    controller: MetricsCollectorController,
     producer: MetricsProducer,
 }
 
@@ -26,7 +27,7 @@
 
     Ok(CollectorHandle {
         thread_handle: handle,
-        _controller: controller,
+        controller,
         producer,
     })
 }
@@ -53,4 +54,17 @@
             .on_thread_stop(flush_fn)
             .on_thread_park(flush_fn)
     }
+
+    /// Export the metrics to be consumed by metrics scraper, e.g. Prometheus
+    ///
+    /// # Panics
+    /// If metrics is not initialized
+    #[must_use]
+    pub fn scrape_metrics(&self) -> Vec<u8> {
+        let mut store = self.controller.snapshot().expect("Metrics must be set up");
+        let mut buff = Vec::new();
+        store.export(&mut buff);
+
+        buff
+    }
 }

ipa-core/src/helpers/transport/handler.rs

@@ -149,6 +149,12 @@ impl<R: AsRef<dyn ProtocolResult>> From<R> for HelperResponse {
     }
 }
 
+impl From<Vec<u8>> for HelperResponse {
+    fn from(value: Vec<u8>) -> Self {
+        Self { body: value }
+    }
+}
+
 /// Union of error types returned by API operations.
 #[derive(thiserror::Error, Debug)]
 pub enum Error {

ipa-core/src/helpers/transport/in_memory/transport.rs

@@ -116,7 +116,8 @@ impl<I: TransportIdentity> InMemoryTransport<I> {
                             | RouteId::QueryInput
                             | RouteId::QueryStatus
                             | RouteId::CompleteQuery
-                            | RouteId::KillQuery => {
+                            | RouteId::KillQuery
+                            | RouteId::Metrics => {
                                 handler
                                     .as_ref()
                                     .expect("Handler is set")

ipa-core/src/helpers/transport/mod.rs

@@ -229,6 +229,26 @@ where
     fn extra(&self) -> Self::Params;
 }
 
+impl RouteParams<RouteId, NoQueryId, NoStep> for RouteId {
+    type Params = &'static str;
+
+    fn resource_identifier(&self) -> RouteId {
+        *self
+    }
+
+    fn query_id(&self) -> NoQueryId {
+        NoQueryId
+    }
+
+    fn gate(&self) -> NoStep {
+        NoStep
+    }
+
+    fn extra(&self) -> Self::Params {
+        ""
+    }
+}
+
 impl RouteParams<NoResourceIdentifier, QueryId, Gate> for (QueryId, Gate) {
     type Params = &'static str;
 

ipa-core/src/helpers/transport/routing.rs

@@ -24,6 +24,7 @@ pub enum RouteId {
     QueryStatus,
     CompleteQuery,
     KillQuery,
+    Metrics,
 }
 
 /// The header/metadata of the incoming request.

ipa-core/src/net/http_serde.rs

@@ -68,6 +68,16 @@ pub mod echo {
     pub const AXUM_PATH: &str = "/echo";
 }
 
+pub mod metrics {
+
+    use serde::{Deserialize, Serialize};
+
+    #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
+    pub struct Request {}
+
+    pub const AXUM_PATH: &str = "/metrics";
+}
+
 pub mod query {
     use std::fmt::{Display, Formatter};
 

ipa-core/src/net/server/handlers/metrics.rs

@@ -0,0 +1,62 @@
+use axum::{routing::get, Extension, Router};
+use hyper::StatusCode;
+
+use crate::{
+    helpers::{routing::RouteId, BodyStream},
+    net::{
+        http_serde::{self},
+        Error, MpcHttpTransport,
+    },
+};
+
+/// Takes details from the HTTP request and creates a `[TransportCommand]::CreateQuery` that is sent
+/// to the [`HttpTransport`].
+async fn handler(transport: Extension<MpcHttpTransport>) -> Result<Vec<u8>, Error> {
+    match transport
+        .dispatch(RouteId::Metrics, BodyStream::empty())
+        .await
+    {
+        Ok(resp) => Ok(resp.into_body()),
+        Err(err) => Err(Error::application(StatusCode::INTERNAL_SERVER_ERROR, err)),
+    }
+}
+
+pub fn router(transport: MpcHttpTransport) -> Router {
+    Router::new()
+        .route(http_serde::metrics::AXUM_PATH, get(handler))
+        .layer(Extension(transport))
+}
+
+#[cfg(all(test, unit_test))]
+mod tests {
+    use axum::{
+        body::Body,
+        http::uri::{self, Authority, Scheme},
+    };
+
+    use super::*;
+    use crate::{
+        helpers::{make_owned_handler, routing::Addr, HelperIdentity, HelperResponse},
+        net::server::handlers::query::test_helpers::assert_success_with,
+    };
+
+    #[tokio::test]
+    async fn happy_case() {
+        let handler = make_owned_handler(
+            move |addr: Addr<HelperIdentity>, _data: BodyStream| async move {
+                let RouteId::Metrics = addr.route else {
+                    panic!("unexpected call");
+                };
+                Ok(HelperResponse::from(Vec::new()))
+            },
+        );
+        let uri = uri::Builder::new()
+            .scheme(Scheme::HTTP)
+            .authority(Authority::from_static("localhost"))
+            .path_and_query(String::from("/metrics"))
+            .build()
+            .unwrap();
+        let req = hyper::Request::get(uri).body(Body::empty()).unwrap();
+        assert_success_with(req, handler).await;
+    }
+}

ipa-core/src/net/server/handlers/mod.rs

@@ -1,4 +1,5 @@
 mod echo;
+mod metrics;
 mod query;
 
 use axum::Router;
@@ -9,12 +10,14 @@ use crate::{
 };
 
 pub fn mpc_router(transport: MpcHttpTransport) -> Router {
-    echo::router().nest(
-        http_serde::query::BASE_AXUM_PATH,
-        Router::new()
-            .merge(query::query_router(transport.clone()))
-            .merge(query::h2h_router(transport.inner_transport)),
-    )
+    echo::router()
+        .merge(metrics::router(transport.clone()))
+        .nest(
+            http_serde::query::BASE_AXUM_PATH,
+            Router::new()
+                .merge(query::query_router(transport.clone()))
+                .merge(query::h2h_router(transport.inner_transport)),
+        )
 }
 
 pub fn shard_router(transport: Arc<HttpTransport<Shard>>) -> Router {

ipa-core/src/net/test.rs

@@ -23,6 +23,8 @@ use once_cell::sync::Lazy;
 use rustls_pki_types::CertificateDer;
 
 use super::{ConnectionFlavor, HttpTransport, Shard};
+#[cfg(all(test, web_test, descriptive_gate))]
+use crate::cli::{install_collector, LoggingHandle};
 use crate::{
     config::{
         ClientConfig, HpkeClientConfig, HpkeServerConfig, NetworkConfig, PeerConfig, ServerConfig,
@@ -263,7 +265,11 @@ impl TestApp {
             shard_server.start_on(&IpaRuntime::current(), self.shard_server.socket.take(), ()),
         )
         .await;
-        setup.connect(transport, shard_transport)
+
+        let metrics_handle = install_collector().unwrap();
+        let logging_handle = LoggingHandle { metrics_handle };
+
+        setup.connect(transport, shard_transport, logging_handle)
     }
 }
 

ipa-core/src/net/transport.rs

@@ -115,7 +115,10 @@
                 let req = serde_json::from_str(route.extra().borrow())?;
                 self.clients[client_ix].status_match(req).await
             }
-            evt @ (RouteId::QueryInput | RouteId::ReceiveQuery | RouteId::KillQuery) => {
+            evt @ (RouteId::QueryInput
+            | RouteId::ReceiveQuery
+            | RouteId::KillQuery
+            | RouteId::Metrics) => {
                 unimplemented!(
                     "attempting to send client-specific request {evt:?} to another helper"
                 )

ipa-core/src/test_fixture/app.rs

@@ -6,6 +6,7 @@ use typenum::Unsigned;
 
 use crate::{
     app::AppConfig,
+    cli::{install_collector, LoggingHandle},
     ff::Serializable,
     helpers::{
         query::{QueryConfig, QueryInput},
@@ -68,8 +69,15 @@ impl Default for TestApp {
 
         let mpc_network = InMemoryMpcNetwork::new(handlers.map(Some));
         let shard_network = InMemoryShardNetwork::with_shards(1);
-        let drivers = zip3(mpc_network.transports().each_ref(), setup)
-            .map(|(t, s)| s.connect(Clone::clone(t), shard_network.transport(t.identity(), 0)));
+        let drivers = zip3(mpc_network.transports().each_ref(), setup).map(|(t, s)| {
+            let metrics_handle = install_collector().unwrap();
+            let logging_handle = LoggingHandle { metrics_handle };
+            s.connect(
+                Clone::clone(t),
+                shard_network.transport(t.identity(), 0),
+                logging_handle,
+            )
+        });
 
         Self {
             drivers,

ipa-metrics-prometheus/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "ipa-metrics-prometheus"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+default = []
+
+[dependencies]
+ipa-metrics = { path = "../ipa-metrics" }
+
+# Open telemetry crates: opentelemetry-prometheus crate implementation is based on Opentelemetry API and SDK 0.23. (TBC)
+opentelemetry = "0.24"
+opentelemetry_sdk = { version = "0.24", features = ["metrics", "rt-tokio"] }
+opentelemetry-prometheus = { version = "0.17" }
+prometheus = "0.13.3"