saluki_core/observability/metrics/
api.rs

1//! Internal telemetry API handler.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use saluki_api::{
7    extract::State,
8    response::IntoResponse,
9    routing::{get, Router},
10    APIHandler, DynamicRoute, EndpointType,
11};
12use saluki_common::sync::shutdown::ShutdownHandle;
13use saluki_error::generic_error;
14use tokio::sync::Mutex;
15
16use super::{get_shared_metrics_state, AggregatedMetricsProcessor, Reflector, RemapperRule, TelemetryProcessor};
17use crate::runtime::{state::DataspaceRegistry, InitializationError, Supervisable, SupervisorFuture};
18
19/// State shared by the telemetry API routes.
20#[derive(Clone)]
21pub struct InternalTelemetryState {
22    metrics: Reflector<AggregatedMetricsProcessor>,
23    raw: Arc<Mutex<TelemetryProcessor>>,
24    compat: Arc<Mutex<TelemetryProcessor>>,
25}
26
27/// API handler that exposes the internal telemetry routes (`/metrics`, and optionally `/compat/metrics`) on the
28/// unprivileged endpoint.
29pub struct InternalTelemetryAPIHandler {
30    state: InternalTelemetryState,
31    compat_enabled: bool,
32}
33
34impl InternalTelemetryAPIHandler {
35    fn new(metrics: Reflector<AggregatedMetricsProcessor>, compat_remapper_rules: Option<Vec<RemapperRule>>) -> Self {
36        let compat_enabled = compat_remapper_rules.is_some();
37        let compat_remapper_rules = compat_remapper_rules.unwrap_or_default();
38
39        Self {
40            state: InternalTelemetryState {
41                metrics,
42                raw: Arc::new(Mutex::new(TelemetryProcessor::new())),
43                compat: Arc::new(Mutex::new(
44                    TelemetryProcessor::new().with_remapper_rules(compat_remapper_rules),
45                )),
46            },
47            compat_enabled,
48        }
49    }
50
51    async fn metrics_handler(State(state): State<InternalTelemetryState>) -> impl IntoResponse {
52        let mut processor = state.raw.lock().await;
53        processor.process(state.metrics.state())
54    }
55
56    async fn compat_metrics_handler(State(state): State<InternalTelemetryState>) -> impl IntoResponse {
57        let mut processor = state.compat.lock().await;
58        processor.process(state.metrics.state())
59    }
60}
61
62impl APIHandler for InternalTelemetryAPIHandler {
63    type State = InternalTelemetryState;
64
65    fn generate_initial_state(&self) -> Self::State {
66        self.state.clone()
67    }
68
69    fn generate_routes(&self) -> Router<Self::State> {
70        let router = Router::new().route("/metrics", get(Self::metrics_handler));
71        if self.compat_enabled {
72            router.route("/compat/metrics", get(Self::compat_metrics_handler))
73        } else {
74            router
75        }
76    }
77}
78
79/// A worker that exposes the process's internal telemetry on the unprivileged API endpoint.
80///
81/// Asserts these routes on the unprivileged endpoint:
82///
83/// - `GET /metrics`: the full aggregated metrics state, rendered as Prometheus text exposition.
84/// - `GET /compat/metrics`, when compatibility remapper rules are configured: the subset of metrics that match the
85///   rules, rendered with their remapped names.
86pub struct InternalTelemetryAPIWorker {
87    compat_remapper_rules: Option<Vec<RemapperRule>>,
88}
89
90impl InternalTelemetryAPIWorker {
91    /// Creates a new `InternalTelemetryAPIWorker`.
92    ///
93    /// When `compat_remapper_rules` is `Some`, the worker also serves `/compat/metrics` with those rules.
94    pub fn new(compat_remapper_rules: Option<Vec<RemapperRule>>) -> Self {
95        Self { compat_remapper_rules }
96    }
97}
98
99#[async_trait]
100impl Supervisable for InternalTelemetryAPIWorker {
101    fn name(&self) -> &str {
102        "internal-telemetry-api"
103    }
104
105    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
106        let metrics = get_shared_metrics_state();
107        let handler = InternalTelemetryAPIHandler::new(metrics, self.compat_remapper_rules.clone());
108        let route = DynamicRoute::http(EndpointType::Unprivileged, &handler);
109
110        Ok(Box::pin(async move {
111            DataspaceRegistry::try_current()
112                .ok_or_else(|| generic_error!("Dataspace not available."))?
113                .assert(route, "internal-telemetry-api");
114
115            process_shutdown.await;
116            Ok(())
117        }))
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use futures::stream;
124    use http::{Request, StatusCode};
125    use http_body_util::Empty;
126    use hyper::body::Bytes;
127    use tower::ServiceExt as _;
128
129    use super::*;
130    use crate::observability::metrics::MetricsSnapshot;
131
132    fn handler(compat_remapper_rules: Option<Vec<RemapperRule>>) -> InternalTelemetryAPIHandler {
133        let (metrics, _worker) = Reflector::new(stream::empty::<Vec<MetricsSnapshot>>(), AggregatedMetricsProcessor);
134        InternalTelemetryAPIHandler::new(metrics, compat_remapper_rules)
135    }
136
137    async fn status(handler: &InternalTelemetryAPIHandler, path: &str) -> StatusCode {
138        let request = Request::get(path).body(Empty::<Bytes>::new()).unwrap();
139        handler
140            .generate_routes()
141            .with_state(handler.generate_initial_state())
142            .oneshot(request)
143            .await
144            .unwrap()
145            .status()
146    }
147
148    #[tokio::test]
149    async fn compat_route_is_only_served_with_compat_rules() {
150        let without_compat = handler(None);
151        assert_eq!(status(&without_compat, "/metrics").await, StatusCode::OK);
152        assert_eq!(status(&without_compat, "/compat/metrics").await, StatusCode::NOT_FOUND);
153
154        let with_compat = handler(Some(Vec::new()));
155        assert_eq!(status(&with_compat, "/metrics").await, StatusCode::OK);
156        assert_eq!(status(&with_compat, "/compat/metrics").await, StatusCode::OK);
157    }
158}
saluki_core/observability/metrics/api.rs

saluki_core/observability/metrics/
api.rs
