use crate::error::LimiteronError;
use dashmap::DashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
#[derive(Debug, Clone, Copy)]
pub struct BulkheadConfig {
pub max_concurrent: usize,
}
impl Default for BulkheadConfig {
fn default() -> Self {
Self { max_concurrent: 32 }
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum BulkheadError {
#[error("bulkhead '{0}' is full (max_concurrent exhausted)")]
Full(String),
#[cfg(feature = "circuit-breaker")]
#[error("bulkhead '{0}' circuit breaker is open")]
CircuitOpen(String),
}
impl crate::i18n::LocalizedMsg for BulkheadError {
fn message_key(&self) -> &'static str {
match self {
BulkheadError::Full(_) => "bulkhead-full",
#[cfg(feature = "circuit-breaker")]
BulkheadError::CircuitOpen(_) => "bulkhead-circuit-open",
}
}
fn message_args(&self) -> Vec<(&'static str, String)> {
match self {
BulkheadError::Full(name) => vec![("name", name.clone())],
#[cfg(feature = "circuit-breaker")]
BulkheadError::CircuitOpen(name) => vec![("name", name.clone())],
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct BulkheadStats {
pub active: i64,
pub admitted: u64,
pub rejected: u64,
pub max_concurrent: usize,
}
pub struct Bulkhead {
name: String,
config: BulkheadConfig,
active: AtomicI64,
admitted: AtomicU64,
rejected: AtomicU64,
#[cfg(feature = "circuit-breaker")]
circuit_breaker: parking_lot::RwLock<Option<Arc<crate::circuit::CircuitBreaker>>>,
}
impl Bulkhead {
pub fn name(&self) -> &str {
&self.name
}
pub fn max_concurrent(&self) -> usize {
self.config.max_concurrent
}
pub fn stats(&self) -> BulkheadStats {
BulkheadStats {
active: self.active.load(Ordering::Relaxed),
admitted: self.admitted.load(Ordering::Relaxed),
rejected: self.rejected.load(Ordering::Relaxed),
max_concurrent: self.config.max_concurrent,
}
}
#[cfg(feature = "circuit-breaker")]
pub fn set_circuit_breaker(&self, cb: Option<Arc<crate::circuit::CircuitBreaker>>) {
*self.circuit_breaker.write() = cb;
}
fn try_acquire(self: &Arc<Self>) -> Result<BulkheadPermit, BulkheadError> {
loop {
let current = self.active.load(Ordering::Relaxed);
if current >= self.config.max_concurrent as i64 {
self.rejected.fetch_add(1, Ordering::Relaxed);
return Err(BulkheadError::Full(self.name.clone()));
}
match self.active.compare_exchange(
current,
current + 1,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => {
self.admitted.fetch_add(1, Ordering::Relaxed);
return Ok(BulkheadPermit {
bulkhead: Arc::clone(self),
});
}
Err(_) => continue,
}
}
}
}
pub struct BulkheadPermit {
bulkhead: Arc<Bulkhead>,
}
impl Drop for BulkheadPermit {
fn drop(&mut self) {
self.bulkhead.active.fetch_sub(1, Ordering::AcqRel);
}
}
pub struct BulkheadRegistry {
pools: DashMap<String, Arc<Bulkhead>>,
default_config: BulkheadConfig,
}
impl Default for BulkheadRegistry {
fn default() -> Self {
Self::new()
}
}
impl BulkheadRegistry {
pub fn new() -> Self {
Self {
pools: DashMap::new(),
default_config: BulkheadConfig::default(),
}
}
pub fn with_default_config(mut self, config: BulkheadConfig) -> Self {
self.default_config = config;
self
}
pub fn with_pool_config(self, name: impl Into<String>, config: BulkheadConfig) -> Self {
let name = name.into();
self.pools.insert(
name.clone(),
Arc::new(Bulkhead {
name,
config,
active: AtomicI64::new(0),
admitted: AtomicU64::new(0),
rejected: AtomicU64::new(0),
#[cfg(feature = "circuit-breaker")]
circuit_breaker: parking_lot::RwLock::new(None),
}),
);
self
}
pub fn bulkhead(&self, name: &str) -> Arc<Bulkhead> {
if let Some(existing) = self.pools.get(name) {
return Arc::clone(existing.value());
}
let mut entry = self.pools.entry(name.to_string()).or_insert_with(|| {
Arc::new(Bulkhead {
name: name.to_string(),
config: self.default_config,
active: AtomicI64::new(0),
admitted: AtomicU64::new(0),
rejected: AtomicU64::new(0),
#[cfg(feature = "circuit-breaker")]
circuit_breaker: parking_lot::RwLock::new(None),
})
});
Arc::clone(entry.value_mut())
}
pub async fn execute<F, T>(&self, pool: &str, operation: F) -> Result<T, BulkheadError>
where
F: std::future::Future<Output = T>,
{
let bulkhead = self.bulkhead(pool);
#[cfg(feature = "circuit-breaker")]
{
let cb = bulkhead.circuit_breaker.read().clone();
if let Some(cb) = cb
&& cb.is_open().await
{
bulkhead.rejected.fetch_add(1, Ordering::Relaxed);
return Err(BulkheadError::CircuitOpen(bulkhead.name().to_string()));
}
}
let _permit = bulkhead.try_acquire()?;
Ok(operation.await)
}
pub fn stats(&self, pool: &str) -> Option<BulkheadStats> {
self.pools.get(pool).map(|b| b.stats())
}
pub fn all_stats(&self) -> Vec<(String, BulkheadStats)> {
let mut out: Vec<_> = self
.pools
.iter()
.map(|e| (e.key().clone(), e.value().stats()))
.collect();
out.sort_by(|a, b| a.0.cmp(&b.0));
out
}
pub fn pool_names(&self) -> Vec<String> {
let mut names: Vec<_> = self.pools.iter().map(|e| e.key().clone()).collect();
names.sort();
names
}
}
impl From<BulkheadError> for LimiteronError {
fn from(e: BulkheadError) -> Self {
LimiteronError::ConcurrencyLimitExceeded(e.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn registry() -> BulkheadRegistry {
BulkheadRegistry::new().with_default_config(BulkheadConfig { max_concurrent: 2 })
}
#[tokio::test]
async fn test_t611_pool_isolation_group_a_full_group_b_unaffected() {
let registry = registry();
let bh_a = registry.bulkhead("group-a");
let p1 = bh_a
.try_acquire()
.expect("an idle pool should be acquirable");
let p2 = bh_a
.try_acquire()
.expect("the second slot should be acquirable");
let err = registry
.execute("group-a", async { 3 })
.await
.expect_err("pool A must fail fast once full");
assert_eq!(err, BulkheadError::Full("group-a".to_string()));
let out = registry
.execute("group-b", async { "ok" })
.await
.expect("pool B must not be affected by pool A saturation");
assert_eq!(out, "ok");
drop((p1, p2));
assert_eq!(registry.stats("group-a").unwrap().active, 0);
}
#[tokio::test]
async fn test_t611_permit_released_after_completion() {
let registry = registry();
registry.execute("p", async {}).await.unwrap();
registry
.execute("p", async { Err::<(), _>(BulkheadError::Full("x".into())) })
.await
.ok(); registry.execute("p", async {}).await.unwrap();
let stats = registry.stats("p").unwrap();
assert_eq!(
stats.active, 0,
"in-flight must return to zero after completion"
);
assert_eq!(stats.admitted, 3);
}
#[tokio::test]
async fn test_t611_rejected_counter_and_stats() {
let registry = registry();
let bh = registry.bulkhead("s");
let hold1 = bh.try_acquire().unwrap();
let hold2 = bh.try_acquire().unwrap();
let _ = registry.execute("s", async {}).await; let stats = registry.stats("s").unwrap();
assert_eq!(stats.active, 2);
assert_eq!(stats.rejected, 1);
assert_eq!(stats.max_concurrent, 2);
drop((hold1, hold2));
}
#[test]
fn test_t611_get_or_create_identity() {
let registry = registry();
let a = registry.bulkhead("x");
let b = registry.bulkhead("x");
assert!(
Arc::ptr_eq(&a, &b),
"same-named pools must get-or-create the same instance"
);
assert_eq!(registry.pool_names(), vec!["x".to_string()]);
}
#[test]
fn test_t611_all_stats_sorted() {
let registry = registry().with_pool_config("b-pool", BulkheadConfig { max_concurrent: 4 });
let _ = registry.bulkhead("a-pool");
let all = registry.all_stats();
let names: Vec<_> = all.iter().map(|(n, _)| n.clone()).collect();
assert_eq!(names, vec!["a-pool".to_string(), "b-pool".to_string()]);
assert_eq!(all[1].1.max_concurrent, 4);
}
#[cfg(feature = "circuit-breaker")]
#[tokio::test]
async fn test_t611_pool_circuit_open_fails_fast() {
use crate::CircuitState;
use crate::circuit::{CircuitBreaker, CircuitBreakerConfig};
let registry = BulkheadRegistry::new()
.with_pool_config("cb-pool", BulkheadConfig { max_concurrent: 4 });
let pool = registry.bulkhead("cb-pool");
let cb = Arc::new(CircuitBreaker::with_dependencies(CircuitBreakerConfig {
failure_threshold: 1,
..CircuitBreakerConfig::default()
}));
pool.set_circuit_breaker(Some(cb.clone()));
let _ = cb
.execute(|| async { Err::<(), _>(LimiteronError::Other("boom".to_string())) })
.await;
assert_eq!(cb.get_state().await, CircuitState::Open);
let err = registry
.execute("cb-pool", async { 1 })
.await
.expect_err("a pool with an open circuit must fail fast");
assert!(
matches!(err, BulkheadError::CircuitOpen(ref n) if n == "cb-pool"),
"should return CircuitOpen, got: {err:?}"
);
let out = registry.execute("other-pool", async { 2 }).await.unwrap();
assert_eq!(out, 2);
}
}