use std::sync::atomic::{AtomicU64, Ordering};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use crate::process::Pid;
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ExitReason {
Normal,
Crash(String),
Killed,
Timeout,
}
impl std::fmt::Display for ExitReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ExitReason::Normal => write!(f, "normal"),
ExitReason::Crash(msg) => write!(f, "crash: {msg}"),
ExitReason::Killed => write!(f, "killed"),
ExitReason::Timeout => write!(f, "timeout"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessLink {
pub pid_a: Pid,
pub pid_b: Pid,
}
#[derive(Debug, Clone)]
pub struct ProcessMonitor {
pub watcher: Pid,
pub target: Pid,
pub ref_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProcessDown {
pub pid: Pid,
pub reason: ExitReason,
pub ref_id: String,
}
pub struct MonitorRegistry {
links: DashMap<Pid, Vec<Pid>>,
monitors: DashMap<Pid, Vec<ProcessMonitor>>,
next_ref: AtomicU64,
}
impl MonitorRegistry {
pub fn new() -> Self {
Self {
links: DashMap::new(),
monitors: DashMap::new(),
next_ref: AtomicU64::new(1),
}
}
pub fn link(&self, pid_a: Pid, pid_b: Pid) {
self.links
.entry(pid_a)
.or_default()
.value_mut()
.retain(|&p| p != pid_b); self.links.entry(pid_a).or_default().push(pid_b);
self.links
.entry(pid_b)
.or_default()
.value_mut()
.retain(|&p| p != pid_a);
self.links.entry(pid_b).or_default().push(pid_a);
}
pub fn unlink(&self, pid_a: Pid, pid_b: Pid) {
if let Some(mut links) = self.links.get_mut(&pid_a) {
links.retain(|&p| p != pid_b);
}
if let Some(mut links) = self.links.get_mut(&pid_b) {
links.retain(|&p| p != pid_a);
}
}
pub fn monitor(&self, watcher: Pid, target: Pid) -> String {
let ref_id = format!("mon-{}", self.next_ref.fetch_add(1, Ordering::Relaxed));
let monitor = ProcessMonitor {
watcher,
target,
ref_id: ref_id.clone(),
};
self.monitors.entry(target).or_default().push(monitor);
ref_id
}
pub fn demonitor(&self, ref_id: &str) {
for mut entry in self.monitors.iter_mut() {
entry.value_mut().retain(|m| m.ref_id != ref_id);
}
}
pub fn process_exited(&self, pid: Pid, reason: &ExitReason) -> (Vec<(Pid, ExitReason)>, Vec<ProcessDown>) {
let mut link_signals = Vec::new();
let mut down_signals = Vec::new();
if let Some((_, linked_pids)) = self.links.remove(&pid) {
for linked_pid in &linked_pids {
link_signals.push((*linked_pid, reason.clone()));
if let Some(mut reverse) = self.links.get_mut(linked_pid) {
reverse.retain(|&p| p != pid);
}
}
}
if let Some((_, monitors)) = self.monitors.remove(&pid) {
for mon in monitors {
down_signals.push(ProcessDown {
pid,
reason: reason.clone(),
ref_id: mon.ref_id,
});
}
}
(link_signals, down_signals)
}
pub fn get_links(&self, pid: Pid) -> Vec<Pid> {
self.links
.get(&pid)
.map(|v| v.value().clone())
.unwrap_or_default()
}
pub fn get_monitors(&self, pid: Pid) -> Vec<ProcessMonitor> {
self.monitors
.get(&pid)
.map(|v| v.value().clone())
.unwrap_or_default()
}
pub fn is_linked(&self, pid_a: Pid, pid_b: Pid) -> bool {
self.links
.get(&pid_a)
.map(|v| v.contains(&pid_b))
.unwrap_or(false)
}
}
impl Default for MonitorRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn link_creates_bidirectional() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
assert!(reg.is_linked(1, 2));
assert!(reg.is_linked(2, 1));
}
#[test]
fn link_idempotent() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
reg.link(1, 2);
assert_eq!(reg.get_links(1).len(), 1);
assert_eq!(reg.get_links(2).len(), 1);
}
#[test]
fn unlink_removes_bidirectional() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
reg.unlink(1, 2);
assert!(!reg.is_linked(1, 2));
assert!(!reg.is_linked(2, 1));
}
#[test]
fn monitor_returns_unique_ref() {
let reg = MonitorRegistry::new();
let ref1 = reg.monitor(10, 20);
let ref2 = reg.monitor(10, 20);
assert_ne!(ref1, ref2);
}
#[test]
fn demonitor_removes_monitor() {
let reg = MonitorRegistry::new();
let ref_id = reg.monitor(10, 20);
assert_eq!(reg.get_monitors(20).len(), 1);
reg.demonitor(&ref_id);
assert_eq!(reg.get_monitors(20).len(), 0);
}
#[test]
fn process_exited_delivers_link_signals() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
reg.link(1, 3);
let (links, downs) = reg.process_exited(1, &ExitReason::Crash("panic".into()));
assert_eq!(links.len(), 2);
assert!(downs.is_empty());
let pids: Vec<Pid> = links.iter().map(|(p, _)| *p).collect();
assert!(pids.contains(&2));
assert!(pids.contains(&3));
assert!(!reg.is_linked(2, 1));
assert!(!reg.is_linked(3, 1));
}
#[test]
fn process_exited_delivers_down_signals() {
let reg = MonitorRegistry::new();
let ref1 = reg.monitor(10, 1);
let ref2 = reg.monitor(20, 1);
let (links, downs) = reg.process_exited(1, &ExitReason::Normal);
assert!(links.is_empty());
assert_eq!(downs.len(), 2);
let refs: Vec<&str> = downs.iter().map(|d| d.ref_id.as_str()).collect();
assert!(refs.contains(&ref1.as_str()));
assert!(refs.contains(&ref2.as_str()));
}
#[test]
fn process_exited_delivers_both_links_and_monitors() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
reg.monitor(10, 1);
let (links, downs) = reg.process_exited(1, &ExitReason::Killed);
assert_eq!(links.len(), 1);
assert_eq!(downs.len(), 1);
assert_eq!(links[0].0, 2);
assert_eq!(downs[0].reason, ExitReason::Killed);
}
#[test]
fn normal_exit_delivers_normal_reason() {
let reg = MonitorRegistry::new();
reg.link(1, 2);
let (links, _) = reg.process_exited(1, &ExitReason::Normal);
assert_eq!(links[0].1, ExitReason::Normal);
}
#[test]
fn multiple_monitors_on_same_target() {
let reg = MonitorRegistry::new();
reg.monitor(10, 1);
reg.monitor(20, 1);
reg.monitor(30, 1);
let (_, downs) = reg.process_exited(1, &ExitReason::Timeout);
assert_eq!(downs.len(), 3);
}
#[test]
fn exit_reason_display() {
assert_eq!(ExitReason::Normal.to_string(), "normal");
assert_eq!(ExitReason::Crash("oom".into()).to_string(), "crash: oom");
assert_eq!(ExitReason::Killed.to_string(), "killed");
assert_eq!(ExitReason::Timeout.to_string(), "timeout");
}
#[test]
fn exit_reason_serde_roundtrip() {
let reasons = vec![
ExitReason::Normal,
ExitReason::Crash("test".into()),
ExitReason::Killed,
ExitReason::Timeout,
];
for reason in reasons {
let json = serde_json::to_string(&reason).unwrap();
let restored: ExitReason = serde_json::from_str(&json).unwrap();
assert_eq!(restored, reason);
}
}
#[test]
fn process_down_serde_roundtrip() {
let down = ProcessDown {
pid: 42,
reason: ExitReason::Crash("segfault".into()),
ref_id: "mon-1".into(),
};
let json = serde_json::to_string(&down).unwrap();
let restored: ProcessDown = serde_json::from_str(&json).unwrap();
assert_eq!(restored.pid, 42);
assert_eq!(restored.ref_id, "mon-1");
}
#[test]
fn get_links_empty() {
let reg = MonitorRegistry::new();
assert!(reg.get_links(999).is_empty());
}
#[test]
fn get_monitors_empty() {
let reg = MonitorRegistry::new();
assert!(reg.get_monitors(999).is_empty());
}
}