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#[derive(Debug, Clone, Copy)]
8pub struct Full;
9
10#[derive(Debug)]
11pub struct Bulkhead {
12    key: String,
13    limit: usize,
14    semaphore: Arc<Semaphore>,
15}
16
17impl Bulkhead {
18    pub fn new(key: impl Into<String>, max_concurrent: usize) -> Self {
19        Self {
20            key: key.into(),
21            limit: max_concurrent,
22            semaphore: Arc::new(Semaphore::new(max_concurrent)),
23        }
24    }
25
26    // The returned permit must be held for the call's duration (and, for
27    // streaming responses, the stream's lifetime).
28    pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, Full> {
29        Arc::clone(&self.semaphore)
30            .try_acquire_owned()
31            .map_err(|_e| {
32                tracing::warn!(
33                    key = %self.key,
34                    limit = self.limit,
35                    "bulkhead saturated, rejecting call",
36                );
37                Full
38            })
39    }
40
41    #[must_use]
42    pub const fn limit(&self) -> usize {
43        self.limit
44    }
45}