Skip to main content

systemprompt_database/resilience/
bulkhead.rs

1//! A non-blocking concurrency limiter.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use std::sync::Arc;
7
8use tokio::sync::{OwnedSemaphorePermit, Semaphore};
9
10#[derive(Debug, Clone, Copy)]
11pub struct Full;
12
13#[derive(Debug)]
14pub struct Bulkhead {
15    key: String,
16    limit: usize,
17    semaphore: Arc<Semaphore>,
18}
19
20impl Bulkhead {
21    pub fn new(key: impl Into<String>, max_concurrent: usize) -> Self {
22        Self {
23            key: key.into(),
24            limit: max_concurrent,
25            semaphore: Arc::new(Semaphore::new(max_concurrent)),
26        }
27    }
28
29    // Why: The returned permit must be held for the call's duration (and, for
30    // streaming responses, the stream's lifetime).
31    pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, Full> {
32        Arc::clone(&self.semaphore)
33            .try_acquire_owned()
34            .map_err(|_e| {
35                tracing::warn!(
36                    key = %self.key,
37                    limit = self.limit,
38                    "bulkhead saturated, rejecting call",
39                );
40                Full
41            })
42    }
43
44    #[must_use]
45    pub const fn limit(&self) -> usize {
46        self.limit
47    }
48}