iota_network_stack/
metrics.rs

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

iota_network_stack/
metrics.rs
