reth_rpc_builder/
rate_limiter.rs

1//! [`jsonrpsee`] helper layer for rate limiting certain methods.
2
3use jsonrpsee::{server::middleware::rpc::RpcServiceT, types::Request, MethodResponse};
4use std::{
5    future::Future,
6    pin::Pin,
7    sync::Arc,
8    task::{ready, Context, Poll},
9};
10use tokio::sync::{OwnedSemaphorePermit, Semaphore};
11use tokio_util::sync::PollSemaphore;
12use tower::Layer;
13
14/// Rate limiter for the RPC server.
15///
16/// Rate limits expensive calls such as debug_ and trace_.
17#[derive(Debug, Clone)]
18pub struct RpcRequestRateLimiter {
19    inner: Arc<RpcRequestRateLimiterInner>,
20}
21
22impl RpcRequestRateLimiter {
23    /// Create a new rate limit layer with the given number of permits.
24    pub fn new(rate_limit: usize) -> Self {
25        Self {
26            inner: Arc::new(RpcRequestRateLimiterInner {
27                call_guard: PollSemaphore::new(Arc::new(Semaphore::new(rate_limit))),
28            }),
29        }
30    }
31}
32
33impl<S> Layer<S> for RpcRequestRateLimiter {
34    type Service = RpcRequestRateLimitingService<S>;
35
36    fn layer(&self, inner: S) -> Self::Service {
37        RpcRequestRateLimitingService::new(inner, self.clone())
38    }
39}
40
41/// Rate Limiter for the RPC server
42#[derive(Debug, Clone)]
43struct RpcRequestRateLimiterInner {
44    /// Semaphore to rate limit calls
45    call_guard: PollSemaphore,
46}
47
48/// A [`RpcServiceT`] middleware that rate limits RPC calls to the server.
49#[derive(Debug, Clone)]
50pub struct RpcRequestRateLimitingService<S> {
51    /// The rate limiter for RPC requests
52    rate_limiter: RpcRequestRateLimiter,
53    /// The inner service being wrapped
54    inner: S,
55}
56
57impl<S> RpcRequestRateLimitingService<S> {
58    /// Create a new rate limited service.
59    pub const fn new(service: S, rate_limiter: RpcRequestRateLimiter) -> Self {
60        Self { inner: service, rate_limiter }
61    }
62}
63
64impl<S> RpcServiceT for RpcRequestRateLimitingService<S>
65where
66    S: RpcServiceT + Send + Sync + Clone + 'static,
67{
68    fn call<'a>(&self, req: Request<'a>) -> impl Future<Output = MethodResponse> + Send + 'a {
69        let method_name = req.method_name();
70        if method_name.starts_with("trace_") || method_name.starts_with("debug_") {
71            RateLimitingRequestFuture {
72                fut: self.inner.call(req),
73                guard: Some(self.rate_limiter.inner.call_guard.clone()),
74                permit: None,
75            }
76        } else {
77            // if we don't need to rate limit, then there
78            // is no need to get a semaphore permit
79            RateLimitingRequestFuture { fut: self.inner.call(req), guard: None, permit: None }
80        }
81    }
82
83    fn batch<'a>(
84        &self,
85        requests: jsonrpsee::core::middleware::Batch<'a>,
86    ) -> impl Future<Output = MethodResponse> + Send + 'a {
87        self.inner.batch(requests)
88    }
89
90    fn notification<'a>(
91        &self,
92        n: jsonrpsee::core::middleware::Notification<'a>,
93    ) -> impl Future<Output = MethodResponse> + Send + 'a {
94        self.inner.notification(n)
95    }
96}
97
98/// Response future.
99#[pin_project::pin_project]
100pub struct RateLimitingRequestFuture<F> {
101    #[pin]
102    fut: F,
103    guard: Option<PollSemaphore>,
104    permit: Option<OwnedSemaphorePermit>,
105}
106
107impl<F> std::fmt::Debug for RateLimitingRequestFuture<F> {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.write_str("RateLimitingRequestFuture")
110    }
111}
112
113impl<F: Future> Future for RateLimitingRequestFuture<F> {
114    type Output = F::Output;
115
116    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
117        let this = self.project();
118        if let Some(guard) = this.guard.as_mut() {
119            *this.permit = ready!(guard.poll_acquire(cx));
120            *this.guard = None;
121        }
122        let res = this.fut.poll(cx);
123        if res.is_ready() {
124            *this.permit = None;
125        }
126        res
127    }
128}
reth_rpc_builder/rate_limiter.rs

reth_rpc_builder/
rate_limiter.rs
