iota_network_stack/
concurrency.rs

1// Copyright (c) 2026 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! Per-service concurrency limiting for gRPC servers.
5//!
6//! Unlike [`tower::limit::GlobalConcurrencyLimitLayer`] applied around a whole
7//! server, [`ServiceConcurrencyLimit`] bounds the in-flight requests of a
8//! single gRPC service, so services sharing one listener cannot crowd each
9//! other out of admission slots.
10//!
11//! Tower's per-service [`tower::limit::ConcurrencyLimitLayer`] cannot be used
12//! here: its `ConcurrencyLimit` wrapper does not implement tonic's
13//! [`NamedService`], which `Routes::add_service` requires for routing, and
14//! shedding through tower's `LoadShed` surfaces as a `BoxError`, incompatible
15//! with the router's `Error = Infallible` bound — over-limit requests must be
16//! answered in-band with a gRPC `RESOURCE_EXHAUSTED` response instead.
17
18use std::{
19    convert::Infallible,
20    num::NonZeroUsize,
21    sync::Arc,
22    task::{Context, Poll},
23};
24
25use futures::future::BoxFuture;
26use tokio::sync::Semaphore;
27use tonic::{
28    body::Body,
29    codegen::http::{Request, Response},
30    server::NamedService,
31};
32use tower::Service;
33
34/// Bounds the number of concurrent in-flight requests to the wrapped gRPC
35/// service, independently of any other service registered on the same server.
36///
37/// With `load_shed` enabled, requests over the limit are rejected immediately
38/// with gRPC `RESOURCE_EXHAUSTED`; otherwise they wait for a slot to free up.
39/// Clones share the same limit.
40#[derive(Clone)]
41pub struct ServiceConcurrencyLimit<S> {
42    inner: S,
43    semaphore: Arc<Semaphore>,
44    load_shed: bool,
45}
46
47impl<S> ServiceConcurrencyLimit<S> {
48    pub fn new(inner: S, limit: NonZeroUsize, load_shed: bool) -> Self {
49        Self {
50            inner,
51            // Clamp: `Semaphore::new` panics above `MAX_PERMITS`, and
52            // effectively-unlimited configs multiply large values by the CPU
53            // core count.
54            semaphore: Arc::new(Semaphore::new(limit.get().min(Semaphore::MAX_PERMITS))),
55            load_shed,
56        }
57    }
58}
59
60impl<S> Service<Request<Body>> for ServiceConcurrencyLimit<S>
61where
62    S: Service<Request<Body>, Response = Response<Body>, Error = Infallible>
63        + Clone
64        + Send
65        + 'static,
66    S::Future: Send + 'static,
67{
68    type Response = Response<Body>;
69    type Error = Infallible;
70    type Future = BoxFuture<'static, Result<Response<Body>, Infallible>>;
71
72    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
73        self.inner.poll_ready(cx)
74    }
75
76    fn call(&mut self, request: Request<Body>) -> Self::Future {
77        // Admission happens here rather than in `poll_ready` (where
78        // `tower::limit::ConcurrencyLimit` acquires its permit): a shed
79        // request must be answered with a gRPC response, and `poll_ready`
80        // can only signal Ready or Pending — converting Pending into an
81        // error via an outer load-shed layer is ruled out by the router's
82        // `Error = Infallible` bound. Readiness-based acquisition would
83        // also buy no upstream backpressure: the axum router dispatches
84        // every request on a fresh clone of this service.
85        //
86        // Take the instance that was driven to readiness and leave the clone
87        // for later calls, as `poll_ready` readiness does not transfer to
88        // clones.
89        let clone = self.inner.clone();
90        let mut inner = std::mem::replace(&mut self.inner, clone);
91        let semaphore = self.semaphore.clone();
92        let load_shed = self.load_shed;
93
94        Box::pin(async move {
95            // The permit is held until the response future resolves, mirroring
96            // `tower::limit::ConcurrencyLimit`.
97            let _permit = if load_shed {
98                match semaphore.try_acquire_owned() {
99                    Ok(permit) => permit,
100                    Err(_) => {
101                        return Ok(tonic::Status::resource_exhausted(
102                            "service concurrency limit reached",
103                        )
104                        .into_http());
105                    }
106                }
107            } else {
108                semaphore
109                    .acquire_owned()
110                    .await
111                    .expect("the semaphore is never closed")
112            };
113            inner.call(request).await
114        })
115    }
116}
117
118impl<S: NamedService> NamedService for ServiceConcurrencyLimit<S> {
119    const NAME: &'static str = S::NAME;
120}
121
122#[cfg(test)]
123mod tests {
124    use std::time::Duration;
125
126    use tower::ServiceExt;
127
128    use super::*;
129
130    /// Inner service whose responses only complete once `release` is
131    /// notified, keeping requests in flight for as long as the test needs.
132    #[derive(Clone)]
133    struct BlockingService {
134        release: Arc<tokio::sync::Notify>,
135    }
136
137    impl Service<Request<Body>> for BlockingService {
138        type Response = Response<Body>;
139        type Error = Infallible;
140        type Future = BoxFuture<'static, Result<Response<Body>, Infallible>>;
141
142        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
143            Poll::Ready(Ok(()))
144        }
145
146        fn call(&mut self, _request: Request<Body>) -> Self::Future {
147            let release = self.release.clone();
148            Box::pin(async move {
149                release.notified().await;
150                Ok(Response::new(Body::default()))
151            })
152        }
153    }
154
155    fn request() -> Request<Body> {
156        Request::new(Body::default())
157    }
158
159    #[tokio::test]
160    async fn load_shedding_rejects_requests_over_the_limit() {
161        let release = Arc::new(tokio::sync::Notify::new());
162        let service = ServiceConcurrencyLimit::new(
163            BlockingService {
164                release: release.clone(),
165            },
166            NonZeroUsize::MIN,
167            true,
168        );
169
170        let in_flight = tokio::spawn(service.clone().oneshot(request()));
171        tokio::task::yield_now().await;
172
173        let shed = service.clone().oneshot(request()).await.unwrap();
174        assert_eq!(
175            shed.headers().get("grpc-status").unwrap(),
176            &(tonic::Code::ResourceExhausted as i32).to_string()
177        );
178
179        release.notify_one();
180        let response = in_flight.await.unwrap().unwrap();
181        assert!(response.headers().get("grpc-status").is_none());
182    }
183
184    #[tokio::test]
185    async fn without_load_shedding_requests_over_the_limit_wait() {
186        let release = Arc::new(tokio::sync::Notify::new());
187        let service = ServiceConcurrencyLimit::new(
188            BlockingService {
189                release: release.clone(),
190            },
191            NonZeroUsize::MIN,
192            false,
193        );
194
195        let first = tokio::spawn(service.clone().oneshot(request()));
196        tokio::task::yield_now().await;
197
198        let mut second = tokio::spawn(service.clone().oneshot(request()));
199        let waiting = tokio::time::timeout(Duration::from_millis(50), &mut second).await;
200        assert!(waiting.is_err(), "second request should wait for a slot");
201
202        release.notify_one();
203        first.await.unwrap().unwrap();
204        release.notify_one();
205        second.await.unwrap().unwrap();
206    }
207
208    #[tokio::test]
209    async fn limits_are_independent_per_service() {
210        let release = Arc::new(tokio::sync::Notify::new());
211        let blocking = BlockingService {
212            release: release.clone(),
213        };
214        let saturated = ServiceConcurrencyLimit::new(blocking.clone(), NonZeroUsize::MIN, true);
215        let other = ServiceConcurrencyLimit::new(blocking, NonZeroUsize::MIN, true);
216
217        let in_flight = tokio::spawn(saturated.clone().oneshot(request()));
218        tokio::task::yield_now().await;
219
220        // The other service has its own semaphore and still admits requests.
221        let admitted = tokio::spawn(other.oneshot(request()));
222        tokio::task::yield_now().await;
223
224        release.notify_waiters();
225        assert!(
226            in_flight
227                .await
228                .unwrap()
229                .unwrap()
230                .headers()
231                .get("grpc-status")
232                .is_none()
233        );
234        assert!(
235            admitted
236                .await
237                .unwrap()
238                .unwrap()
239                .headers()
240                .get("grpc-status")
241                .is_none()
242        );
243    }
244}
iota_network_stack/concurrency.rs

iota_network_stack/
concurrency.rs
