mysten_network/
metrics.rs

1// Copyright (c) 2022, Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3use std::time::Duration;
4use tonic::codegen::http::header::HeaderName;
5use tonic::codegen::http::{HeaderValue, Request, Response};
6use tonic::{Code, Status};
7use tower_http::classify::GrpcFailureClass;
8use tower_http::trace::{OnFailure, OnRequest, OnResponse};
9use tracing::Span;
10
11pub(crate) static GRPC_ENDPOINT_PATH_HEADER: HeaderName = HeaderName::from_static("grpc-path-req");
12
13/// The trait to be implemented when want to be notified about
14/// a new request and related metrics around it. When a request
15/// is performed (up to the point that a response is created) the
16/// on_response method is called with the corresponding metrics
17/// details. The on_request method will be called when the request
18/// is received, but not further processing has happened at this
19/// point.
20pub trait MetricsCallbackProvider: Send + Sync + Clone + 'static {
21    /// Method will be called when a request has been received.
22    /// `path`: the endpoint uri path
23    fn on_request(&self, path: String);
24
25    /// Method to be called from the server when a request is performed.
26    /// `path`: the endpoint uri path
27    /// `latency`: the time when the request was received and when the response was created
28    /// `status`: the http status code of the response
29    /// `grpc_status_code`: the grpc status code (see https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc)
30    fn on_response(&self, path: String, latency: Duration, status: u16, grpc_status_code: Code);
31
32    /// Called when request call is started
33    fn on_start(&self, _path: &str) {}
34
35    /// Called when request call is dropped.
36    /// It is guaranteed that for each on_start there will be corresponding on_drop
37    fn on_drop(&self, _path: &str) {}
38}
39
40#[derive(Clone, Default)]
41pub struct DefaultMetricsCallbackProvider {}
42impl MetricsCallbackProvider for DefaultMetricsCallbackProvider {
43    fn on_request(&self, _path: String) {}
44
45    fn on_response(
46        &self,
47        _path: String,
48        _latency: Duration,
49        _status: u16,
50        _grpc_status_code: Code,
51    ) {
52    }
53}
54
55#[derive(Clone)]
56pub(crate) struct MetricsHandler<M: MetricsCallbackProvider> {
57    metrics_provider: M,
58}
59
60impl<M: MetricsCallbackProvider> MetricsHandler<M> {
61    pub(crate) fn new(metrics_provider: M) -> Self {
62        Self { metrics_provider }
63    }
64}
65
66impl<B, M: MetricsCallbackProvider> OnResponse<B> for MetricsHandler<M> {
67    fn on_response(self, response: &Response<B>, latency: Duration, _span: &Span) {
68        let grpc_status = Status::from_header_map(response.headers());
69        let grpc_status_code = grpc_status.map_or(Code::Ok, |s| s.code());
70
71        let path: HeaderValue = response
72            .headers()
73            .get(&GRPC_ENDPOINT_PATH_HEADER)
74            .unwrap()
75            .clone();
76
77        self.metrics_provider.on_response(
78            path.to_str().unwrap().to_string(),
79            latency,
80            response.status().as_u16(),
81            grpc_status_code,
82        );
83    }
84}
85
86impl<B, M: MetricsCallbackProvider> OnRequest<B> for MetricsHandler<M> {
87    fn on_request(&mut self, request: &Request<B>, _span: &Span) {
88        self.metrics_provider
89            .on_request(request.uri().path().to_string());
90    }
91}
92
93impl<M: MetricsCallbackProvider> OnFailure<GrpcFailureClass> for MetricsHandler<M> {
94    fn on_failure(
95        &mut self,
96        _failure_classification: GrpcFailureClass,
97        _latency: Duration,
98        _span: &Span,
99    ) {
100        // just do nothing for now so we avoid printing unnecessary logs
101    }
102}