fedimint_server/metrics/
jsonrpsee.rs

1//! jsonrpsee/tower rpc layer that collects rpc stats
2//!
3//! Based on implementation of logger from:
4//!
5//! <https://github.com/paritytech/jsonrpsee/blob/bf5952fb663bdb8193b9f8a43182454c143b0e7d/server/src/middleware/rpc/layer/logger.rs#L1>
6
7use 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 {
127    use std::borrow::Cow;
128
129    use fedimint_metrics::prometheus::core::Collector;
130    use fedimint_metrics::prometheus::proto::Metric;
131    use futures::future::{Ready, ready};
132    use jsonrpsee::types::{Id, ResponsePayload};
133
134    use super::*;
135
136    const REGISTERED_METHOD: &str = "metrics_test_registered";
137    const UNREGISTERED_METHODS: [&str; 3] = [
138        "metrics_test_unregistered",
139        "metrics_test_unregistered_!@#$%^&*()",
140        "metrics_test_unregistered_with_a_very_long_attacker_controlled_suffix",
141    ];
142
143    struct SuccessService;
144
145    impl<'a> RpcServiceT<'a> for SuccessService {
146        type Future = Ready<MethodResponse>;
147
148        fn call(&self, _request: Request<'a>) -> Self::Future {
149            ready(MethodResponse::response(
150                Id::Number(1),
151                ResponsePayload::success("ok").into(),
152                usize::MAX,
153            ))
154        }
155    }
156
157    fn has_method_label(metric: &Metric, method: &str) -> bool {
158        metric
159            .get_label()
160            .iter()
161            .any(|label| label.name() == "method" && label.value() == method)
162    }
163
164    fn duration_count(method: &str) -> u64 {
165        JSONRPC_API_REQUEST_DURATION_SECONDS
166            .collect()
167            .into_iter()
168            .flat_map(|family| family.metric)
169            .filter(|metric| has_method_label(metric, method))
170            .map(|metric| metric.histogram.sample_count())
171            .sum()
172    }
173
174    fn response_count(method: &str) -> u64 {
175        JSONRPC_API_REQUEST_RESPONSE_CODE
176            .collect()
177            .into_iter()
178            .flat_map(|family| family.metric)
179            .filter(|metric| has_method_label(metric, method))
180            .map(|metric| metric.counter.value() as u64)
181            .sum()
182    }
183
184    #[tokio::test]
185    async fn bounds_method_labels_for_all_jsonrpc_metrics() {
186        let service = MetricsService {
187            service: SuccessService,
188            methods: Arc::new([REGISTERED_METHOD].into_iter().collect()),
189        };
190        let duration_before = [
191            duration_count(REGISTERED_METHOD),
192            duration_count(UNKNOWN_METHOD),
193        ];
194        let response_before = [
195            response_count(REGISTERED_METHOD),
196            response_count(UNKNOWN_METHOD),
197        ];
198
199        for method in std::iter::once(REGISTERED_METHOD).chain(UNREGISTERED_METHODS) {
200            service
201                .call(Request::new(Cow::Borrowed(method), None, Id::Number(1)))
202                .await;
203        }
204
205        assert_eq!(duration_count(REGISTERED_METHOD) - duration_before[0], 1);
206        assert_eq!(
207            duration_count(UNKNOWN_METHOD) - duration_before[1],
208            UNREGISTERED_METHODS.len() as u64
209        );
210        assert_eq!(response_count(REGISTERED_METHOD) - response_before[0], 1);
211        assert_eq!(
212            response_count(UNKNOWN_METHOD) - response_before[1],
213            UNREGISTERED_METHODS.len() as u64
214        );
215
216        for method in UNREGISTERED_METHODS {
217            assert_eq!(duration_count(method), 0);
218            assert_eq!(response_count(method), 0);
219        }
220    }
221}
fedimint_server/metrics/jsonrpsee.rs

fedimint_server/metrics/
jsonrpsee.rs
