use std::sync::Arc;
use std::time::Duration;
use parking_lot::Mutex;
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use crate::adapter::net::behavior::fold::{IslandId, IslandQuery, NodeId};
use crate::adapter::net::behavior::meshos::event_loop::{MeshOsHandle, MeshOsSnapshotReader};
use crate::adapter::net::behavior::meshos::{
DaemonIntent, DaemonIntentUpdate, DaemonLifecycleSignal, DaemonRef, MeshOsEvent,
};
use crate::adapter::net::compute::{DaemonLifecycleEvent, DaemonLifecycleObserver};
use crate::adapter::net::cortex::workflow::{ActiveClaim, TaskId, WorkflowAdapter};
use crate::adapter::net::MeshNode;
use super::lifecycle::LifecycleTransition;
use super::liveness::project_liveness_from_snapshot;
use super::runtime::SchedulerBridge;
fn resolve_island_host(mesh: &MeshNode, island: IslandId) -> Option<NodeId> {
mesh.island_fold()
.query(IslandQuery::Get(island))
.first()
.map(|(_, record)| record.host)
}
fn event_to_signal(event: DaemonLifecycleEvent) -> (DaemonRef, DaemonLifecycleSignal) {
match event {
DaemonLifecycleEvent::Registered { id, name, at } => (
DaemonRef { id, name },
DaemonLifecycleSignal::Started { at },
),
DaemonLifecycleEvent::Unregistered { id, name, at } => (
DaemonRef { id, name },
DaemonLifecycleSignal::ExitedCleanly { at },
),
DaemonLifecycleEvent::Crashed {
id,
name,
at,
reason,
} => (
DaemonRef { id, name },
DaemonLifecycleSignal::Crashed { at, reason },
),
DaemonLifecycleEvent::HealthChanged {
id,
name,
at,
health,
} => (
DaemonRef { id, name },
DaemonLifecycleSignal::HealthChanged { at, health },
),
DaemonLifecycleEvent::SaturationChanged {
id,
name,
at,
saturation,
} => (
DaemonRef { id, name },
DaemonLifecycleSignal::SaturationChanged { at, saturation },
),
}
}
struct BridgeLifecycleObserver {
bridge: Arc<Mutex<SchedulerBridge>>,
workflow: Arc<WorkflowAdapter>,
}
impl DaemonLifecycleObserver for BridgeLifecycleObserver {
fn observe(&self, event: DaemonLifecycleEvent) {
let (daemon, signal) = event_to_signal(event);
let transition = {
let state = self.workflow.state();
let guard = state.read();
self.bridge
.lock()
.lifecycle_transition(&signal, &daemon, &guard)
};
match transition {
Some(LifecycleTransition::ConfirmRunning(task)) => {
let _ = self.workflow.start(task);
}
Some(LifecycleTransition::FailStep(task)) => {
let _ = self.workflow.fail(task);
self.bridge.lock().on_released(task);
}
None => {}
}
}
}
struct FanOutLifecycleObserver {
observers: Vec<Arc<dyn DaemonLifecycleObserver>>,
}
impl DaemonLifecycleObserver for FanOutLifecycleObserver {
fn observe(&self, event: DaemonLifecycleEvent) {
for obs in &self.observers {
obs.observe(event.clone());
}
}
}
pub fn fan_out_lifecycle(
observers: Vec<Arc<dyn DaemonLifecycleObserver>>,
) -> Arc<dyn DaemonLifecycleObserver> {
Arc::new(FanOutLifecycleObserver { observers })
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct TickReport {
pub published: usize,
pub down: usize,
}
pub struct SchedulerBridgeDriver {
bridge: Arc<Mutex<SchedulerBridge>>,
workflow: Arc<WorkflowAdapter>,
mesh: Arc<MeshNode>,
handle: MeshOsHandle,
snapshot: MeshOsSnapshotReader,
last_published: Mutex<std::collections::HashMap<DaemonRef, DaemonIntentUpdate>>,
shutdown: Arc<Notify>,
}
impl SchedulerBridgeDriver {
pub fn new(
workflow: Arc<WorkflowAdapter>,
mesh: Arc<MeshNode>,
handle: MeshOsHandle,
snapshot: MeshOsSnapshotReader,
) -> Self {
Self {
bridge: Arc::new(Mutex::new(SchedulerBridge::new())),
workflow,
mesh,
handle,
snapshot,
last_published: Mutex::new(std::collections::HashMap::new()),
shutdown: Arc::new(Notify::new()),
}
}
pub fn on_running(&self, task: TaskId, claim: ActiveClaim) {
self.bridge.lock().on_running(task, claim);
}
pub fn on_released(&self, task: TaskId) -> Option<ActiveClaim> {
self.bridge.lock().on_released(task)
}
pub fn tick(&self) -> TickReport {
let delta = {
let snapshot = self.snapshot.load();
project_liveness_from_snapshot(&snapshot)
};
let down = delta.down.len();
self.mesh
.set_liveness_down(delta.down.into_iter().collect());
let intents: Vec<DaemonIntentUpdate> = {
let state = self.workflow.state();
let guard = state.read();
let mesh = &self.mesh;
self.bridge
.lock()
.desired_intents(&guard, |island| resolve_island_host(mesh, island))
};
let mut published = 0;
let mut last = self.last_published.lock();
let mut current: std::collections::HashSet<DaemonRef> = std::collections::HashSet::new();
for intent in intents {
current.insert(intent.daemon.clone());
if last.get(&intent.daemon) == Some(&intent) {
continue;
}
if self
.handle
.try_publish(MeshOsEvent::DaemonIntentUpdate(intent.clone()))
.is_ok()
{
last.insert(intent.daemon.clone(), intent);
published += 1;
}
}
let vanished: Vec<DaemonRef> = last
.keys()
.filter(|daemon| !current.contains(*daemon))
.cloned()
.collect();
for daemon in vanished {
let was_run = last.get(&daemon).map(|u| u.intent) == Some(DaemonIntent::Run);
if !was_run {
last.remove(&daemon);
continue;
}
let stop = DaemonIntentUpdate {
daemon: daemon.clone(),
intent: DaemonIntent::Stop,
node: None,
};
if self
.handle
.try_publish(MeshOsEvent::DaemonIntentUpdate(stop))
.is_ok()
{
last.remove(&daemon);
published += 1;
}
}
drop(last);
TickReport { published, down }
}
pub fn lifecycle_observer(&self) -> Arc<dyn DaemonLifecycleObserver> {
Arc::new(BridgeLifecycleObserver {
bridge: Arc::clone(&self.bridge),
workflow: Arc::clone(&self.workflow),
})
}
pub fn spawn(self: Arc<Self>, interval: Duration) -> JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = ticker.tick() => {
self.tick();
}
_ = self.shutdown.notified() => break,
}
}
})
}
pub fn shutdown(&self) {
self.shutdown.notify_one();
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use super::*;
#[test]
fn event_to_signal_mirrors_the_meshos_sink_mapping() {
let at = Instant::now();
let (d, s) = event_to_signal(DaemonLifecycleEvent::Registered {
id: 7,
name: "task/1".into(),
at,
});
assert_eq!(
d,
DaemonRef {
id: 7,
name: "task/1".into()
}
);
assert_eq!(s, DaemonLifecycleSignal::Started { at });
let (_, s) = event_to_signal(DaemonLifecycleEvent::Unregistered {
id: 7,
name: "task/1".into(),
at,
});
assert_eq!(s, DaemonLifecycleSignal::ExitedCleanly { at });
let (_, s) = event_to_signal(DaemonLifecycleEvent::Crashed {
id: 7,
name: "task/1".into(),
at,
reason: "oom".into(),
});
assert_eq!(
s,
DaemonLifecycleSignal::Crashed {
at,
reason: "oom".into(),
}
);
}
struct Counter(Arc<AtomicUsize>);
impl DaemonLifecycleObserver for Counter {
fn observe(&self, _event: DaemonLifecycleEvent) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
#[test]
fn fan_out_forwards_to_every_observer() {
let a = Arc::new(AtomicUsize::new(0));
let b = Arc::new(AtomicUsize::new(0));
let fan = fan_out_lifecycle(vec![
Arc::new(Counter(a.clone())),
Arc::new(Counter(b.clone())),
]);
fan.observe(DaemonLifecycleEvent::Registered {
id: 1,
name: "x".into(),
at: Instant::now(),
});
assert_eq!(a.load(Ordering::Relaxed), 1);
assert_eq!(b.load(Ordering::Relaxed), 1);
}
}