iota_network_stack/
metrics.rs1use 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
19pub trait MetricsCallbackProvider: Send + Sync + Clone + 'static {
27 fn on_request(&self, path: String);
30
31 fn on_response(&self, path: String, latency: Duration, status: u16, grpc_status_code: Code);
37
38 fn on_error(&self, _latency: Duration, _grpc_status_code: Code) {}
45
46 fn on_start(&self, _path: &str) {}
48
49 fn on_drop(&self, _path: &str) {}
53}
54
55#[derive(Clone, Default)]
56pub struct DefaultMetricsCallbackProvider {}
57impl MetricsCallbackProvider for DefaultMetricsCallbackProvider {
58 fn on_request(&self, _path: String) {}
59
60 fn on_response(
61 &self,
62 _path: String,
63 _latency: Duration,
64 _status: u16,
65 _grpc_status_code: Code,
66 ) {
67 }
68}
69
70#[derive(Clone)]
71pub(crate) struct MetricsHandler<M: MetricsCallbackProvider> {
72 metrics_provider: M,
73}
74
75impl<M: MetricsCallbackProvider> MetricsHandler<M> {
76 pub(crate) fn new(metrics_provider: M) -> Self {
77 Self { metrics_provider }
78 }
79}
80
81impl<B, M: MetricsCallbackProvider> OnResponse<B> for MetricsHandler<M> {
82 fn on_response(self, response: &Response<B>, latency: Duration, _span: &Span) {
83 let grpc_status = Status::from_header_map(response.headers());
84 let grpc_status_code = grpc_status.map_or(Code::Ok, |s| s.code());
85
86 let path: HeaderValue = response
87 .headers()
88 .get(&GRPC_ENDPOINT_PATH_HEADER)
89 .unwrap()
90 .clone();
91
92 self.metrics_provider.on_response(
93 path.to_str().unwrap().to_string(),
94 latency,
95 response.status().as_u16(),
96 grpc_status_code,
97 );
98 }
99}
100
101impl<B, M: MetricsCallbackProvider> OnRequest<B> for MetricsHandler<M> {
102 fn on_request(&mut self, request: &Request<B>, _span: &Span) {
103 self.metrics_provider
104 .on_request(request.uri().path().to_string());
105 }
106}
107
108impl<M: MetricsCallbackProvider> OnFailure<GrpcFailureClass> for MetricsHandler<M> {
109 fn on_failure(
110 &mut self,
111 failure_classification: GrpcFailureClass,
112 latency: Duration,
113 _span: &Span,
114 ) {
115 if let GrpcFailureClass::Error(_) = failure_classification {
119 self.metrics_provider.on_error(latency, Code::Internal);
120 }
121 }
122}