mysten_network/
metrics.rs1use 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
13pub trait MetricsCallbackProvider: Send + Sync + Clone + 'static {
21 fn on_request(&self, path: String);
24
25 fn on_response(&self, path: String, latency: Duration, status: u16, grpc_status_code: Code);
31
32 fn on_start(&self, _path: &str) {}
34
35 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 }
102}