fedimint_server/metrics/
jsonrpsee.rs1use std::borrow::Cow;
8use std::collections::HashSet;
9use std::pin::Pin;
10use std::sync::Arc;
11use std::task;
12use std::task::Poll;
13
14use fedimint_metrics::prometheus::HistogramTimer;
15use futures::Future;
16use jsonrpsee::MethodResponse;
17use jsonrpsee::server::middleware::rpc::RpcServiceT;
18use jsonrpsee::types::Request;
19use pin_project::pin_project;
20
21use super::{JSONRPC_API_REQUEST_DURATION_SECONDS, JSONRPC_API_REQUEST_RESPONSE_CODE};
22
23#[pin_project]
24pub struct ResponseFuture<F> {
25 method: &'static str,
26 #[pin]
27 fut: F,
28 timer: Option<HistogramTimer>,
29}
30
31impl<F> std::fmt::Debug for ResponseFuture<F> {
32 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 f.write_str("ResponseFuture")
34 }
35}
36
37impl<F: Future<Output = MethodResponse>> Future for ResponseFuture<F> {
38 type Output = F::Output;
39
40 fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
41 let projected = self.project();
42 let res = projected.fut.poll(cx);
43 if let Poll::Ready(res) = &res
44 && let Some(timer) = projected.timer.take()
45 {
46 timer.observe_duration();
47
48 JSONRPC_API_REQUEST_RESPONSE_CODE
49 .with_label_values(&[
50 *projected.method,
51 &if let Some(code) = res.as_error_code() {
52 Cow::Owned(code.to_string())
53 } else {
54 Cow::Borrowed("0")
55 },
56 if res.is_subscription() {
57 "subscription"
58 } else if res.is_batch() {
59 "batch"
60 } else {
61 "default"
62 },
63 ])
64 .inc();
65 }
66 res
67 }
68}
69
70const UNKNOWN_METHOD: &str = "unknown";
71
72#[derive(Clone, Debug)]
73pub struct MetricsLayer {
74 methods: Arc<HashSet<&'static str>>,
75}
76
77impl MetricsLayer {
78 pub fn new(methods: impl IntoIterator<Item = &'static str>) -> Self {
79 Self {
80 methods: Arc::new(methods.into_iter().collect()),
81 }
82 }
83}
84
85impl<S> tower::Layer<S> for MetricsLayer {
86 type Service = MetricsService<S>;
87
88 fn layer(&self, service: S) -> Self::Service {
89 MetricsService {
90 service,
91 methods: self.methods.clone(),
92 }
93 }
94}
95
96pub struct MetricsService<S> {
97 pub(crate) service: S,
98 methods: Arc<HashSet<&'static str>>,
99}
100
101impl<'a, S> RpcServiceT<'a> for MetricsService<S>
102where
103 S: RpcServiceT<'a> + Send + Sync,
104{
105 type Future = ResponseFuture<S::Future>;
106
107 fn call(&self, req: Request<'a>) -> Self::Future {
108 let method = self
109 .methods
110 .get(req.method_name())
111 .copied()
112 .unwrap_or(UNKNOWN_METHOD);
113 let timer = JSONRPC_API_REQUEST_DURATION_SECONDS
114 .with_label_values(&[method])
115 .start_timer();
116
117 ResponseFuture {
118 method,
119 fut: self.service.call(req),
120 timer: Some(timer),
121 }
122 }
123}
124
125#[cfg(test)]
126mod tests;