use std::collections::HashMap;
use super::observation::{ExitObservation, ObservationError};
use super::{
CommandTags, ComponentId, ExitReason, FailureReason, HostCommand, HostSignal, Instant,
LifecycleConfig, LifecycleState, LiveChild, MonitorContext, MonitorMessage, Receiver,
RegistryError, Scheduler, SharedStatus, TombstoneKind, VecDeque, process_failure,
publish_children, send_host_command, spawn_child,
};
pub(super) fn run_monitor(mut context: MonitorContext) -> Result<(), RegistryError> {
let mut restarts = VecDeque::new();
loop {
match context.receiver.recv() {
Ok(MonitorMessage::Signal(signal)) => {
handle_runtime_signal(&mut context, signal, &mut restarts)?;
if lifecycle(&context.status)? == LifecycleState::Failed {
return Ok(());
}
}
Ok(MonitorMessage::Stop(reply)) => {
let result = drain_children(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
&context.receiver,
&mut context.live,
context.config,
)
.and_then(|()| {
stop_supervisor(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
context.config,
)
});
let _delivered = reply.send(result);
return Ok(());
}
Ok(MonitorMessage::Probe { child, reply_tx }) => {
let result = probe_live_child(&mut context, &child, &mut restarts);
let _delivered = reply_tx.send(result);
}
Ok(MonitorMessage::Send {
child,
message,
reply_tx,
}) => {
let result = find_child(&context.live, &child)
.ok_or_else(|| RegistryError::SupervisorProtocol {
id: context.id,
detail: format!("unknown child {child}"),
})
.and_then(|pid| {
send_host_command(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
HostCommand::Send {
pid,
message,
expects_reply: false,
},
)
});
let _delivered = reply_tx.send(result);
}
Err(_) => return Ok(()),
}
}
}
fn handle_runtime_signal(
context: &mut MonitorContext,
signal: HostSignal,
restarts: &mut VecDeque<Instant>,
) -> Result<(), RegistryError> {
let HostSignal::Exit { pid, kind } = signal else {
return match signal {
HostSignal::Protocol(detail) => Err(RegistryError::SupervisorProtocol {
id: context.id,
detail,
}),
HostSignal::Reply { .. } | HostSignal::UncorrelatedReply => {
Err(RegistryError::SupervisorProtocol {
id: context.id,
detail: "unsolicited child reply".to_owned(),
})
}
HostSignal::Live | HostSignal::Barrier => Err(RegistryError::SupervisorProtocol {
id: context.id,
detail: "unexpected supervisor control acknowledgement".to_owned(),
}),
HostSignal::Exit { .. } => Ok(()),
};
};
let Some(index) = context.live.iter().position(|child| child.pid == pid) else {
return Err(RegistryError::SupervisorProtocol {
id: context.id,
detail: format!("EXIT from unknown child pid {pid}"),
});
};
let exited = context.live.remove(index);
publish_children(&context.status, &context.live)?;
let now = Instant::now();
while restarts
.front()
.is_some_and(|at| now.duration_since(*at) > context.policy.window)
{
let _expired = restarts.pop_front();
}
if restarts.len() >= context.policy.max_restarts {
let reason = process_failure(&context.scheduler, exited.spec.name, pid, kind);
set_failed(context, reason)?;
drain_children(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
&context.receiver,
&mut context.live,
context.config,
)?;
publish_children(&context.status, &context.live)?;
stop_supervisor(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
context.config,
)?;
return Ok(());
}
restarts.push_back(now);
let replacement = spawn_child(
context.id,
&context.scheduler,
context.supervisor,
&exited.spec,
)?;
context.live.insert(
index,
LiveChild {
spec: exited.spec,
pid: replacement,
},
);
publish_children(&context.status, &context.live)
}
fn set_failed(context: &MonitorContext, reason: FailureReason) -> Result<(), RegistryError> {
let from = {
let mut status = context
.status
.lifecycle
.lock()
.map_err(|_| RegistryError::SynchronizationPoisoned)?;
let from = status.state;
status.state = LifecycleState::Failed;
status.failure = Some(reason);
from
};
context
.events
.publish_transition(context.id, Some(from), LifecycleState::Failed)
.map_err(|_| RegistryError::SynchronizationPoisoned)
}
fn lifecycle(status: &SharedStatus) -> Result<LifecycleState, RegistryError> {
status
.lifecycle
.lock()
.map(|state| state.state)
.map_err(|_| RegistryError::SynchronizationPoisoned)
}
fn probe_live_child(
context: &mut MonitorContext,
child: &str,
restarts: &mut VecDeque<Instant>,
) -> Result<i64, RegistryError> {
synchronize_mailbox(context, restarts)?;
let (pid, liveness_message) = context
.live
.iter()
.find(|live| live.spec.name == child)
.map(|live| (live.pid, live.spec.liveness_message))
.ok_or_else(|| RegistryError::SupervisorProtocol {
id: context.id,
detail: format!("unknown child {child}"),
})?;
send_host_command(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
HostCommand::Send {
pid,
message: liveness_message,
expects_reply: true,
},
)?;
match context
.receiver
.recv_timeout(context.config.operation_timeout)
{
Ok(MonitorMessage::Signal(HostSignal::Reply {
pid: reply_pid,
value,
})) if reply_pid == pid => Ok(value),
Ok(MonitorMessage::Signal(signal)) => {
handle_runtime_signal(context, signal, restarts)?;
Err(RegistryError::LivenessTimeout {
id: context.id,
child: child.to_owned(),
})
}
Ok(message) => {
refuse_during_probe(context.id, message);
Err(RegistryError::LivenessTimeout {
id: context.id,
child: child.to_owned(),
})
}
Err(_) => Err(RegistryError::LivenessTimeout {
id: context.id,
child: child.to_owned(),
}),
}
}
fn synchronize_mailbox(
context: &mut MonitorContext,
restarts: &mut VecDeque<Instant>,
) -> Result<(), RegistryError> {
send_host_command(
context.id,
&context.scheduler,
context.supervisor,
context.tags,
HostCommand::Barrier,
)?;
let deadline = Instant::now() + context.config.operation_timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
match context.receiver.recv_timeout(remaining) {
Ok(MonitorMessage::Signal(HostSignal::Barrier)) => return Ok(()),
Ok(MonitorMessage::Signal(signal)) => {
handle_runtime_signal(context, signal, restarts)?;
}
Ok(message) => refuse_during_probe(context.id, message),
Err(error) => {
return Err(RegistryError::SupervisorProtocol {
id: context.id,
detail: format!("supervisor mailbox barrier failed: {error}"),
});
}
}
}
}
fn refuse_during_probe(id: ComponentId, message: MonitorMessage) {
match message {
MonitorMessage::Stop(reply) => {
let _delivered = reply.send(Err(RegistryError::AlreadyStopping { id }));
}
MonitorMessage::Probe { child, reply_tx } => {
let _delivered = reply_tx.send(Err(RegistryError::SupervisorProtocol {
id,
detail: format!("probe for {child} arrived during another probe"),
}));
}
MonitorMessage::Send {
child, reply_tx, ..
} => {
let _delivered = reply_tx.send(Err(RegistryError::SupervisorProtocol {
id,
detail: format!("send for {child} arrived during a probe"),
}));
}
MonitorMessage::Signal(_) => {}
}
}
fn find_child(live: &[LiveChild], name: &str) -> Option<u64> {
live.iter()
.find(|child| child.spec.name == name)
.map(|child| child.pid)
}
pub(super) fn drain_children(
id: ComponentId,
scheduler: &Scheduler,
supervisor: u64,
tags: CommandTags,
receiver: &Receiver<MonitorMessage>,
live: &mut Vec<LiveChild>,
config: LifecycleConfig,
) -> Result<(), RegistryError> {
let mut early_exits = HashMap::new();
for child in live.clone() {
send_host_command(
id,
scheduler,
supervisor,
tags,
HostCommand::Send {
pid: child.pid,
message: child.spec.stop_message,
expects_reply: false,
},
)?;
let kind = early_exits.remove(&child.pid).map_or_else(
|| {
wait_for_child_exit(
id,
child.pid,
receiver,
config.operation_timeout,
&mut early_exits,
)
},
Ok,
)?;
if kind != TombstoneKind::Normal {
return Err(RegistryError::StopFailed {
id,
reason: process_failure(scheduler, child.spec.name, child.pid, kind),
});
}
live.retain(|item| item.pid != child.pid);
}
Ok(())
}
fn wait_for_child_exit(
id: ComponentId,
pid: u64,
receiver: &Receiver<MonitorMessage>,
timeout: std::time::Duration,
early_exits: &mut HashMap<u64, TombstoneKind>,
) -> Result<TombstoneKind, RegistryError> {
let deadline = Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
match receiver.recv_timeout(remaining) {
Ok(MonitorMessage::Signal(HostSignal::Exit { pid: exited, kind })) if exited == pid => {
return Ok(kind);
}
Ok(MonitorMessage::Signal(HostSignal::Exit { pid: exited, kind })) => {
early_exits.insert(exited, kind);
}
Ok(MonitorMessage::Signal(
HostSignal::Reply { .. } | HostSignal::UncorrelatedReply,
)) => {
}
Ok(MonitorMessage::Signal(HostSignal::Protocol(detail))) => {
return Err(RegistryError::SupervisorProtocol { id, detail });
}
Ok(MonitorMessage::Signal(HostSignal::Live | HostSignal::Barrier)) => {
return Err(RegistryError::SupervisorProtocol {
id,
detail: "unexpected supervisor control event during drain".to_owned(),
});
}
Ok(message) => refuse_during_drain(id, message),
Err(_) => {
return Err(RegistryError::StopFailed {
id,
reason: FailureReason {
child: None,
tombstone: None,
detail: format!("timed out awaiting process {pid} EXIT event"),
},
});
}
}
}
}
fn refuse_during_drain(id: ComponentId, message: MonitorMessage) {
match message {
MonitorMessage::Stop(reply) => {
let _delivered = reply.send(Err(RegistryError::AlreadyStopping { id }));
}
MonitorMessage::Probe { reply_tx, .. } => {
let _delivered = reply_tx.send(Err(RegistryError::AlreadyStopping { id }));
}
MonitorMessage::Send { reply_tx, .. } => {
let _delivered = reply_tx.send(Err(RegistryError::AlreadyStopping { id }));
}
MonitorMessage::Signal(_) => {}
}
}
pub(super) fn stop_supervisor(
id: ComponentId,
scheduler: &std::sync::Arc<Scheduler>,
supervisor: u64,
tags: CommandTags,
config: LifecycleConfig,
) -> Result<(), RegistryError> {
let observation = ExitObservation::register(scheduler, supervisor, config.operation_timeout)
.map_err(|error| observation_error(id, None, error))?;
if let Err(error) = send_host_command(id, scheduler, supervisor, tags, HostCommand::Stop) {
observation
.cancel()
.map_err(|cancel| observation_error(id, None, cancel))?;
return Err(error);
}
let reason = observation
.wait()
.map_err(|error| observation_error(id, None, error))?;
if reason == ExitReason::Normal {
Ok(())
} else {
Err(RegistryError::StopFailed {
id,
reason: FailureReason {
child: None,
tombstone: Some(exit_reason_kind(reason)),
detail: "host supervisor exited abnormally".to_owned(),
},
})
}
}
fn observation_error(
id: ComponentId,
child: Option<String>,
error: ObservationError,
) -> RegistryError {
let detail = match error {
ObservationError::Spawn(detail) => format!("tombstone observer spawn failed: {detail}"),
ObservationError::ReadyDeadline => "tombstone observer readiness timed out".to_owned(),
ObservationError::Registration(detail) => {
format!("tombstone monitor registration failed: {detail}")
}
ObservationError::TombstoneDeadline => "tombstone observation timed out".to_owned(),
ObservationError::Protocol(detail) => detail,
ObservationError::CleanupDeadline => "tombstone observer cleanup timed out".to_owned(),
};
RegistryError::StopFailed {
id,
reason: FailureReason {
child,
tombstone: None,
detail,
},
}
}
const fn exit_reason_kind(reason: ExitReason) -> TombstoneKind {
match reason {
ExitReason::Normal => TombstoneKind::Normal,
ExitReason::Kill => TombstoneKind::Kill,
ExitReason::Killed => TombstoneKind::Killed,
ExitReason::Error => TombstoneKind::Error,
ExitReason::NoConnection => TombstoneKind::NoConnection,
ExitReason::NoProc => TombstoneKind::NoProcess,
}
}