rpc/
metrics.rs

1use http::{Request, Response};
2use http_body::{Body, Frame};
3use pin_project_lite::pin_project;
4use rust_zero_core::{
5    CounterVec, GaugeVec, HistogramOptions, HistogramVec, Metrics, MetricsError, VectorOptions,
6};
7use std::{
8    collections::BTreeSet,
9    future::Future,
10    pin::Pin,
11    sync::Arc,
12    task::{Context, Poll},
13    time::Instant,
14};
15use tower::{Layer, Service};
16
17const UNKNOWN_METHOD: &str = "unknown";
18
19/// The side of a gRPC transport observed by [`RpcMetrics`].
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum RpcMetricMode {
22    Client,
23    Server,
24}
25
26impl RpcMetricMode {
27    fn name(self) -> &'static str {
28        match self {
29            Self::Client => "client",
30            Self::Server => "server",
31        }
32    }
33}
34
35/// Registered gRPC request, latency, and in-flight metrics.
36///
37/// Callers provide the generated RPC method paths up front (for example
38/// `/users.Users/Get`). Requests for any other path share the `unknown` label, preventing an
39/// attacker or a misconfigured client from creating an unbounded number of metric series.
40#[derive(Clone)]
41pub struct RpcMetrics {
42    requests: CounterVec,
43    duration: HistogramVec,
44    in_flight: GaugeVec,
45    methods: Arc<BTreeSet<String>>,
46}
47
48impl RpcMetrics {
49    pub fn new<I, S>(
50        metrics: &Metrics,
51        namespace: impl Into<String>,
52        mode: RpcMetricMode,
53        methods: I,
54    ) -> Result<Self, MetricsError>
55    where
56        I: IntoIterator<Item = S>,
57        S: Into<String>,
58    {
59        let namespace = namespace.into();
60        let subsystem = format!("rpc_{}", mode.name());
61        let options = |name, help| {
62            VectorOptions::new(name, help)
63                .with_namespace(namespace.clone())
64                .with_subsystem(subsystem.clone())
65        };
66
67        Ok(Self {
68            requests: metrics.counter_vec(
69                options("requests_total", "Completed gRPC requests")
70                    .with_labels(["method", "code"]),
71            )?,
72            duration: metrics.histogram_vec(
73                HistogramOptions::new("", "").with_vector_options(
74                    options("request_duration_seconds", "gRPC request duration")
75                        .with_labels(["method", "code"]),
76                ),
77            )?,
78            in_flight: metrics.gauge_vec(
79                options("requests_in_flight", "In-flight gRPC requests").with_labels(["method"]),
80            )?,
81            methods: Arc::new(methods.into_iter().map(Into::into).collect()),
82        })
83    }
84
85    fn method(&self, path: &str) -> String {
86        if self.methods.contains(path) {
87            path.to_owned()
88        } else {
89            UNKNOWN_METHOD.to_owned()
90        }
91    }
92
93    fn start(&self, path: &str) -> RpcObservation {
94        let method = self.method(path);
95        self.in_flight
96            .inc(&[&method])
97            .expect("gRPC in-flight metric labels are fixed");
98        RpcObservation {
99            metrics: self.clone(),
100            method,
101            started_at: Instant::now(),
102            finished: false,
103        }
104    }
105}
106
107struct RpcObservation {
108    metrics: RpcMetrics,
109    method: String,
110    started_at: Instant,
111    finished: bool,
112}
113
114impl RpcObservation {
115    fn finish(&mut self, code: &str) {
116        if self.finished {
117            return;
118        }
119        self.finished = true;
120        let labels = [self.method.as_str(), code];
121        self.metrics
122            .requests
123            .inc(&labels)
124            .expect("gRPC request metric labels are fixed");
125        self.metrics
126            .duration
127            .observe(self.started_at.elapsed().as_secs_f64(), &labels)
128            .expect("gRPC duration metric labels and observation are valid");
129        self.metrics
130            .in_flight
131            .dec(&[&self.method])
132            .expect("gRPC in-flight metric labels are fixed");
133    }
134}
135
136impl Drop for RpcObservation {
137    fn drop(&mut self) {
138        self.finish("cancelled");
139    }
140}
141
142/// Tower layer that records complete unary and streaming gRPC calls.
143#[derive(Clone)]
144pub struct RpcMetricsLayer {
145    metrics: RpcMetrics,
146}
147
148impl RpcMetricsLayer {
149    pub fn new(metrics: RpcMetrics) -> Self {
150        Self { metrics }
151    }
152}
153
154impl<S> Layer<S> for RpcMetricsLayer {
155    type Service = RpcMetricsService<S>;
156
157    fn layer(&self, inner: S) -> Self::Service {
158        RpcMetricsService {
159            inner,
160            metrics: self.metrics.clone(),
161        }
162    }
163}
164
165#[derive(Clone)]
166pub struct RpcMetricsService<S> {
167    inner: S,
168    metrics: RpcMetrics,
169}
170
171impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RpcMetricsService<S>
172where
173    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Send + 'static,
174    S::Future: Send + 'static,
175    S::Error: Send + 'static,
176    ReqBody: Send + 'static,
177    ResBody: Body + Send + 'static,
178{
179    type Response = Response<RpcMetricsBody<ResBody>>;
180    type Error = S::Error;
181    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
182
183    fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
184        self.inner.poll_ready(context)
185    }
186
187    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
188        let mut observation = self.metrics.start(request.uri().path());
189        let future = self.inner.call(request);
190        Box::pin(async move {
191            match future.await {
192                Ok(response) => {
193                    let header_code = grpc_code(response.headers())
194                        .map(str::to_owned)
195                        .or_else(|| (!response.status().is_success()).then(|| "http_error".into()));
196                    let (parts, body) = response.into_parts();
197                    let mut wrapped = RpcMetricsBody {
198                        inner: body,
199                        observation: Some(observation),
200                    };
201                    if let Some(code) = header_code {
202                        wrapped.finish(&code);
203                    }
204                    Ok(Response::from_parts(parts, wrapped))
205                }
206                Err(error) => {
207                    observation.finish("transport_error");
208                    Err(error)
209                }
210            }
211        })
212    }
213}
214
215pin_project! {
216    pub struct RpcMetricsBody<B> {
217        #[pin]
218        inner: B,
219        observation: Option<RpcObservation>,
220    }
221
222    impl<B> PinnedDrop for RpcMetricsBody<B> {
223        fn drop(this: Pin<&mut Self>) {
224            let this = this.project();
225            // Dropping an unfinished response body is a cancelled RPC.
226            this.observation.take();
227        }
228    }
229}
230
231impl<B> RpcMetricsBody<B> {
232    fn finish(&mut self, code: &str) {
233        if let Some(mut observation) = self.observation.take() {
234            observation.finish(code);
235        }
236    }
237}
238
239impl<B> Body for RpcMetricsBody<B>
240where
241    B: Body,
242{
243    type Data = B::Data;
244    type Error = B::Error;
245
246    fn poll_frame(
247        self: Pin<&mut Self>,
248        context: &mut Context<'_>,
249    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
250        let mut this = self.project();
251        match this.inner.as_mut().poll_frame(context) {
252            Poll::Ready(Some(Ok(frame))) => {
253                if let Some(code) = frame.trailers_ref().and_then(grpc_code) {
254                    if let Some(mut observation) = this.observation.take() {
255                        observation.finish(code);
256                    }
257                }
258                Poll::Ready(Some(Ok(frame)))
259            }
260            Poll::Ready(Some(Err(error))) => {
261                if let Some(mut observation) = this.observation.take() {
262                    observation.finish("body_error");
263                }
264                Poll::Ready(Some(Err(error)))
265            }
266            Poll::Ready(None) => {
267                if let Some(mut observation) = this.observation.take() {
268                    observation.finish("0");
269                }
270                Poll::Ready(None)
271            }
272            Poll::Pending => Poll::Pending,
273        }
274    }
275
276    fn is_end_stream(&self) -> bool {
277        self.inner.is_end_stream()
278    }
279
280    fn size_hint(&self) -> http_body::SizeHint {
281        self.inner.size_hint()
282    }
283}
284
285fn grpc_code(headers: &http::HeaderMap) -> Option<&str> {
286    headers
287        .get("grpc-status")
288        .and_then(|value| value.to_str().ok())
289        .map(|code| match code {
290            "0" | "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" | "10" | "11" | "12"
291            | "13" | "14" | "15" | "16" => code,
292            _ => "invalid",
293        })
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299    use http_body::Frame;
300    use std::{convert::Infallible, future::Ready};
301
302    #[derive(Clone)]
303    struct Reply {
304        code: &'static str,
305    }
306
307    impl Service<Request<()>> for Reply {
308        type Response = Response<OneFrame>;
309        type Error = Infallible;
310        type Future = Ready<Result<Self::Response, Self::Error>>;
311
312        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
313            Poll::Ready(Ok(()))
314        }
315
316        fn call(&mut self, _: Request<()>) -> Self::Future {
317            std::future::ready(Ok(Response::new(OneFrame(Some(self.code)))))
318        }
319    }
320
321    struct OneFrame(Option<&'static str>);
322
323    impl Body for OneFrame {
324        type Data = &'static [u8];
325        type Error = Infallible;
326
327        fn poll_frame(
328            mut self: Pin<&mut Self>,
329            _: &mut Context<'_>,
330        ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
331            Poll::Ready(self.0.take().map(|code| {
332                let mut trailers = http::HeaderMap::new();
333                trailers.insert("grpc-status", code.parse().unwrap());
334                Ok(Frame::trailers(trailers))
335            }))
336        }
337    }
338
339    #[tokio::test]
340    async fn records_trailer_status_and_bounds_unknown_methods() {
341        use tower::ServiceExt;
342
343        let registry = Metrics::new();
344        let metrics = RpcMetrics::new(
345            &registry,
346            "users",
347            RpcMetricMode::Server,
348            ["/users.Users/Get"],
349        )
350        .unwrap();
351        let mut service = RpcMetricsLayer::new(metrics).layer(Reply { code: "5" });
352        let response = service
353            .ready()
354            .await
355            .unwrap()
356            .call(Request::builder().uri("/attacker/value").body(()).unwrap())
357            .await
358            .unwrap();
359        let mut body = Box::pin(response.into_body());
360        std::future::poll_fn(|context| body.as_mut().poll_frame(context)).await;
361
362        let rendered = registry.render();
363        assert!(
364            rendered.contains("users_rpc_server_requests_total{method=\"unknown\",code=\"5\"} 1")
365        );
366        assert!(rendered.contains("users_rpc_server_requests_in_flight{method=\"unknown\"} 0"));
367        assert!(!rendered.contains("attacker"));
368    }
369}
rpc/metrics.rs

rpc/
metrics.rs
