Skip to main content

systemprompt_database/resilience/
bulkhead.rs

1//! A non-blocking concurrency limiter.
2
3use std::sync::Arc;
4
5use tokio::sync::{OwnedSemaphorePermit, Semaphore};
6
7/// Returned by [`Bulkhead::try_acquire`] when the concurrency limit is reached.
8#[derive(Debug, Clone, Copy)]
9pub struct Full;
10
11/// A concurrency cap for one dependency.
12///
13/// Acquisition is non-blocking: when the limit is reached the call is rejected
14/// with [`Full`] rather than queued, so a slow dependency fast-fails callers
15/// instead of letting them pile up and exhaust workers.
16#[derive(Debug)]
17pub struct Bulkhead {
18    key: String,
19    limit: usize,
20    semaphore: Arc<Semaphore>,
21}
22
23impl Bulkhead {
24    pub fn new(key: impl Into<String>, max_concurrent: usize) -> Self {
25        Self {
26            key: key.into(),
27            limit: max_concurrent,
28            semaphore: Arc::new(Semaphore::new(max_concurrent)),
29        }
30    }
31
32    // The returned permit must be held for the call's duration (and, for
33    // streaming responses, the stream's lifetime).
34    pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, Full> {
35        Arc::clone(&self.semaphore)
36            .try_acquire_owned()
37            .map_err(|_e| {
38                tracing::warn!(
39                    key = %self.key,
40                    limit = self.limit,
41                    "bulkhead saturated, rejecting call",
42                );
43                Full
44            })
45    }
46
47    /// The configured concurrency limit.
48    #[must_use]
49    pub const fn limit(&self) -> usize {
50        self.limit
51    }
52}