use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use crate::bytecode::Value;
use super::error::{report_fault, SpawnError};
use super::handle::FlowHandle;
use super::process::{FlowId, FlowOutcome, RestartPolicy};
use super::runtime::RuntimeSpawner;
use super::sync_lock;
const DEFAULT_MAX_RESTARTS: u32 = 3;
const DEFAULT_MAX_PERIOD: Duration = Duration::from_secs(5);
#[derive(Clone, Debug)]
pub struct ChildSpec {
pub name: String,
pub function: u32,
pub args: Vec<Value>,
pub restart: RestartPolicy,
}
impl ChildSpec {
pub fn new(name: impl Into<String>, function: u32) -> Self {
ChildSpec {
name: name.into(),
function,
args: Vec::new(),
restart: RestartPolicy::OnFailure,
}
}
pub fn args(mut self, args: Vec<Value>) -> Self {
self.args = args;
self
}
pub fn restart(mut self, restart: RestartPolicy) -> Self {
self.restart = restart;
self
}
}
#[derive(Clone, Debug)]
pub struct SupervisorConfig {
pub max_restarts: u32,
pub max_period: Duration,
}
impl Default for SupervisorConfig {
fn default() -> Self {
SupervisorConfig {
max_restarts: DEFAULT_MAX_RESTARTS,
max_period: DEFAULT_MAX_PERIOD,
}
}
}
struct ChildExit {
id: FlowId,
outcome: FlowOutcome,
}
struct LiveChild {
spec: ChildSpec,
}
struct Inner {
spawner: RuntimeSpawner,
config: SupervisorConfig,
events: Mutex<VecDeque<ChildExit>>,
cvar: Condvar,
children: Mutex<HashMap<FlowId, LiveChild>>,
restart_times: Mutex<VecDeque<Instant>>,
intensity_exceeded: AtomicBool,
shutdown: AtomicBool,
}
#[derive(Clone)]
pub(crate) struct SupervisorLink {
inner: Arc<Inner>,
}
impl SupervisorLink {
pub(crate) fn notify(&self, id: FlowId, outcome: FlowOutcome) {
match sync_lock::lock(&self.inner.events, "SupervisorLink::notify") {
Ok(mut events) => {
events.push_back(ChildExit { id, outcome });
self.inner.cvar.notify_one();
}
Err(e) => report_fault(e),
}
}
}
pub struct Supervisor {
inner: Arc<Inner>,
thread: Option<JoinHandle<()>>,
}
impl Supervisor {
pub fn new(spawner: RuntimeSpawner) -> Result<Self, SpawnError> {
Self::with_config(spawner, SupervisorConfig::default())
}
pub fn with_config(spawner: RuntimeSpawner, config: SupervisorConfig) -> Result<Self, SpawnError> {
let inner = Arc::new(Inner {
spawner,
config,
events: Mutex::new(VecDeque::new()),
cvar: Condvar::new(),
children: Mutex::new(HashMap::new()),
restart_times: Mutex::new(VecDeque::new()),
intensity_exceeded: AtomicBool::new(false),
shutdown: AtomicBool::new(false),
});
let drive_inner = inner.clone();
let thread = std::thread::Builder::new()
.name("byteflow-supervisor".into())
.spawn(move || drive(drive_inner))
.map_err(|e| SpawnError::ThreadSpawnFailed(e.to_string()))?;
Ok(Supervisor {
inner,
thread: Some(thread),
})
}
pub fn start_child(&self, spec: ChildSpec) -> Result<FlowHandle, SpawnError> {
spawn_child(&self.inner, spec)
}
pub fn live_children(&self) -> usize {
match sync_lock::lock(&self.inner.children, "Supervisor::live_children") {
Ok(g) => g.len(),
Err(e) => {
report_fault(e);
0
}
}
}
pub fn intensity_exceeded(&self) -> bool {
self.inner.intensity_exceeded.load(Ordering::Acquire)
}
pub fn shutdown(mut self) {
self.inner.shutdown.store(true, Ordering::Release);
self.inner.cvar.notify_all();
if let Some(t) = self.thread.take() {
let _ = t.join();
}
}
}
fn spawn_child(inner: &Arc<Inner>, spec: ChildSpec) -> Result<FlowHandle, SpawnError> {
let link = SupervisorLink {
inner: inner.clone(),
};
let mut children = match sync_lock::lock(&inner.children, "spawn_child") {
Ok(c) => c,
Err(e) => {
report_fault(e);
return Err(SpawnError::VmInit(
"supervisor child table poisoned".into(),
));
}
};
let handle = inner.spawner.spawn_linked(
spec.function,
&spec.args,
spec.restart,
link,
)?;
children.insert(handle.id(), LiveChild { spec });
Ok(handle)
}
fn should_restart(policy: RestartPolicy, outcome: &FlowOutcome) -> bool {
match policy {
RestartPolicy::Always => true,
RestartPolicy::OnFailure => matches!(outcome, FlowOutcome::Failed(_)),
RestartPolicy::Never => false,
}
}
fn intensity_hit(inner: &Inner) -> bool {
let now = Instant::now();
let mut times = match sync_lock::lock(&inner.restart_times, "intensity_hit") {
Ok(t) => t,
Err(e) => {
report_fault(e);
return true;
}
};
times.push_back(now);
let window_start = now.checked_sub(inner.config.max_period).unwrap_or(now);
while times.front().map(|t| *t < window_start).unwrap_or(false) {
times.pop_front();
}
if times.len() as u32 > inner.config.max_restarts {
inner.intensity_exceeded.store(true, Ordering::Release);
true
} else {
false
}
}
fn drive(inner: Arc<Inner>) {
loop {
if inner.shutdown.load(Ordering::Acquire) {
return;
}
let exit = {
let mut events = match sync_lock::lock(&inner.events, "supervisor::drive") {
Ok(e) => e,
Err(e) => {
report_fault(e);
return;
}
};
loop {
if inner.shutdown.load(Ordering::Acquire) {
return;
}
if let Some(exit) = events.pop_front() {
break exit;
}
match sync_lock::wait_timeout(
&inner.cvar,
events,
Duration::from_millis(100),
"supervisor::wait",
) {
Ok((guard, _)) => events = guard,
Err(e) => {
report_fault(e);
return;
}
}
}
};
handle_exit(&inner, exit);
}
}
fn handle_exit(inner: &Arc<Inner>, exit: ChildExit) {
let spec = {
let mut children = match sync_lock::lock(&inner.children, "handle_exit") {
Ok(c) => c,
Err(e) => {
report_fault(e);
return;
}
};
match children.remove(&exit.id) {
Some(live) => live.spec,
None => return,
}
};
if !should_restart(spec.restart, &exit.outcome) {
return;
}
if inner.intensity_exceeded.load(Ordering::Acquire) || intensity_hit(inner) {
return;
}
let _ = spawn_child(inner, spec);
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::{Duration, Instant};
use crate::bytecode::{Chunk, ChunkBuilder, Value};
use crate::scheduler::runtime::{Runtime, RuntimeConfig};
fn trap_chunk() -> Chunk {
let mut b = ChunkBuilder::new("trap");
b.begin_function("boom", 0, 1);
b.emit_trap(1);
b.finish()
}
fn ok_chunk() -> Chunk {
let mut b = ChunkBuilder::new("ok");
b.begin_function("main", 0, 1);
b.emit_load_imm(0, 7);
b.emit_return(0);
b.finish()
}
fn tiny_runtime(chunk: Chunk) -> Runtime {
Runtime::with_config(
chunk,
RuntimeConfig {
workers: 1,
quantum: 1_000,
},
)
.expect("runtime")
}
fn wait_until(mut pred: impl FnMut() -> bool) {
let start = Instant::now();
while !pred() {
assert!(
start.elapsed() < Duration::from_secs(2),
"supervisor test timed out"
);
std::thread::sleep(Duration::from_millis(5));
}
}
#[test]
fn on_failure_does_not_restart_a_clean_exit() {
let rt = tiny_runtime(ok_chunk());
let sup = Supervisor::new(rt.spawner()).expect("supervisor");
let outcome = sup
.start_child(ChildSpec::new("main", 0).restart(RestartPolicy::OnFailure))
.expect("start_child")
.join();
wait_until(|| sup.live_children() == 0);
let spawned = rt.metrics().processes_spawned;
sup.shutdown();
rt.shutdown();
assert!(matches!(outcome, FlowOutcome::Completed(_)));
assert_eq!(spawned, 1);
}
#[test]
fn on_failure_restarts_until_intensity() {
let rt = tiny_runtime(trap_chunk());
let sup = Supervisor::with_config(
rt.spawner(),
SupervisorConfig {
max_restarts: 2,
max_period: Duration::from_secs(5),
},
)
.expect("supervisor");
let _first = sup
.start_child(ChildSpec::new("boom", 0).restart(RestartPolicy::OnFailure))
.expect("start_child");
wait_until(|| sup.intensity_exceeded() && rt.metrics().processes_failed >= 3);
let spawned = rt.metrics().processes_spawned;
let failed = rt.metrics().processes_failed;
sup.shutdown();
rt.shutdown();
assert_eq!(spawned, 3);
assert_eq!(failed, 3);
}
#[test]
fn always_restarts_a_clean_exit_until_intensity() {
let rt = tiny_runtime(ok_chunk());
let sup = Supervisor::with_config(
rt.spawner(),
SupervisorConfig {
max_restarts: 2,
max_period: Duration::from_secs(5),
},
)
.expect("supervisor");
let _ = sup
.start_child(ChildSpec::new("main", 0).restart(RestartPolicy::Always))
.expect("start_child");
wait_until(|| sup.intensity_exceeded() && rt.metrics().processes_completed >= 3);
let spawned = rt.metrics().processes_spawned;
sup.shutdown();
rt.shutdown();
assert_eq!(spawned, 3);
}
#[test]
fn policy_table() {
let ok = FlowOutcome::Completed(Value::Unit);
let fail = FlowOutcome::Failed("boom".into());
assert!(should_restart(RestartPolicy::Always, &ok));
assert!(should_restart(RestartPolicy::Always, &fail));
assert!(!should_restart(RestartPolicy::OnFailure, &ok));
assert!(should_restart(RestartPolicy::OnFailure, &fail));
assert!(!should_restart(RestartPolicy::Never, &ok));
assert!(!should_restart(RestartPolicy::Never, &fail));
}
}