Skip to main content

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;