Skip to main content

salvo_otel/
metrics.rs

1use std::time::Instant;
2
3use opentelemetry::metrics::{Counter, Histogram};
4use opentelemetry::{KeyValue, global};
5use opentelemetry_semantic_conventions::trace;
6use salvo_core::http::ResBody;
7use salvo_core::prelude::*;
8
9/// Middleware for metrics with OpenTelemetry.
10#[derive(Debug)]
11pub struct Metrics {
12    request_count: Counter<u64>,
13    error_count: Counter<u64>,
14    duration: Histogram<f64>,
15}
16
17impl Default for Metrics {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl Metrics {
24    /// Create `Metrics` middleware with `meter`.
25    #[must_use]
26    pub fn new() -> Self {
27        let meter = global::meter("salvo");
28        Self {
29            request_count: meter
30                .u64_counter("salvo_request_count")
31                .with_description("total request count (since start of service)")
32                .build(),
33            error_count: meter
34                .u64_counter("salvo_error_count")
35                .with_description("failed request count (since start of service)")
36                .build(),
37            duration: meter
38                .f64_histogram("salvo_request_duration_ms")
39                .with_unit("milliseconds")
40                .with_description(
41                    "request duration histogram (in milliseconds, since start of service)",
42                )
43                .build(),
44        }
45    }
46}
47
48#[async_trait]
49impl Handler for Metrics {
50    async fn handle(
51        &self,
52        req: &mut Request,
53        depot: &mut Depot,
54        res: &mut Response,
55        ctrl: &mut FlowCtrl,
56    ) {
57        let mut labels = Vec::with_capacity(3);
58        labels.push(KeyValue::new(
59            trace::HTTP_REQUEST_METHOD,
60            req.method().to_string(),
61        ));
62        labels.push(KeyValue::new(trace::URL_FULL, req.uri().to_string()));
63
64        let s = Instant::now();
65        ctrl.call_next(req, depot, res).await;
66        let elapsed = s.elapsed();
67
68        let status = res.status_code.unwrap_or_else(|| {
69            tracing::info!("[otel::Metrics] Treat status_code=none as 200(OK).");
70            StatusCode::OK
71        });
72        labels.push(KeyValue::new(
73            trace::HTTP_RESPONSE_STATUS_CODE,
74            status.as_u16() as i64,
75        ));
76        if status.is_client_error() || status.is_server_error() {
77            self.error_count.add(1, &labels);
78            let msg = if let ResBody::Error(body) = &res.body {
79                body.to_string()
80            } else {
81                format!("ErrorCode: {}", status.as_u16())
82            };
83            labels.push(KeyValue::new(trace::EXCEPTION_MESSAGE, msg));
84        }
85
86        self.request_count.add(1, &labels);
87        self.duration
88            .record(elapsed.as_secs_f64() * 1000.0, &labels);
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use salvo_core::prelude::*;
95    use salvo_core::test::{ResponseExt, TestClient};
96
97    use super::*;
98    #[tokio::test]
99    async fn test_metrics_default() {
100        let metrics = Metrics::default();
101        assert_eq!(format!("{:?}", metrics).contains("Metrics"), true);
102    }
103
104    #[handler]
105    async fn hello() -> &'static str {
106        "Hello"
107    }
108
109    #[tokio::test]
110    async fn test_metrics_handle() {
111        let metrics = Metrics::default();
112
113        let router = Router::new().hoop(metrics).goal(hello);
114        let service = Service::new(router);
115
116        let content = TestClient::get("http://127.0.0.1:8698")
117            .send(&service)
118            .await
119            .take_string()
120            .await
121            .unwrap();
122        assert!(content == "Hello");
123    }
124}