siderite_core/middleware/
concurrency.rs

1//! Bounded handler admission and optional response-body ownership.
2
3use super::adapt::{BoxService, Next, from_fn, impl_layer};
4use crate::{ApiError, IntoResponse, Request, Response};
5use http::{StatusCode, header};
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{Duration, Instant};
9use tokio::sync::{OwnedSemaphorePermit, Semaphore};
10use tower::Layer;
11
12mod body;
13
14/// Invalid concurrency or queue-deadline configuration.
15#[derive(Debug, Clone, Copy, thiserror::Error)]
16#[error("invalid concurrency capacity or queue deadline")]
17pub struct ConcurrencyLimitError;
18
19/// Admission counters shared by all clones of a limit.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct ConcurrencyStats {
22    /// Requests or bodies currently holding execution permits.
23    pub active: usize,
24    /// Requests waiting for execution permits.
25    pub waiting: usize,
26    /// Requests admitted to execution since construction.
27    pub admitted: u64,
28    /// Requests rejected because the queue is full or closed.
29    pub rejected: u64,
30    /// Requests whose queue deadline expired.
31    pub timed_out: u64,
32    /// Cumulative queue time of completed admission attempts, nanoseconds.
33    pub wait_nanos: u64,
34}
35
36#[derive(Debug)]
37struct Shared {
38    active: Arc<Semaphore>,
39    waiting: Arc<Semaphore>,
40    max_active: usize,
41    max_waiting: usize,
42    admitted: AtomicU64,
43    rejected: AtomicU64,
44    timed_out: AtomicU64,
45    wait_nanos: AtomicU64,
46}
47
48/// Bounds execution and waiting, returning 503 with Retry-After on overload.
49///
50/// Defaults to one queue slot per execution slot and a five-second queue
51/// deadline. Permits normally end at response production. `hold_body(true)`
52/// retains them until body completion, error or drop. Upgraded connections
53/// require their own lifecycle budget. Clones share admission and counters.
54#[derive(Debug, Clone)]
55pub struct ConcurrencyLimit {
56    shared: Arc<Shared>,
57    wait: Duration,
58    body: bool,
59    valid: bool,
60}
61
62impl ConcurrencyLimit {
63    /// Configure execution capacity, treating zero as one for compatibility.
64    ///
65    /// Args:
66    ///     max: Maximum active requests, at most Semaphore::MAX_PERMITS.
67    ///
68    /// Returns:
69    ///     Bounded limit; invalid oversized configuration fails requests.
70    #[must_use]
71    pub fn new(max: usize) -> Self {
72        let capacity = max.clamp(1, Semaphore::MAX_PERMITS);
73        Self {
74            shared: Arc::new(Shared::new(capacity, capacity)),
75            wait: Duration::from_secs(5),
76            body: false,
77            valid: max <= Semaphore::MAX_PERMITS,
78        }
79    }
80
81    /// Validate a strictly positive execution capacity immediately.
82    ///
83    /// Args:
84    ///     max: Maximum active requests.
85    ///
86    /// Returns:
87    ///     Configured limit or invalid-capacity error.
88    pub fn try_new(max: usize) -> Result<Self, ConcurrencyLimitError> {
89        if max == 0 || max > Semaphore::MAX_PERMITS {
90            return Err(ConcurrencyLimitError);
91        }
92        Ok(Self::new(max))
93    }
94
95    /// Set bounded queue capacity and deadline, including a zero-slot queue.
96    ///
97    /// Args:
98    ///     max: Maximum waiting requests, at most Semaphore::MAX_PERMITS.
99    ///     timeout: Maximum admission wait; zero means immediate admission.
100    ///
101    /// Returns:
102    ///     Fresh admission state; existing clones keep their previous state.
103    #[must_use]
104    pub fn queue(mut self, max: usize, timeout: Duration) -> Self {
105        let capacity = max.min(Semaphore::MAX_PERMITS);
106        let deadline = Instant::now().checked_add(timeout).is_some();
107        self.valid &= max <= Semaphore::MAX_PERMITS && deadline;
108        self.shared = Arc::new(Shared::new(self.shared.max_active, capacity));
109        self.wait = timeout;
110        self
111    }
112
113    /// Include response-body consumption in the execution limit.
114    ///
115    /// Args:
116    ///     enabled: Retain permits until body completion, error or drop.
117    ///
118    /// Returns:
119    ///     Layer sharing its original admission state.
120    #[must_use]
121    pub fn hold_body(mut self, enabled: bool) -> Self {
122        self.body = enabled;
123        self
124    }
125
126    /// Close admission and wake waiters; already active work can complete.
127    ///
128    /// Returns:
129    ///     No value; all clones observe the closed state.
130    pub fn close(&self) {
131        self.shared.active.close();
132        self.shared.waiting.close();
133    }
134
135    /// Read an approximate concurrent snapshot without locking.
136    ///
137    /// Returns:
138    ///     Current capacities and cumulative admission counters.
139    pub fn stats(&self) -> ConcurrencyStats {
140        let state = &self.shared;
141        ConcurrencyStats {
142            active: state.max_active - state.active.available_permits(),
143            waiting: state.max_waiting - state.waiting.available_permits(),
144            admitted: state.admitted.load(Ordering::Relaxed),
145            rejected: state.rejected.load(Ordering::Relaxed),
146            timed_out: state.timed_out.load(Ordering::Relaxed),
147            wait_nanos: state.wait_nanos.load(Ordering::Relaxed),
148        }
149    }
150
151    pub(crate) fn validate(&self) -> Result<(), ConcurrencyLimitError> {
152        if self.valid {
153            Ok(())
154        } else {
155            Err(ConcurrencyLimitError)
156        }
157    }
158
159    async fn acquire(&self) -> Option<OwnedSemaphorePermit> {
160        let state = &self.shared;
161        if let Ok(permit) = state.active.clone().try_acquire_owned() {
162            state.admitted.fetch_add(1, Ordering::Relaxed);
163            return Some(permit);
164        }
165        let waiting = state.waiting.clone().try_acquire_owned();
166        let Ok(_waiting) = waiting else {
167            state.rejected.fetch_add(1, Ordering::Relaxed);
168            return None;
169        };
170        let started = Instant::now();
171        let future = state.active.clone().acquire_owned();
172        let result = tokio::time::timeout(self.wait, future).await;
173        let nanos = started.elapsed().as_nanos().min(u64::MAX as u128);
174        state.wait_nanos.fetch_add(nanos as u64, Ordering::Relaxed);
175        match result {
176            Ok(Ok(permit)) => {
177                state.admitted.fetch_add(1, Ordering::Relaxed);
178                Some(permit)
179            }
180            Ok(Err(_)) => {
181                state.rejected.fetch_add(1, Ordering::Relaxed);
182                None
183            }
184            Err(_) => {
185                state.timed_out.fetch_add(1, Ordering::Relaxed);
186                None
187            }
188        }
189    }
190
191    fn wrap(&self, inner: BoxService) -> BoxService {
192        let layer = self.clone();
193        from_fn(move |req: Request, next: Next| {
194            let layer = layer.clone();
195            async move {
196                if layer.validate().is_err() {
197                    let error = ApiError::internal("Invalid admission.");
198                    return error.into_response();
199                }
200                let Some(permit) = layer.acquire().await else {
201                    return overload();
202                };
203                let response = next.run(req).await;
204                if layer.body {
205                    body::retain(response, permit)
206                } else {
207                    drop(permit);
208                    response
209                }
210            }
211        })
212        .layer(inner)
213    }
214}
215
216impl Shared {
217    fn new(active: usize, waiting: usize) -> Self {
218        Self {
219            active: Arc::new(Semaphore::new(active)),
220            waiting: Arc::new(Semaphore::new(waiting)),
221            max_active: active,
222            max_waiting: waiting,
223            admitted: AtomicU64::new(0),
224            rejected: AtomicU64::new(0),
225            timed_out: AtomicU64::new(0),
226            wait_nanos: AtomicU64::new(0),
227        }
228    }
229}
230
231fn overload() -> Response {
232    let mut response = ApiError::new(
233        StatusCode::SERVICE_UNAVAILABLE,
234        "Request admission unavailable; retry later.",
235    )
236    .into_response();
237    response
238        .headers_mut()
239        .insert(header::RETRY_AFTER, http::HeaderValue::from_static("1"));
240    response
241}
242
243impl_layer!(ConcurrencyLimit);
siderite_core/middleware/concurrency.rs

siderite_core/middleware/
concurrency.rs
