use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::Arc;
use std::time::Duration;
use tokio::{sync, sync::mpsc};
use super::{
alive::AliveTracker, registry::Registry, runtime::SupervisorCore, supervisor::Supervisor,
};
use crate::{
core::SupervisorConfig,
events::Bus,
policies::{BackoffPolicy, RestartPolicy},
subscribers::{Subscribe, SubscriberSet},
};
pub struct SupervisorBuilder {
cfg: SupervisorConfig,
subscribers: Vec<Arc<dyn Subscribe>>,
#[cfg(feature = "controller")]
controller_config: Option<crate::controller::ControllerConfig>,
}
impl SupervisorBuilder {
pub fn new(cfg: SupervisorConfig) -> Self {
Self {
cfg,
subscribers: Vec::new(),
#[cfg(feature = "controller")]
controller_config: None,
}
}
pub fn with_grace(mut self, grace: Duration) -> Self {
self.cfg.grace = grace;
self
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.cfg.timeout = Some(timeout).filter(|d| !d.is_zero());
self
}
pub fn with_max_retries(mut self, max_retries: u32) -> Self {
assert!(
max_retries > 0,
"with_max_retries(0) is invalid: zero retries cannot be represented; \
use RestartPolicy::Never to stop after the first failure"
);
self.cfg.max_retries = NonZeroU32::new(max_retries);
self
}
pub fn with_max_concurrent(mut self, max_concurrent: usize) -> Self {
assert!(
max_concurrent > 0,
"with_max_concurrent(0) is invalid: no task could ever run; \
omit the call for unlimited concurrency"
);
self.cfg.max_concurrent = NonZeroUsize::new(max_concurrent);
self
}
pub fn with_bus_capacity(mut self, bus_capacity: usize) -> Self {
self.cfg.bus_capacity = bus_capacity;
self
}
pub fn with_restart(mut self, restart: RestartPolicy) -> Self {
self.cfg.restart = restart;
self
}
pub fn with_backoff(mut self, backoff: BackoffPolicy) -> Self {
self.cfg.backoff = backoff;
self
}
pub fn with_subscribers(mut self, subscribers: Vec<Arc<dyn Subscribe>>) -> Self {
self.subscribers = subscribers;
self
}
#[cfg(feature = "controller")]
pub fn with_controller(mut self, config: crate::controller::ControllerConfig) -> Self {
self.controller_config = Some(config);
self
}
pub fn build(self) -> Arc<Supervisor> {
let bus = Bus::new(self.cfg.bus_capacity_clamped());
let subs = Arc::new(SubscriberSet::new(self.subscribers, bus.clone()));
let runtime_token = tokio_util::sync::CancellationToken::new();
let semaphore = self
.cfg
.max_concurrent
.map(|n| Arc::new(sync::Semaphore::new(n.get())));
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
let registry = Registry::new(
bus.clone(),
runtime_token.clone(),
semaphore,
self.cfg.grace,
cmd_rx,
);
let alive = Arc::new(AliveTracker::new());
let core = SupervisorCore::new_internal(
self.cfg,
bus.clone(),
subs,
alive,
registry,
runtime_token.clone(),
cmd_tx,
);
#[cfg(feature = "controller")]
let controller = self
.controller_config
.map(|ctrl_cfg| crate::controller::Controller::new(ctrl_cfg, &core, bus.clone()));
Supervisor::from_parts(
core,
#[cfg(feature = "controller")]
controller,
#[cfg(feature = "controller")]
runtime_token,
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::policies::{BackoffPolicy, RestartPolicy};
use std::time::Duration;
#[test]
fn setters_override_config_fields() {
let backoff = BackoffPolicy::exponential(Duration::from_millis(200));
let b = SupervisorBuilder::new(SupervisorConfig::default())
.with_grace(Duration::from_secs(30))
.with_timeout(Duration::from_secs(5))
.with_max_retries(10)
.with_max_concurrent(4)
.with_bus_capacity(2048)
.with_restart(RestartPolicy::Never)
.with_backoff(backoff);
assert_eq!(b.cfg.grace, Duration::from_secs(30));
assert_eq!(b.cfg.timeout, Some(Duration::from_secs(5)));
assert_eq!(b.cfg.max_retries.map(|n| n.get()), Some(10));
assert_eq!(b.cfg.max_concurrent.map(|n| n.get()), Some(4));
assert_eq!(b.cfg.bus_capacity, 2048);
assert!(
matches!(b.cfg.restart, RestartPolicy::Never),
"with_restart must store the given policy"
);
assert_eq!(
b.cfg.backoff.factor(),
2.0,
"with_backoff must store the given policy"
);
}
#[test]
fn with_timeout_zero_means_no_timeout() {
let b = SupervisorBuilder::new(SupervisorConfig {
timeout: Some(Duration::from_secs(5)),
..Default::default()
})
.with_timeout(Duration::ZERO);
assert_eq!(
b.cfg.timeout, None,
"a zero timeout must normalize to None (no timeout)"
);
}
#[test]
#[should_panic(expected = "with_max_retries(0)")]
fn with_max_retries_zero_panics() {
let _ = SupervisorBuilder::new(SupervisorConfig::default()).with_max_retries(0);
}
#[test]
#[should_panic(expected = "with_max_concurrent(0)")]
fn with_max_concurrent_zero_panics() {
let _ = SupervisorBuilder::new(SupervisorConfig::default()).with_max_concurrent(0);
}
}