use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{self, Receiver, Sender, SyncSender};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use beamr::atom::Atom;
use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use beamr::scheduler::Scheduler;
use beamr::term::Term;
use beamr::term::boxed::Tuple;
use crate::component::{ChildArgument, ChildSpec, ComponentId, SupervisionPolicy};
use crate::error::{FailureReason, RegistryError, TombstoneKind};
use crate::event::{EventHub, LifecycleState};
use crate::status::ChildStatus;
pub(crate) struct TreeHandle {
tx: Sender<MonitorMessage>,
join: Option<JoinHandle<Result<(), RegistryError>>>,
}
impl TreeHandle {
pub(crate) fn stop(mut self, id: ComponentId) -> Result<(), RegistryError> {
let (reply_tx, reply_rx) = mpsc::sync_channel(1);
self.tx
.send(MonitorMessage::Stop(reply_tx))
.map_err(|_| RegistryError::MonitorTerminated { id })?;
let result = reply_rx
.recv()
.map_err(|_| RegistryError::MonitorTerminated { id })?;
let joined = self
.join
.take()
.ok_or(RegistryError::MonitorTerminated { id })?
.join()
.map_err(|_| RegistryError::MonitorTerminated { id })?;
result.and(joined)
}
pub(crate) fn probe(&self, id: ComponentId, child: String) -> Result<i64, RegistryError> {
let (reply_tx, reply_rx) = mpsc::sync_channel(1);
self.tx
.send(MonitorMessage::Probe { child, reply_tx })
.map_err(|_| RegistryError::MonitorTerminated { id })?;
reply_rx
.recv()
.map_err(|_| RegistryError::MonitorTerminated { id })?
}
pub(crate) fn send(
&self,
id: ComponentId,
child: String,
message: i64,
) -> Result<(), RegistryError> {
let (reply_tx, reply_rx) = mpsc::sync_channel(1);
self.tx
.send(MonitorMessage::Send {
child,
message,
reply_tx,
})
.map_err(|_| RegistryError::MonitorTerminated { id })?;
reply_rx
.recv()
.map_err(|_| RegistryError::MonitorTerminated { id })?
}
pub(crate) fn join_finished(mut self, id: ComponentId) -> Result<(), RegistryError> {
self.join
.take()
.ok_or(RegistryError::MonitorTerminated { id })?
.join()
.map_err(|_| RegistryError::MonitorTerminated { id })?
}
}
pub(crate) fn start_tree(
id: ComponentId,
scheduler: Arc<Scheduler>,
specs: &[ChildSpec],
policy: SupervisionPolicy,
config: LifecycleConfig,
status: &Arc<SharedStatus>,
events: Arc<EventHub>,
) -> Result<TreeHandle, RegistryError> {
let (signal_tx, signal_rx) = mpsc::channel();
let control_tx = signal_tx.clone();
let tags = CommandTags::new(&scheduler);
let supervisor = spawn_supervisor(id, &scheduler, signal_tx, tags)?;
match recv_signal(id, &signal_rx, config.operation_timeout)? {
HostSignal::Live => {}
other => {
return Err(RegistryError::SupervisorProtocol {
id,
detail: format!("expected supervisor live event, got {other:?}"),
});
}
}
let live = start_children(
id, &scheduler, supervisor, specs, status, &signal_rx, config,
)?;
*status
.supervisor
.lock()
.map_err(|_| RegistryError::SynchronizationPoisoned)? = Some(supervisor);
let monitor_status = Arc::clone(status);
let monitor = thread::Builder::new()
.name(format!("frame-component-{id}"))
.spawn(move || {
run_monitor(MonitorContext {
id,
scheduler,
supervisor,
tags,
receiver: signal_rx,
live,
policy,
config,
status: monitor_status,
events,
})
})
.map_err(|error| RegistryError::Spawn {
id,
process: "monitor thread".to_owned(),
detail: error.to_string(),
})?;
Ok(TreeHandle {
tx: control_tx,
join: Some(monitor),
})
}
fn start_children(
id: ComponentId,
scheduler: &Arc<Scheduler>,
supervisor: u64,
specs: &[ChildSpec],
status: &SharedStatus,
receiver: &Receiver<MonitorMessage>,
config: LifecycleConfig,
) -> Result<Vec<LiveChild>, RegistryError> {
let tags = CommandTags::new(scheduler);
let mut live = Vec::with_capacity(specs.len());
for spec in specs {
let pid = spawn_child(id, scheduler, supervisor, spec)?;
live.push(LiveChild {
spec: spec.clone(),
pid,
});
publish_children(status, &live)?;
send_host_command(
id,
scheduler,
supervisor,
tags,
HostCommand::Send {
pid,
message: spec.liveness_message,
expects_reply: true,
},
)?;
match recv_signal(id, receiver, config.operation_timeout)? {
HostSignal::Reply { pid: reply_pid, .. } if reply_pid == pid => {}
HostSignal::Exit {
pid: exit_pid,
kind,
} => {
let reason = process_failure(scheduler, spec.name.clone(), exit_pid, kind);
live.retain(|child| child.pid != exit_pid);
publish_children(status, &live)?;
cleanup_start_failure(
id, scheduler, supervisor, tags, receiver, &mut live, config,
)?;
return Err(RegistryError::StartFailed { id, reason });
}
other => {
let reason = FailureReason {
child: Some(spec.name.clone()),
tombstone: None,
detail: format!("unexpected liveness signal: {other:?}"),
};
cleanup_start_failure(
id, scheduler, supervisor, tags, receiver, &mut live, config,
)?;
return Err(RegistryError::StartFailed { id, reason });
}
}
}
Ok(live)
}
#[derive(Debug)]
enum HostSignal {
Live,
Barrier,
Reply { pid: u64, value: i64 },
UncorrelatedReply,
Exit { pid: u64, kind: TombstoneKind },
Protocol(String),
}
enum MonitorMessage {
Signal(HostSignal),
Stop(SyncSender<Result<(), RegistryError>>),
Probe {
child: String,
reply_tx: SyncSender<Result<i64, RegistryError>>,
},
Send {
child: String,
message: i64,
reply_tx: SyncSender<Result<(), RegistryError>>,
},
}
#[derive(Clone)]
struct HostMailbox {
tx: Sender<MonitorMessage>,
tags: CommandTags,
}
struct HostSupervisor {
mailbox: HostMailbox,
pending_replies: VecDeque<u64>,
announced_live: bool,
}
impl NativeHandler for HostSupervisor {
fn handle(&mut self, context: &mut NativeContext<'_>) -> NativeOutcome {
context.set_trap_exit(true);
if !self.announced_live {
self.announced_live = true;
if self
.mailbox
.tx
.send(MonitorMessage::Signal(HostSignal::Live))
.is_err()
{
return NativeOutcome::Stop(ExitReason::Error);
}
}
while let Some(message) = context.recv() {
match decode_command(message, self.mailbox.tags) {
CommandDecode::Command(HostCommand::Barrier) => {
if self.send_signal(HostSignal::Barrier).is_err() {
return NativeOutcome::Stop(ExitReason::Error);
}
}
CommandDecode::Command(HostCommand::Send {
pid,
message,
expects_reply,
}) => {
if expects_reply {
self.pending_replies.push_back(pid);
}
context.send(pid, Term::small_int(message));
}
CommandDecode::Command(HostCommand::Stop) => {
return NativeOutcome::Stop(ExitReason::Normal);
}
CommandDecode::Malformed(detail) => {
if self.send_signal(HostSignal::Protocol(detail)).is_err() {
return NativeOutcome::Stop(ExitReason::Error);
}
}
CommandDecode::NotCommand => {
let signal = decode_message(message, &mut self.pending_replies);
let observed_exit = matches!(signal, HostSignal::Exit { .. });
if self.send_signal(signal).is_err() && !observed_exit {
return NativeOutcome::Stop(ExitReason::Error);
}
}
}
}
NativeOutcome::Wait
}
}
impl HostSupervisor {
fn send_signal(&self, signal: HostSignal) -> Result<(), ()> {
self.mailbox
.tx
.send(MonitorMessage::Signal(signal))
.map_err(|_| ())
}
}
#[derive(Clone)]
struct LiveChild {
spec: ChildSpec,
pid: u64,
}
struct MonitorContext {
id: ComponentId,
scheduler: Arc<Scheduler>,
supervisor: u64,
tags: CommandTags,
receiver: Receiver<MonitorMessage>,
live: Vec<LiveChild>,
policy: SupervisionPolicy,
config: LifecycleConfig,
status: Arc<SharedStatus>,
events: Arc<EventHub>,
}
fn spawn_supervisor(
id: ComponentId,
scheduler: &Arc<Scheduler>,
tx: Sender<MonitorMessage>,
tags: CommandTags,
) -> Result<u64, RegistryError> {
let mailbox = HostMailbox { tx, tags };
scheduler
.spawn_native(Box::new(move || {
Box::new(HostSupervisor {
mailbox: mailbox.clone(),
pending_replies: VecDeque::new(),
announced_live: false,
})
}))
.map_err(|error| RegistryError::Spawn {
id,
process: "host supervisor".to_owned(),
detail: error.to_string(),
})
}
fn spawn_child(
id: ComponentId,
scheduler: &Scheduler,
supervisor: u64,
spec: &ChildSpec,
) -> Result<u64, RegistryError> {
let module = scheduler.atom_table().intern(&spec.module);
let function = scheduler.atom_table().intern(&spec.function);
let args = spec
.arguments
.iter()
.map(|argument| match argument {
ChildArgument::SupervisorPid => Term::pid(supervisor),
ChildArgument::Integer(value) => Term::small_int(*value),
ChildArgument::Atom(value) => Term::atom(scheduler.atom_table().intern(value)),
})
.collect();
scheduler
.spawn_link(supervisor, module, function, args)
.map_err(|error| RegistryError::Spawn {
id,
process: spec.name.clone(),
detail: error.to_string(),
})
}
fn decode_message(message: Term, pending: &mut VecDeque<u64>) -> HostSignal {
if let Some(value) = message.as_small_int() {
return pending
.pop_front()
.map_or(HostSignal::UncorrelatedReply, |pid| HostSignal::Reply {
pid,
value,
});
}
let Some(tuple) = Tuple::new(message) else {
return HostSignal::Protocol("supervisor received a non-tuple message".to_owned());
};
if tuple.arity() != 3 || tuple.get(0).and_then(Term::as_atom) != Some(Atom::EXIT) {
return HostSignal::Protocol("supervisor received a non-EXIT tuple".to_owned());
}
match (
tuple.get(1).and_then(Term::as_pid),
tuple.get(2).and_then(Term::as_atom),
) {
(Some(pid), Some(reason)) => reason_to_tombstone(reason).map_or_else(
|| HostSignal::Protocol("EXIT carried an unknown reason atom".to_owned()),
|kind| HostSignal::Exit { pid, kind },
),
_ => HostSignal::Protocol("EXIT carried malformed fields".to_owned()),
}
}
fn reason_to_tombstone(reason: Atom) -> Option<TombstoneKind> {
if reason == Atom::NORMAL {
Some(TombstoneKind::Normal)
} else if reason == Atom::KILL {
Some(TombstoneKind::Kill)
} else if reason == Atom::KILLED {
Some(TombstoneKind::Killed)
} else if reason == Atom::ERROR {
Some(TombstoneKind::Error)
} else if reason == Atom::NOCONNECTION {
Some(TombstoneKind::NoConnection)
} else if reason == Atom::NOPROC {
Some(TombstoneKind::NoProcess)
} else {
None
}
}
fn recv_signal(
id: ComponentId,
receiver: &Receiver<MonitorMessage>,
timeout: Duration,
) -> Result<HostSignal, RegistryError> {
match receiver.recv_timeout(timeout) {
Ok(MonitorMessage::Signal(signal)) => Ok(signal),
Ok(_) => Err(RegistryError::SupervisorProtocol {
id,
detail: "lifecycle request arrived before monitor start".to_owned(),
}),
Err(error) => Err(RegistryError::SupervisorProtocol {
id,
detail: error.to_string(),
}),
}
}
fn publish_children(status: &SharedStatus, live: &[LiveChild]) -> Result<(), RegistryError> {
*status
.children
.lock()
.map_err(|_| RegistryError::SynchronizationPoisoned)? = live
.iter()
.map(|child| ChildStatus {
name: child.spec.name.clone(),
pid: child.pid,
})
.collect();
Ok(())
}
fn process_failure(
scheduler: &Scheduler,
child: String,
pid: u64,
kind: TombstoneKind,
) -> FailureReason {
let error = scheduler.take_exit_error(pid);
let exception = scheduler.take_exit_exception(pid);
FailureReason {
child: Some(child),
tombstone: Some(kind),
detail: format!("error={error:?}, exception={exception:?}"),
}
}
fn cleanup_start_failure(
id: ComponentId,
scheduler: &Arc<Scheduler>,
supervisor: u64,
tags: CommandTags,
receiver: &Receiver<MonitorMessage>,
live: &mut Vec<LiveChild>,
config: LifecycleConfig,
) -> Result<(), RegistryError> {
drain_children(id, scheduler, supervisor, tags, receiver, live, config)?;
stop_supervisor(id, scheduler, supervisor, tags, config)
}
mod command;
mod monitor;
pub(crate) mod observation;
#[cfg(test)]
mod tests;
mod types;
use command::{
CommandDecode, CommandTags, HostCommand, decode as decode_command, send_host_command,
};
use monitor::{drain_children, run_monitor, stop_supervisor};
pub use types::LifecycleConfig;
pub(crate) use types::{SharedStatus, StatusState};