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    pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, Full> {
30        Arc::clone(&self.semaphore)
31            .try_acquire_owned()
32            .map_err(|_e| {
33                tracing::warn!(
34                    key = %self.key,
35                    limit = self.limit,
36                    "bulkhead saturated, rejecting call",
37                );
38                Full
39            })
40    }
41
42    #[must_use]
43    pub const fn limit(&self) -> usize {
44        self.limit
45    }
46}