use std::{
future::Future,
sync::{Arc, Weak},
};
use dashmap::DashMap;
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use tokio::sync::broadcast;
use crate::{
RuntimeError, TaskSpec,
core::{OutcomeTx, SupervisorCore, TaskOutcome},
events::{Bus, Event, EventKind},
identity::TaskId,
};
use super::{
admission::AdmissionPolicy,
config::ControllerConfig,
error::ControllerError,
slot::{SlotState, SlotStatus},
spec::ControllerSpec,
};
mod introspect;
mod recovery;
mod shutdown;
fn schedule_recovery(recovery_at: &mut Option<Instant>, candidate: Instant) -> Option<Instant> {
match recovery_at {
Some(current) if *current <= candidate => None,
_ => {
*recovery_at = Some(candidate);
Some(candidate)
}
}
}
struct Submission {
id: TaskId,
spec: ControllerSpec,
done: Option<OutcomeTx>,
}
#[derive(Clone)]
pub(crate) struct ControllerHandle {
tx: mpsc::Sender<Submission>,
}
impl ControllerHandle {
pub async fn submit(&self, spec: ControllerSpec) -> Result<TaskId, ControllerError> {
let id = TaskId::next();
self.tx
.send(Submission {
id,
spec,
done: None,
})
.await
.map_err(|_| ControllerError::Closed)?;
Ok(id)
}
pub fn try_submit(&self, spec: ControllerSpec) -> Result<TaskId, ControllerError> {
let id = TaskId::next();
self.tx
.try_send(Submission {
id,
spec,
done: None,
})
.map_err(|e| match e {
mpsc::error::TrySendError::Full(_) => ControllerError::Full,
mpsc::error::TrySendError::Closed(_) => ControllerError::Closed,
})?;
Ok(id)
}
pub async fn submit_and_watch(
&self,
spec: ControllerSpec,
) -> Result<(TaskId, oneshot::Receiver<TaskOutcome>), ControllerError> {
let id = TaskId::next();
let (tx, rx) = oneshot::channel();
self.tx
.send(Submission {
id,
spec,
done: Some(tx),
})
.await
.map_err(|_| ControllerError::Closed)?;
Ok((id, rx))
}
}
pub(crate) struct Controller {
config: ControllerConfig,
supervisor: Weak<SupervisorCore>,
bus: Bus,
slots: DashMap<Arc<str>, Arc<Mutex<SlotState>>>,
running: DashMap<TaskId, Arc<str>>,
watchers: DashMap<TaskId, OutcomeTx>,
tx: mpsc::Sender<Submission>,
rx: RwLock<Option<mpsc::Receiver<Submission>>>,
shutting_down: std::sync::atomic::AtomicBool,
}
impl Controller {
pub fn new(config: ControllerConfig, supervisor: &Arc<SupervisorCore>, bus: Bus) -> Arc<Self> {
let (tx, rx) = mpsc::channel(config.queue_capacity.max(1));
Arc::new(Self {
config,
supervisor: Arc::downgrade(supervisor),
bus,
slots: DashMap::new(),
running: DashMap::new(),
watchers: DashMap::new(),
tx,
rx: RwLock::new(Some(rx)),
shutting_down: std::sync::atomic::AtomicBool::new(false),
})
}
fn finalize_rejected(&self, id: TaskId, reason: &str) {
if let Some((_, tx)) = self.watchers.remove(&id) {
let _ = tx.send(TaskOutcome::Rejected {
reason: Arc::from(reason),
});
}
}
fn is_shutting_down(&self) -> bool {
self.shutting_down
.load(std::sync::atomic::Ordering::Acquire)
}
pub fn handle(&self) -> ControllerHandle {
ControllerHandle {
tx: self.tx.clone(),
}
}
pub fn run(self: Arc<Self>, token: CancellationToken) {
let bus = self.bus.clone();
tokio::spawn(async move {
if let Err(e) = self.run_inner(token).await {
bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task("controller")
.with_reason(format!("controller_loop_exited: {e}")),
);
}
});
}
async fn run_inner(&self, token: CancellationToken) -> Result<(), ControllerError> {
let mut rx = self
.rx
.write()
.await
.take()
.ok_or(ControllerError::AlreadyStarted)?;
let mut bus_rx = self.bus.subscribe();
let mut recovery_at = None;
let recovery_timer = tokio::time::sleep(recovery::RECOVERY_DELAY);
tokio::pin!(recovery_timer);
loop {
tokio::select! {
_ = token.cancelled() => break,
_ = &mut recovery_timer, if recovery_at.is_some() => {
recovery_at = None;
if !self.is_shutting_down()
&& let Some(next_recovery) = self
.guarded("recover_stale_slots", self.recover_stale_slots())
.await
.flatten()
&& let Some(deadline) = schedule_recovery(
&mut recovery_at,
next_recovery,
)
{
recovery_timer.as_mut().reset(deadline);
}
}
Some(sub) = rx.recv() => {
let _ = self.guarded("handle_submission", self.handle_submission(sub)).await;
}
result = bus_rx.recv() => {
match result {
Ok(event) => {
let shutdown_requested = event.kind == EventKind::ShutdownRequested;
let _ = self.guarded("handle_event", self.handle_event(event)).await;
if shutdown_requested {
recovery_at = None;
}
}
Err(broadcast::error::RecvError::Lagged(_)) => {
if !self.is_shutting_down() {
if let Some(deadline) = schedule_recovery(
&mut recovery_at,
Instant::now() + recovery::RECOVERY_DELAY,
) {
recovery_timer.as_mut().reset(deadline);
}
}
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
}
}
self.finalize_pending_on_shutdown(&mut rx);
Ok(())
}
async fn guarded<T>(&self, who: &'static str, fut: impl Future<Output = T>) -> Option<T> {
match crate::core::panic_guard::guarded(fut).await {
Ok(output) => Some(output),
Err(msg) => {
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task("controller")
.with_reason(format!("{who}_panicked: {msg}")),
);
None
}
}
}
async fn handle_submission(&self, sub: Submission) {
let Some(sup) = self.supervisor.upgrade() else {
return;
};
let Submission { id, spec, done } = sub;
if let Some(tx) = done {
self.watchers.insert(id, tx);
}
if self.is_shutting_down() {
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(spec.slot_name().to_owned())
.with_id(id)
.with_reason(crate::reasons::CONTROLLER_SHUTTING_DOWN),
);
self.finalize_rejected(id, crate::reasons::CONTROLLER_SHUTTING_DOWN);
return;
}
let slot_name: Arc<str> = Arc::from(spec.slot_name());
let admission = spec.admission;
let task_spec = spec.task_spec;
let slot_arc = self.get_or_create_slot(&slot_name);
let mut slot = slot_arc.lock().await;
match (&slot.status, admission) {
(SlotStatus::Idle, _) => {
match self.start_in_slot(&sup, &mut slot, &slot_name, id, task_spec) {
Ok(()) => {
let reason: &'static str = match admission {
AdmissionPolicy::Queue => "admission=Queue status=admitting",
AdmissionPolicy::Replace => "admission=Replace status=admitting",
AdmissionPolicy::DropIfRunning => {
"admission=DropIfRunning status=admitting"
}
};
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(reason),
);
}
Err(e) => {
let reason = format!("add_failed: {e}");
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(reason.clone()),
);
self.finalize_rejected(id, &reason);
self.gc_if_idle(&slot_name, slot);
}
}
}
(SlotStatus::Running { .. }, AdmissionPolicy::Replace) => {
if let Some(rid) = slot.running_id
&& let Err(e) = sup.remove(rid)
{
let reason = format!("remove_failed: {e}");
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(reason.clone()),
);
self.finalize_rejected(id, &reason);
return;
}
self.replace_head_or_push(&mut slot, &slot_name, id, task_spec);
slot.status = SlotStatus::Terminating {
cancelled_at: Instant::now(),
};
self.bus.publish(
Event::new(EventKind::ControllerSlotTransition)
.with_task(Arc::clone(&slot_name))
.with_reason("running→terminating (replace)"),
);
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(format!("admission=Replace depth={}", slot.queue.len())),
);
}
(SlotStatus::Admitting { .. }, AdmissionPolicy::Replace) => {
self.replace_head_or_push(&mut slot, &slot_name, id, task_spec);
slot.status = SlotStatus::Terminating {
cancelled_at: Instant::now(),
};
self.bus.publish(
Event::new(EventKind::ControllerSlotTransition)
.with_task(Arc::clone(&slot_name))
.with_reason("admitting→terminating (replace)"),
);
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(format!(
"admission=Replace status=admitting depth={}",
slot.queue.len()
)),
);
}
(SlotStatus::Terminating { .. }, AdmissionPolicy::Replace) => {
self.replace_head_or_push(&mut slot, &slot_name, id, task_spec);
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(format!(
"admission=Replace status=terminating depth={}",
slot.queue.len()
)),
);
}
(
SlotStatus::Admitting { .. }
| SlotStatus::Running { .. }
| SlotStatus::Terminating { .. },
AdmissionPolicy::Queue,
) => {
if self.reject_if_full(&slot_name, id, slot.queue.len()) {
return;
}
slot.queue.push_back((id, task_spec));
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(format!("admission=Queue depth={}", slot.queue.len())),
);
}
(
SlotStatus::Admitting { .. }
| SlotStatus::Running { .. }
| SlotStatus::Terminating { .. },
AdmissionPolicy::DropIfRunning,
) => {
let reason = format!("dropped: slot busy ({})", slot.status.label());
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(reason.clone()),
);
self.finalize_rejected(id, &reason);
}
}
}
async fn handle_event(&self, event: Arc<Event>) {
match event.kind {
EventKind::TaskAdded => self.on_task_added(&event).await,
EventKind::TaskRemoved => self.on_task_finished(&event).await,
EventKind::TaskAddFailed => self.on_task_add_failed(&event).await,
EventKind::TaskRemoveRequested => self.on_remove_requested(&event).await,
EventKind::ShutdownRequested => {
self.shutting_down
.store(true, std::sync::atomic::Ordering::Release);
}
_ => {}
}
}
async fn on_remove_requested(&self, event: &Event) {
let Some(id) = event.id else {
return;
};
let slot_keys: Vec<Arc<str>> = self
.slots
.iter()
.map(|entry| Arc::clone(entry.key()))
.collect();
for slot_name in slot_keys {
let Some(slot_arc) = self.slots.get(&*slot_name).map(|e| e.clone()) else {
continue;
};
let mut slot = slot_arc.lock().await;
let Some(pos) = slot.queue.iter().position(|(qid, _)| *qid == id) else {
continue;
};
slot.queue.remove(pos);
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(&slot_name))
.with_id(id)
.with_reason(crate::reasons::REMOVED_FROM_QUEUE),
);
self.finalize_rejected(id, crate::reasons::REMOVED_FROM_QUEUE);
self.gc_if_idle(&slot_name, slot);
return;
}
}
async fn on_task_added(&self, event: &Event) {
let Some(id) = event.id else {
return;
};
let Some(sup) = self.supervisor.upgrade() else {
return;
};
let Some(slot_name) = self.running.get(&id).map(|e| e.clone()) else {
return;
};
let Some(slot_arc) = self.slots.get(&*slot_name).map(|e| e.clone()) else {
return;
};
let mut slot = slot_arc.lock().await;
if slot.running_id != Some(id) {
return;
}
match slot.status {
SlotStatus::Admitting { .. } => {
slot.status = SlotStatus::Running {
started_at: Instant::now(),
};
self.bus.publish(
Event::new(EventKind::ControllerSlotTransition)
.with_task(slot_name)
.with_reason("admitting→running"),
);
}
SlotStatus::Terminating { .. } => {
if let Some(rid) = slot.running_id
&& let Err(e) = sup.remove(rid)
{
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(slot_name)
.with_reason(format!("remove_failed: {e}")),
);
}
}
_ => {}
}
}
async fn on_task_add_failed(&self, event: &Event) {
self.free_and_advance(event).await;
}
async fn on_task_finished(&self, event: &Event) {
self.free_and_advance(event).await;
}
async fn free_and_advance(&self, event: &Event) {
let Some(id) = event.id else {
return;
};
let Some(sup) = self.supervisor.upgrade() else {
return;
};
let Some((_, slot_name)) = self.running.remove(&id) else {
return;
};
let Some(slot_arc) = self.slots.get(&*slot_name).map(|e| e.clone()) else {
return;
};
let mut slot = slot_arc.lock().await;
if slot.running_id != Some(id) {
return;
}
slot.running_id = None;
slot.status = SlotStatus::Idle;
if !self.is_shutting_down() {
self.start_next_from_queue(&sup, &mut slot, &slot_name);
}
self.gc_if_idle(&slot_name, slot);
}
fn start_in_slot(
&self,
sup: &Arc<SupervisorCore>,
slot: &mut SlotState,
slot_name: &Arc<str>,
id: TaskId,
task_spec: TaskSpec,
) -> Result<(), RuntimeError> {
let done = self.watchers.remove(&id).map(|(_, tx)| tx);
match sup.add_task_with_id_watched(id, task_spec, done) {
Ok(_) => {
slot.status = SlotStatus::Admitting {
since: Instant::now(),
};
slot.running_id = Some(id);
self.running.insert(id, Arc::clone(slot_name));
Ok(())
}
Err((e, done)) => {
if let Some(tx) = done {
self.watchers.insert(id, tx);
}
Err(e)
}
}
}
fn start_next_from_queue(
&self,
sup: &Arc<SupervisorCore>,
slot: &mut SlotState,
slot_name: &Arc<str>,
) {
while let Some((next_id, next_spec)) = slot.queue.pop_front() {
match self.start_in_slot(sup, slot, slot_name, next_id, next_spec) {
Ok(()) => {
self.bus.publish(
Event::new(EventKind::ControllerSubmitted)
.with_task(Arc::clone(slot_name))
.with_id(next_id)
.with_reason(format!("started_from_queue depth={}", slot.queue.len())),
);
return;
}
Err(e) => {
let reason = format!("queue_start_failed: {e}");
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(slot_name))
.with_id(next_id)
.with_reason(reason.clone()),
);
self.finalize_rejected(next_id, &reason);
}
}
}
}
#[inline]
fn get_or_create_slot(&self, slot_name: &str) -> Arc<Mutex<SlotState>> {
if let Some(slot) = self.slots.get(slot_name) {
return slot.clone();
}
self.slots
.entry(Arc::from(slot_name))
.or_insert_with(|| Arc::new(Mutex::new(SlotState::new())))
.clone()
}
#[inline]
fn gc_if_idle(&self, slot_name: &Arc<str>, slot: tokio::sync::MutexGuard<'_, SlotState>) {
let collect = matches!(slot.status, SlotStatus::Idle) && slot.queue.is_empty();
drop(slot);
if collect {
self.slots.remove(&**slot_name);
}
}
#[inline]
fn reject_if_full(&self, slot_name: &str, id: TaskId, slot_len: usize) -> bool {
if slot_len >= self.config.max_slot_queue {
let reason = format!(
"{}: {}/{}",
crate::reasons::QUEUE_FULL,
slot_len,
self.config.max_slot_queue
);
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(slot_name)
.with_id(id)
.with_reason(reason.clone()),
);
self.finalize_rejected(id, &reason);
true
} else {
false
}
}
fn replace_head_or_push(
&self,
slot: &mut SlotState,
slot_name: &Arc<str>,
id: TaskId,
task_spec: crate::TaskSpec,
) {
if let Some(head) = slot.queue.front_mut() {
let (displaced_id, _) = std::mem::replace(head, (id, task_spec));
self.bus.publish(
Event::new(EventKind::ControllerRejected)
.with_task(Arc::clone(slot_name))
.with_id(displaced_id)
.with_reason(crate::reasons::SUPERSEDED_BY_REPLACE),
);
self.finalize_rejected(displaced_id, crate::reasons::SUPERSEDED_BY_REPLACE);
} else {
slot.queue.push_front((id, task_spec));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Supervisor;
use crate::TaskContext;
use crate::{BackoffPolicy, RestartPolicy, TaskFn, TaskRef, TaskSpec};
use std::time::Duration;
fn make_spec(name: &str) -> TaskSpec {
let task: TaskRef = TaskFn::arc(name, |_ctx: TaskContext| async { Ok(()) });
TaskSpec::new(task, RestartPolicy::Never, BackoffPolicy::default(), None)
}
fn slot_arc_name() -> Arc<str> {
Arc::from("s")
}
#[test]
fn replace_head_or_push_into_empty_queue() {
let ctrl = make_controller(ControllerConfig::default(), Bus::new(64));
let mut slot = SlotState::new();
ctrl.replace_head_or_push(
&mut slot,
&slot_arc_name(),
TaskId::next(),
make_spec("first"),
);
assert_eq!(slot.queue.len(), 1);
assert_eq!(slot.queue.front().unwrap().1.name(), "first");
}
#[test]
fn replace_head_or_push_replaces_existing_head_and_rejects_displaced() {
let ctrl = make_controller(ControllerConfig::default(), Bus::new(64));
let mut rx = ctrl.bus.subscribe();
let mut slot = SlotState::new();
let displaced = TaskId::next();
slot.queue.push_back((displaced, make_spec("old-head")));
slot.queue.push_back((TaskId::next(), make_spec("tail")));
ctrl.replace_head_or_push(
&mut slot,
&slot_arc_name(),
TaskId::next(),
make_spec("new-head"),
);
assert_eq!(slot.queue.len(), 2, "queue depth should not grow");
assert_eq!(slot.queue.front().unwrap().1.name(), "new-head");
assert_eq!(slot.queue.back().unwrap().1.name(), "tail");
let ev = rx.try_recv().expect("displaced head must be rejected");
assert_eq!(ev.kind, EventKind::ControllerRejected);
assert_eq!(ev.id, Some(displaced));
assert_eq!(
ev.reason.as_deref(),
Some(crate::reasons::SUPERSEDED_BY_REPLACE)
);
}
#[test]
fn replace_head_multiple_times_keeps_depth_1() {
let ctrl = make_controller(ControllerConfig::default(), Bus::new(64));
let mut slot = SlotState::new();
let name = slot_arc_name();
ctrl.replace_head_or_push(&mut slot, &name, TaskId::next(), make_spec("v1"));
ctrl.replace_head_or_push(&mut slot, &name, TaskId::next(), make_spec("v2"));
ctrl.replace_head_or_push(&mut slot, &name, TaskId::next(), make_spec("v3"));
assert_eq!(slot.queue.len(), 1);
assert_eq!(slot.queue.front().unwrap().1.name(), "v3");
}
#[test]
fn reject_if_full_returns_false_below_capacity() {
let bus = Bus::new(64);
let config = ControllerConfig {
queue_capacity: 16,
max_slot_queue: 3,
};
let ctrl = make_controller(config, bus);
assert!(!ctrl.reject_if_full("slot", TaskId::next(), 0));
assert!(!ctrl.reject_if_full("slot", TaskId::next(), 2));
}
#[test]
fn reject_if_full_returns_true_at_capacity() {
let bus = Bus::new(64);
let config = ControllerConfig {
queue_capacity: 16,
max_slot_queue: 3,
};
let ctrl = make_controller(config, bus);
assert!(ctrl.reject_if_full("slot", TaskId::next(), 3));
assert!(ctrl.reject_if_full("slot", TaskId::next(), 10));
}
#[test]
fn get_or_create_slot_creates_idle_slot() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let slot_arc = ctrl.get_or_create_slot("my-slot");
let slot = slot_arc.blocking_lock();
assert_eq!(slot.status, SlotStatus::Idle);
assert!(slot.queue.is_empty());
}
#[test]
fn get_or_create_slot_returns_same_arc() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let s1 = ctrl.get_or_create_slot("x");
let s2 = ctrl.get_or_create_slot("x");
assert!(Arc::ptr_eq(&s1, &s2), "same slot name must return same Arc");
}
#[test]
fn get_or_create_slot_different_names_different_arcs() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let s1 = ctrl.get_or_create_slot("a");
let s2 = ctrl.get_or_create_slot("b");
assert!(!Arc::ptr_eq(&s1, &s2));
}
#[tokio::test]
async fn handle_event_ignores_non_task_removed() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let slot_arc = ctrl.get_or_create_slot("t");
{
let mut slot = slot_arc.lock().await;
slot.status = SlotStatus::Running {
started_at: Instant::now(),
};
}
let event = Arc::new(Event::new(EventKind::TaskFailed).with_task("t"));
ctrl.handle_event(event).await;
let slot = slot_arc.lock().await;
assert!(
matches!(slot.status, SlotStatus::Running { .. }),
"non-TaskRemoved events should not affect slot state"
);
}
#[tokio::test]
async fn shutdown_finalizes_buffered_submission_as_rejected() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let task: TaskRef = TaskFn::arc("buffered", |_ctx: TaskContext| async { Ok(()) });
let (_id, waiter) = ctrl
.handle()
.submit_and_watch(ControllerSpec::queue(TaskSpec::once(task)).with_slot("s"))
.await
.expect("submission accepted into channel");
let mut rx = ctrl.rx.write().await.take().expect("rx present");
ctrl.finalize_pending_on_shutdown(&mut rx);
drop(rx);
let outcome = tokio::time::timeout(Duration::from_secs(1), waiter)
.await
.expect("waiter must resolve, not hang")
.expect("waiter must resolve to an outcome, not a dropped sender");
assert!(
matches!(outcome, TaskOutcome::Rejected { .. }),
"a buffered submission on shutdown must resolve Rejected, got {outcome:?}"
);
}
#[tokio::test]
async fn submit_after_shutdown_finalize_is_rejected_not_leaked() {
let bus = Bus::new(64);
let ctrl = make_controller(ControllerConfig::default(), bus);
let mut rx = ctrl.rx.write().await.take().expect("rx present");
ctrl.finalize_pending_on_shutdown(&mut rx);
let task: TaskRef = TaskFn::arc("late", |_ctx: TaskContext| async { Ok(()) });
let result = ctrl
.handle()
.submit_and_watch(ControllerSpec::queue(TaskSpec::once(task)).with_slot("s"))
.await;
assert!(
result.is_err(),
"a submission after shutdown finalization must be rejected, not handed a doomed waiter"
);
drop(rx);
}
fn make_controller(config: ControllerConfig, bus: Bus) -> Controller {
let (tx, rx) = mpsc::channel(config.queue_capacity.max(1));
Controller {
config,
supervisor: Weak::new(),
bus,
slots: DashMap::new(),
running: DashMap::new(),
watchers: DashMap::new(),
tx,
rx: RwLock::new(Some(rx)),
shutting_down: std::sync::atomic::AtomicBool::new(false),
}
}
#[tokio::test]
async fn guarded_converts_panic_to_diagnostic_and_survives() {
let ctrl = make_controller(ControllerConfig::default(), Bus::new(64));
let mut rx = ctrl.bus.subscribe();
let _ = ctrl.guarded("unit", async { panic!("boom {}", 1) }).await;
let ev = rx
.try_recv()
.expect("a panicking work-unit must publish a diagnostic");
assert_eq!(ev.kind, EventKind::ControllerRejected);
assert!(
ev.reason.as_deref().unwrap_or_default().contains("boom 1"),
"diagnostic must carry the panic message, got {:?}",
ev.reason
);
}
#[tokio::test]
async fn zero_queue_capacity_is_clamped_not_panicking() {
let sup = Supervisor::builder(crate::SupervisorConfig::default())
.with_controller(ControllerConfig {
queue_capacity: 0,
max_slot_queue: 1,
})
.build();
let handle = sup.serve();
let task: TaskRef = TaskFn::arc("clamped", |_ctx: TaskContext| async { Ok(()) });
handle
.submit(ControllerSpec::queue(TaskSpec::once(task)))
.await
.expect("submission must work with capacity clamped to 1");
let _ = handle.shutdown().await;
}
fn long_ago() -> Instant {
Instant::now()
.checked_sub(Duration::from_secs(60))
.expect("test host uptime must exceed one minute")
}
fn insert_admitting_slot(
ctrl: &Controller,
name: &str,
id: TaskId,
since: Instant,
) -> Arc<Mutex<SlotState>> {
let slot_name: Arc<str> = Arc::from(name);
let slot = Arc::new(Mutex::new(SlotState {
status: SlotStatus::Admitting { since },
running_id: Some(id),
queue: std::collections::VecDeque::new(),
}));
ctrl.slots.insert(Arc::clone(&slot_name), Arc::clone(&slot));
ctrl.running.insert(id, slot_name);
slot
}
async fn wait_for_slot_status(
slot: &Arc<Mutex<SlotState>>,
expected: impl Fn(SlotStatus) -> bool,
) -> bool {
for _ in 0..100 {
if expected(slot.lock().await.status) {
return true;
}
tokio::task::yield_now().await;
}
false
}
async fn start_controller_loop(
ctrl: &Arc<Controller>,
token: &CancellationToken,
) -> tokio::task::JoinHandle<Result<(), ControllerError>> {
let runner_ctrl = Arc::clone(ctrl);
let runner_token = token.clone();
let runner = tokio::spawn(async move { runner_ctrl.run_inner(runner_token).await });
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if ctrl.rx.read().await.is_none() {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("controller loop must take its receiver and subscribe to the bus");
runner
}
async fn stop_controller_loop(
token: CancellationToken,
runner: tokio::task::JoinHandle<Result<(), ControllerError>>,
) {
token.cancel();
tokio::time::timeout(Duration::from_secs(1), runner)
.await
.expect("controller loop must stop after cancellation")
.expect("controller loop task must not panic")
.expect("controller loop must exit cleanly");
}
async fn induce_lag_with_task_added(
bus: &Bus,
tail_name: &str,
tail_id: TaskId,
tail_slot: &Arc<Mutex<SlotState>>,
) -> bool {
bus.publish(Event::new(EventKind::TaskStarting).with_task("lag-seed"));
bus.publish(
Event::new(EventKind::TaskAdded)
.with_task(tail_name.to_owned())
.with_id(tail_id),
);
wait_for_slot_status(tail_slot, |status| {
matches!(status, SlotStatus::Running { .. })
})
.await
}
#[tokio::test(flavor = "current_thread")]
async fn lag_recovery_keeps_retained_task_added_event() {
let bus = Bus::new(1);
let sup = Supervisor::new(crate::SupervisorConfig::default(), vec![]);
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), bus.clone());
let id = TaskId::next();
let slot = insert_admitting_slot(&ctrl, "s", id, Instant::now());
let token = CancellationToken::new();
let runner = start_controller_loop(&ctrl, &token).await;
let reached_running = induce_lag_with_task_added(&bus, "s", id, &slot).await;
stop_controller_loop(token, runner).await;
assert!(
reached_running,
"lag recovery must preserve and process the retained TaskAdded event"
);
}
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn repeated_lag_keeps_first_deadline_and_rearms_young_slots() {
let bus = Bus::new(1);
let sup = Supervisor::new(crate::SupervisorConfig::default(), vec![]);
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), bus.clone());
let first = insert_admitting_slot(&ctrl, "first", TaskId::next(), Instant::now());
let first_tail_id = TaskId::next();
let first_tail = insert_admitting_slot(&ctrl, "first-tail", first_tail_id, Instant::now());
let token = CancellationToken::new();
let runner = start_controller_loop(&ctrl, &token).await;
assert!(
induce_lag_with_task_added(&bus, "first-tail", first_tail_id, &first_tail).await,
"first retained event must be processed"
);
tokio::time::advance(Duration::from_secs(4)).await;
assert!(
matches!(first.lock().await.status, SlotStatus::Admitting { .. }),
"recovery must not run before the safety delay"
);
let stale_since = Instant::now()
.checked_sub(recovery::RECOVERY_DELAY)
.expect("paused clock must support a five-second lookback");
let second = insert_admitting_slot(&ctrl, "second", TaskId::next(), stale_since);
let late = insert_admitting_slot(&ctrl, "late", TaskId::next(), Instant::now());
let second_tail_id = TaskId::next();
let second_tail =
insert_admitting_slot(&ctrl, "second-tail", second_tail_id, Instant::now());
assert!(
induce_lag_with_task_added(&bus, "second-tail", second_tail_id, &second_tail).await,
"second retained event must be processed"
);
assert!(
matches!(second.lock().await.status, SlotStatus::Admitting { .. }),
"repeated lag must not run stale-slot recovery immediately"
);
tokio::time::advance(Duration::from_secs(1)).await;
assert!(
wait_for_slot_status(&first, |status| matches!(status, SlotStatus::Idle)).await
&& wait_for_slot_status(&second, |status| matches!(status, SlotStatus::Idle)).await,
"repeated lag must keep the first recovery deadline"
);
assert!(
matches!(late.lock().await.status, SlotStatus::Admitting { .. }),
"a newer slot must not be recovered before its safety delay"
);
tokio::time::advance(Duration::from_secs(4)).await;
assert!(
wait_for_slot_status(&late, |status| matches!(status, SlotStatus::Idle)).await,
"a newer slot must schedule its own delayed recovery"
);
stop_controller_loop(token, runner).await;
}
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn shutdown_request_cancels_pending_lag_recovery() {
let bus = Bus::new(1);
let sup = Supervisor::new(crate::SupervisorConfig::default(), vec![]);
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), bus.clone());
let target = insert_admitting_slot(&ctrl, "target", TaskId::next(), Instant::now());
let tail_id = TaskId::next();
let tail = insert_admitting_slot(&ctrl, "tail", tail_id, Instant::now());
let token = CancellationToken::new();
let runner = start_controller_loop(&ctrl, &token).await;
assert!(
induce_lag_with_task_added(&bus, "tail", tail_id, &tail).await,
"retained event must confirm that lag was handled"
);
bus.publish(Event::new(EventKind::ShutdownRequested));
let mut shutdown_observed = false;
for _ in 0..100 {
if ctrl.is_shutting_down() {
shutdown_observed = true;
break;
}
tokio::task::yield_now().await;
}
tokio::time::advance(recovery::RECOVERY_DELAY + Duration::from_secs(1)).await;
for _ in 0..10 {
tokio::task::yield_now().await;
}
let target_unchanged = matches!(target.lock().await.status, SlotStatus::Admitting { .. });
stop_controller_loop(token, runner).await;
assert!(
shutdown_observed,
"controller must observe shutdown request"
);
assert!(
target_unchanged,
"pending recovery must not run after shutdown is requested"
);
}
async fn sup_with_live_task() -> (Arc<Supervisor>, crate::core::SupervisorHandle, TaskId) {
let sup = Supervisor::new(crate::SupervisorConfig::default(), vec![]);
let handle = sup.serve();
let task: TaskRef = TaskFn::arc("occupant", |ctx: TaskContext| async move {
ctx.cancelled().await;
Ok(())
});
let id = handle
.add_and_wait(TaskSpec::restartable(task), Duration::from_secs(1))
.await
.expect("task should register");
(sup, handle, id)
}
#[tokio::test]
async fn recover_promotes_admitting_slot_with_alive_task() {
let (sup, handle, id) = sup_with_live_task().await;
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), Bus::new(64));
ctrl.slots.insert(
Arc::from("s"),
Arc::new(Mutex::new(SlotState {
status: SlotStatus::Admitting { since: long_ago() },
running_id: Some(id),
queue: std::collections::VecDeque::new(),
})),
);
ctrl.running.insert(id, Arc::from("s"));
let _ = ctrl.recover_stale_slots().await;
let slot_arc = ctrl.slots.get("s").map(|e| e.clone()).expect("slot exists");
let slot = slot_arc.lock().await;
assert!(
matches!(slot.status, SlotStatus::Running { .. }),
"an Admitting slot whose task is alive must be promoted to Running, got {:?}",
slot.status
);
drop(slot);
let _ = handle.shutdown().await;
}
#[tokio::test]
async fn recover_reissues_removal_for_terminating_slot_with_alive_task() {
let (sup, handle, id) = sup_with_live_task().await;
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), Bus::new(64));
ctrl.slots.insert(
Arc::from("s"),
Arc::new(Mutex::new(SlotState {
status: SlotStatus::Terminating {
cancelled_at: long_ago(),
},
running_id: Some(id),
queue: std::collections::VecDeque::new(),
})),
);
ctrl.running.insert(id, Arc::from("s"));
let _ = ctrl.recover_stale_slots().await;
let removed = poll_until(Duration::from_secs(2), || async {
!sup.core().contains_id(id).await
})
.await;
assert!(
removed,
"recovery must re-issue the deferred removal for a Terminating slot"
);
let _ = handle.shutdown().await;
}
#[tokio::test]
async fn no_queue_advancement_after_shutdown_requested() {
let (sup, handle, id) = sup_with_live_task().await;
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), Bus::new(64));
let queued: TaskRef = TaskFn::arc("queued", |ctx: TaskContext| async move {
ctx.cancelled().await;
Ok(())
});
let mut queue = std::collections::VecDeque::new();
queue.push_back((TaskId::next(), TaskSpec::restartable(queued)));
ctrl.slots.insert(
Arc::from("s"),
Arc::new(Mutex::new(SlotState {
status: SlotStatus::Running {
started_at: Instant::now(),
},
running_id: Some(id),
queue,
})),
);
ctrl.running.insert(id, Arc::from("s"));
ctrl.handle_event(Arc::new(Event::new(EventKind::ShutdownRequested)))
.await;
ctrl.handle_event(Arc::new(
Event::new(EventKind::TaskRemoved)
.with_task("occupant")
.with_id(id),
))
.await;
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(
sup.core().id_for_label("queued").await.is_none(),
"controller must not start queued tasks once shutdown has been requested"
);
let _ = handle.shutdown().await;
}
#[tokio::test]
async fn replace_supersedes_in_same_slot() {
let sup = Supervisor::builder(crate::SupervisorConfig::default())
.with_controller(ControllerConfig::default())
.build();
let handle = sup.serve();
let mk = |name: &'static str| -> ControllerSpec {
let task: TaskRef = TaskFn::arc(name, |ctx: TaskContext| async move {
ctx.cancelled().await;
Ok(())
});
ControllerSpec::replace(TaskSpec::restartable(task)).with_slot("s")
};
handle.submit(mk("run-1")).await.unwrap();
handle.submit(mk("run-2")).await.unwrap();
let superseded = poll_until(std::time::Duration::from_secs(3), || async {
let alive = handle.snapshot().await;
alive.iter().any(|n| &**n == "run-2") && alive.iter().all(|n| &**n != "run-1")
})
.await;
assert!(
superseded,
"Replace must supersede run-1 with run-2 in the shared slot, not run both"
);
let _ = handle.shutdown().await;
}
#[tokio::test]
async fn snapshot_reports_status_running_and_queue_depth() {
use crate::controller::SlotStatusKind;
let (sup, handle, id) = sup_with_live_task().await;
let ctrl = Controller::new(ControllerConfig::default(), sup.core(), Bus::new(64));
let queued: TaskRef = TaskFn::arc("queued", |ctx: TaskContext| async move {
ctx.cancelled().await;
Ok(())
});
let mut queue = std::collections::VecDeque::new();
queue.push_back((TaskId::next(), TaskSpec::restartable(queued)));
ctrl.slots.insert(
Arc::from("s"),
Arc::new(Mutex::new(SlotState {
status: SlotStatus::Running {
started_at: Instant::now(),
},
running_id: Some(id),
queue,
})),
);
let snap = ctrl.snapshot().await;
assert_eq!(snap.len(), 1, "one slot tracked");
assert_eq!(snap.running_count(), 1);
assert_eq!(snap.total_queued(), 1);
let view = snap.slot("s").expect("slot 's' must be present");
assert_eq!(view.status, SlotStatusKind::Running);
assert_eq!(view.queue_depth, 1);
assert_eq!(view.running, Some(id));
let _ = handle.shutdown().await;
}
async fn poll_until<F, Fut>(within: std::time::Duration, mut cond: F) -> bool
where
F: FnMut() -> Fut,
Fut: Future<Output = bool>,
{
let deadline = tokio::time::Instant::now() + within;
loop {
if cond().await {
return true;
}
if tokio::time::Instant::now() >= deadline {
return false;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
}