use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;
use tokio::{sync, sync::mpsc};
use super::{
alive::AliveTracker,
registry::Registry,
runtime::{CoreSettings, SupervisorCore},
supervisor::Supervisor,
};
use crate::{
core::{ConfigError, SupervisorConfig, TaskDefaults},
events::Bus,
subscribers::{Subscribe, SubscriberSet},
};
#[must_use]
pub struct SupervisorBuilder {
runtime: SupervisorConfig,
task_defaults: TaskDefaults,
subscribers: Vec<Arc<dyn Subscribe>>,
#[cfg(feature = "controller")]
controller_config: Option<crate::controller::ControllerConfig>,
}
impl SupervisorBuilder {
pub fn new(runtime: SupervisorConfig) -> Self {
Self {
runtime,
task_defaults: TaskDefaults::default(),
subscribers: Vec::new(),
#[cfg(feature = "controller")]
controller_config: None,
}
}
pub fn with_runtime_config(mut self, runtime: SupervisorConfig) -> Self {
self.runtime = runtime;
self
}
pub fn with_task_defaults(mut self, task_defaults: TaskDefaults) -> Self {
self.task_defaults = task_defaults;
self
}
pub fn with_grace(mut self, grace: Duration) -> Self {
self.runtime = self.runtime.with_grace(grace);
self
}
pub fn with_subscriber_shutdown_timeout(mut self, timeout: Duration) -> Self {
self.runtime = self.runtime.with_subscriber_shutdown_timeout(timeout);
self
}
pub fn with_max_concurrent(mut self, max_concurrent: impl Into<Option<NonZeroUsize>>) -> Self {
self.runtime = self.runtime.with_max_concurrent(max_concurrent.into());
self
}
pub fn try_with_max_concurrent(mut self, max_concurrent: usize) -> Result<Self, ConfigError> {
self.runtime = self.runtime.try_with_max_concurrent(max_concurrent)?;
Ok(self)
}
pub fn try_with_bus_capacity(mut self, bus_capacity: usize) -> Result<Self, ConfigError> {
self.runtime = self.runtime.try_with_bus_capacity(bus_capacity)?;
Ok(self)
}
pub fn with_registry_queue_capacity(mut self, registry_queue_capacity: NonZeroUsize) -> Self {
self.runtime = self
.runtime
.with_registry_queue_capacity(registry_queue_capacity);
self
}
pub fn try_with_registry_queue_capacity(
mut self,
registry_queue_capacity: usize,
) -> Result<Self, ConfigError> {
self.runtime = self
.runtime
.try_with_registry_queue_capacity(registry_queue_capacity)?;
Ok(self)
}
pub fn with_bus_capacity(mut self, bus_capacity: NonZeroUsize) -> Self {
self.runtime = self.runtime.with_bus_capacity(bus_capacity);
self
}
pub fn with_subscribers(mut self, subscribers: Vec<Arc<dyn Subscribe>>) -> Self {
self.subscribers = subscribers;
self
}
#[cfg(feature = "controller")]
#[cfg_attr(docsrs, doc(cfg(feature = "controller")))]
pub fn with_controller(mut self, config: crate::controller::ControllerConfig) -> Self {
self.controller_config = Some(config);
self
}
#[must_use]
pub fn build(self) -> Arc<Supervisor> {
let bus = Bus::new(self.runtime.bus_capacity().get());
let subs = Arc::new(SubscriberSet::new_with_shutdown_timeout(
self.subscribers,
bus.clone(),
self.runtime.subscriber_shutdown_timeout(),
));
let runtime_token = tokio_util::sync::CancellationToken::new();
let semaphore = self
.runtime
.max_concurrent()
.map(|limit| Arc::new(sync::Semaphore::new(limit.get())));
let (cmd_tx, cmd_rx) = mpsc::channel(self.runtime.registry_queue_capacity().get());
let registry = Registry::new(
bus.clone(),
runtime_token.clone(),
semaphore,
self.runtime.grace(),
self.task_defaults.clone(),
cmd_rx,
);
let alive = Arc::new(AliveTracker::new());
let core = SupervisorCore::new_internal(
CoreSettings::new(self.runtime, self.task_defaults),
bus.clone(),
subs,
alive,
registry,
runtime_token,
cmd_tx,
);
#[cfg(feature = "controller")]
let controller = self
.controller_config
.map(|config| crate::controller::Controller::new(config, &core, bus.clone()));
#[cfg(feature = "controller")]
if let Some(controller) = &controller {
core.attach_controller(controller);
}
Supervisor::from_parts(
core,
#[cfg(feature = "controller")]
controller,
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{BackoffPolicy, RestartPolicy};
use std::num::NonZeroU32;
#[test]
fn builder_keeps_runtime_and_task_defaults_separate() {
let runtime = SupervisorConfig::default()
.with_grace(Duration::from_secs(30))
.with_subscriber_shutdown_timeout(Duration::from_secs(2))
.with_max_concurrent(NonZeroUsize::new(4))
.with_bus_capacity(NonZeroUsize::new(2048).unwrap())
.with_registry_queue_capacity(NonZeroUsize::new(256).unwrap());
let task_defaults = TaskDefaults::default()
.with_restart(RestartPolicy::Never)
.with_backoff(BackoffPolicy::constant(Duration::from_millis(50)))
.with_timeout(Duration::from_secs(5))
.with_max_retries(NonZeroU32::new(10));
let builder =
SupervisorBuilder::new(runtime.clone()).with_task_defaults(task_defaults.clone());
assert_eq!(builder.runtime.grace(), runtime.grace());
assert_eq!(
builder.runtime.subscriber_shutdown_timeout(),
runtime.subscriber_shutdown_timeout()
);
assert_eq!(builder.runtime.max_concurrent(), runtime.max_concurrent());
assert_eq!(builder.runtime.bus_capacity(), runtime.bus_capacity());
assert_eq!(
builder.runtime.registry_queue_capacity(),
runtime.registry_queue_capacity()
);
assert!(matches!(
builder.task_defaults.restart(),
RestartPolicy::Never
));
assert_eq!(
builder.task_defaults.backoff().first(),
Duration::from_millis(50)
);
assert_eq!(
builder.task_defaults.timeout(),
Some(Duration::from_secs(5))
);
assert_eq!(
builder.task_defaults.max_retries().map(NonZeroU32::get),
Some(10)
);
}
#[test]
fn raw_zero_values_return_errors_instead_of_panicking() {
type RawSetter = fn(SupervisorBuilder, usize) -> Result<SupervisorBuilder, ConfigError>;
let cases: [(&str, RawSetter); 3] = [
("max_concurrent", SupervisorBuilder::try_with_max_concurrent),
("bus_capacity", SupervisorBuilder::try_with_bus_capacity),
(
"registry_queue_capacity",
SupervisorBuilder::try_with_registry_queue_capacity,
),
];
for (field, set) in cases {
assert!(matches!(
set(SupervisorBuilder::new(SupervisorConfig::default()), 0),
Err(ConfigError::Zero { field: actual }) if actual == field
));
}
}
#[test]
fn max_concurrent_accepts_a_limit_or_an_option() {
let limit = NonZeroUsize::new(4).unwrap();
let direct = SupervisorBuilder::new(SupervisorConfig::default()).with_max_concurrent(limit);
let optional =
SupervisorBuilder::new(SupervisorConfig::default()).with_max_concurrent(Some(limit));
let cleared = SupervisorBuilder::new(SupervisorConfig::default()).with_max_concurrent(None);
assert_eq!(direct.runtime.max_concurrent(), Some(limit));
assert_eq!(optional.runtime.max_concurrent(), Some(limit));
assert_eq!(cleared.runtime.max_concurrent(), None);
}
}