use std::sync::{Arc, Mutex as StdMutex, MutexGuard as StdMutexGuard, OnceLock, Weak};
use tokio::sync::{Mutex, RwLock, mpsc};
use tokio_util::sync::CancellationToken;
use crate::{
core::{OutcomeTx, SupervisorCore, TaskOutcome, deferred_drop::DropDomain},
events::{Bus, Event, EventKind, RejectionKind},
identity::TaskId,
};
use super::config::ControllerConfig;
mod state;
use state::{CapacityPending, ControllerState, SlotState};
mod command;
use command::{
AdmissionResult, CompletionResult, ControllerCommand, IdentityOperation, IdentityReply,
RemovalResult, Submission,
};
mod handle;
pub(crate) use handle::ControllerHandle;
mod admission;
mod identity;
mod lifecycle;
use lifecycle::ControllerTask;
mod operations;
mod queue;
#[cfg(test)]
use operations::OperationSet;
use operations::TrackedOperations;
mod snapshot;
pub(crate) struct Controller {
config: ControllerConfig,
supervisor: Weak<SupervisorCore>,
drop_domain: DropDomain,
bus: Bus,
shutdown_token: CancellationToken,
state: StdMutex<ControllerState>,
tx: mpsc::Sender<ControllerCommand>,
rx: RwLock<Option<mpsc::Receiver<ControllerCommand>>>,
shutting_down: std::sync::atomic::AtomicBool,
task: OnceLock<ControllerTask>,
}
impl Controller {
fn state(&self) -> StdMutexGuard<'_, ControllerState> {
self.state.lock().unwrap_or_else(|error| error.into_inner())
}
fn slot(&self, name: &str) -> Option<Arc<Mutex<SlotState>>> {
self.state().slots.get(name).cloned()
}
pub fn new(config: ControllerConfig, supervisor: &Arc<SupervisorCore>, bus: Bus) -> Arc<Self> {
let (tx, rx) = mpsc::channel(config.queue_capacity().get());
let shutdown_token = supervisor.shutdown_started_token();
Arc::new(Self {
config,
supervisor: Arc::downgrade(supervisor),
drop_domain: supervisor.drop_domain().clone(),
bus,
shutdown_token,
state: StdMutex::new(ControllerState::default()),
tx,
rx: RwLock::new(Some(rx)),
shutting_down: std::sync::atomic::AtomicBool::new(false),
task: OnceLock::new(),
})
}
fn finalize_rejected(
&self,
id: TaskId,
kind: RejectionKind,
reason: &str,
) -> Option<TaskOutcome> {
let tx = self.state().watchers.remove(&id)?;
Self::send_rejected(Some(tx), kind, reason)
}
fn send_rejected(
done: Option<OutcomeTx>,
kind: RejectionKind,
reason: &str,
) -> Option<TaskOutcome> {
done?
.send(TaskOutcome::Rejected {
kind,
reason: Arc::from(reason),
})
.err()
}
fn mark_shutting_down(&self) {
self.shutting_down
.store(true, std::sync::atomic::Ordering::Release);
}
fn is_shutting_down(&self) -> bool {
self.shutdown_token.is_cancelled()
|| self
.shutting_down
.load(std::sync::atomic::Ordering::Acquire)
}
fn finalize_remaining_watchers(&self) {
let pending: Vec<TaskId> = self.state().watchers.keys().copied().collect();
for id in pending {
self.bus.publish_lazy(|| {
Event::new(EventKind::ControllerRejected)
.with_id(id)
.with_rejection_kind(RejectionKind::ControllerShuttingDown)
.with_reason(crate::reasons::CONTROLLER_SHUTTING_DOWN)
});
drop(self.finalize_rejected(
id,
RejectionKind::ControllerShuttingDown,
crate::reasons::CONTROLLER_SHUTTING_DOWN,
));
}
}
pub fn handle(&self) -> ControllerHandle {
ControllerHandle::new(self.tx.clone(), self.bus.clone(), self.drop_domain.clone())
}
}
#[cfg(test)]
mod tests;