1use std::sync::atomic::{AtomicU64, Ordering};
11
12use dashmap::DashMap;
13use serde::{Deserialize, Serialize};
14
15use crate::process::Pid;
16
17#[non_exhaustive]
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub enum ExitReason {
21 Normal,
23 Crash(String),
25 Killed,
27 Timeout,
29}
30
31impl std::fmt::Display for ExitReason {
32 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33 match self {
34 ExitReason::Normal => write!(f, "normal"),
35 ExitReason::Crash(msg) => write!(f, "crash: {msg}"),
36 ExitReason::Killed => write!(f, "killed"),
37 ExitReason::Timeout => write!(f, "timeout"),
38 }
39 }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
44pub struct ProcessLink {
45 pub pid_a: Pid,
46 pub pid_b: Pid,
47}
48
49#[derive(Debug, Clone)]
51pub struct ProcessMonitor {
52 pub watcher: Pid,
53 pub target: Pid,
54 pub ref_id: String,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
59pub struct ProcessDown {
60 pub pid: Pid,
61 pub reason: ExitReason,
62 pub ref_id: String,
63}
64
65pub struct MonitorRegistry {
70 links: DashMap<Pid, Vec<Pid>>,
72 monitors: DashMap<Pid, Vec<ProcessMonitor>>,
74 next_ref: AtomicU64,
76}
77
78impl MonitorRegistry {
79 pub fn new() -> Self {
81 Self {
82 links: DashMap::new(),
83 monitors: DashMap::new(),
84 next_ref: AtomicU64::new(1),
85 }
86 }
87
88 pub fn link(&self, pid_a: Pid, pid_b: Pid) {
92 self.links
94 .entry(pid_a)
95 .or_default()
96 .value_mut()
97 .retain(|&p| p != pid_b); self.links.entry(pid_a).or_default().push(pid_b);
99
100 self.links
102 .entry(pid_b)
103 .or_default()
104 .value_mut()
105 .retain(|&p| p != pid_a);
106 self.links.entry(pid_b).or_default().push(pid_a);
107 }
108
109 pub fn unlink(&self, pid_a: Pid, pid_b: Pid) {
111 if let Some(mut links) = self.links.get_mut(&pid_a) {
112 links.retain(|&p| p != pid_b);
113 }
114 if let Some(mut links) = self.links.get_mut(&pid_b) {
115 links.retain(|&p| p != pid_a);
116 }
117 }
118
119 pub fn monitor(&self, watcher: Pid, target: Pid) -> String {
123 let ref_id = format!("mon-{}", self.next_ref.fetch_add(1, Ordering::Relaxed));
124 let monitor = ProcessMonitor {
125 watcher,
126 target,
127 ref_id: ref_id.clone(),
128 };
129 self.monitors.entry(target).or_default().push(monitor);
130 ref_id
131 }
132
133 pub fn demonitor(&self, ref_id: &str) {
135 for mut entry in self.monitors.iter_mut() {
138 entry.value_mut().retain(|m| m.ref_id != ref_id);
139 }
140 }
141
142 pub fn process_exited(
146 &self,
147 pid: Pid,
148 reason: &ExitReason,
149 ) -> (Vec<(Pid, ExitReason)>, Vec<ProcessDown>) {
150 let mut link_signals = Vec::new();
151 let mut down_signals = Vec::new();
152
153 if let Some((_, linked_pids)) = self.links.remove(&pid) {
155 for linked_pid in &linked_pids {
156 link_signals.push((*linked_pid, reason.clone()));
157 if let Some(mut reverse) = self.links.get_mut(linked_pid) {
159 reverse.retain(|&p| p != pid);
160 }
161 }
162 }
163
164 if let Some((_, monitors)) = self.monitors.remove(&pid) {
166 for mon in monitors {
167 down_signals.push(ProcessDown {
168 pid,
169 reason: reason.clone(),
170 ref_id: mon.ref_id,
171 });
172 }
173 }
174
175 (link_signals, down_signals)
176 }
177
178 pub fn get_links(&self, pid: Pid) -> Vec<Pid> {
180 self.links
181 .get(&pid)
182 .map(|v| v.value().clone())
183 .unwrap_or_default()
184 }
185
186 pub fn get_monitors(&self, pid: Pid) -> Vec<ProcessMonitor> {
188 self.monitors
189 .get(&pid)
190 .map(|v| v.value().clone())
191 .unwrap_or_default()
192 }
193
194 pub fn is_linked(&self, pid_a: Pid, pid_b: Pid) -> bool {
196 self.links
197 .get(&pid_a)
198 .map(|v| v.contains(&pid_b))
199 .unwrap_or(false)
200 }
201}
202
203impl Default for MonitorRegistry {
204 fn default() -> Self {
205 Self::new()
206 }
207}
208
209#[cfg(test)]
210mod tests {
211 use super::*;
212
213 #[test]
214 fn link_creates_bidirectional() {
215 let reg = MonitorRegistry::new();
216 reg.link(1, 2);
217 assert!(reg.is_linked(1, 2));
218 assert!(reg.is_linked(2, 1));
219 }
220
221 #[test]
222 fn link_idempotent() {
223 let reg = MonitorRegistry::new();
224 reg.link(1, 2);
225 reg.link(1, 2);
226 assert_eq!(reg.get_links(1).len(), 1);
228 assert_eq!(reg.get_links(2).len(), 1);
229 }
230
231 #[test]
232 fn unlink_removes_bidirectional() {
233 let reg = MonitorRegistry::new();
234 reg.link(1, 2);
235 reg.unlink(1, 2);
236 assert!(!reg.is_linked(1, 2));
237 assert!(!reg.is_linked(2, 1));
238 }
239
240 #[test]
241 fn monitor_returns_unique_ref() {
242 let reg = MonitorRegistry::new();
243 let ref1 = reg.monitor(10, 20);
244 let ref2 = reg.monitor(10, 20);
245 assert_ne!(ref1, ref2);
246 }
247
248 #[test]
249 fn demonitor_removes_monitor() {
250 let reg = MonitorRegistry::new();
251 let ref_id = reg.monitor(10, 20);
252 assert_eq!(reg.get_monitors(20).len(), 1);
253 reg.demonitor(&ref_id);
254 assert_eq!(reg.get_monitors(20).len(), 0);
255 }
256
257 #[test]
258 fn process_exited_delivers_link_signals() {
259 let reg = MonitorRegistry::new();
260 reg.link(1, 2);
261 reg.link(1, 3);
262
263 let (links, downs) = reg.process_exited(1, &ExitReason::Crash("panic".into()));
264
265 assert_eq!(links.len(), 2);
266 assert!(downs.is_empty());
267
268 let pids: Vec<Pid> = links.iter().map(|(p, _)| *p).collect();
270 assert!(pids.contains(&2));
271 assert!(pids.contains(&3));
272
273 assert!(!reg.is_linked(2, 1));
275 assert!(!reg.is_linked(3, 1));
276 }
277
278 #[test]
279 fn process_exited_delivers_down_signals() {
280 let reg = MonitorRegistry::new();
281 let ref1 = reg.monitor(10, 1);
282 let ref2 = reg.monitor(20, 1);
283
284 let (links, downs) = reg.process_exited(1, &ExitReason::Normal);
285
286 assert!(links.is_empty());
287 assert_eq!(downs.len(), 2);
288
289 let refs: Vec<&str> = downs.iter().map(|d| d.ref_id.as_str()).collect();
290 assert!(refs.contains(&ref1.as_str()));
291 assert!(refs.contains(&ref2.as_str()));
292 }
293
294 #[test]
295 fn process_exited_delivers_both_links_and_monitors() {
296 let reg = MonitorRegistry::new();
297 reg.link(1, 2);
298 reg.monitor(10, 1);
299
300 let (links, downs) = reg.process_exited(1, &ExitReason::Killed);
301
302 assert_eq!(links.len(), 1);
303 assert_eq!(downs.len(), 1);
304 assert_eq!(links[0].0, 2);
305 assert_eq!(downs[0].reason, ExitReason::Killed);
306 }
307
308 #[test]
309 fn normal_exit_delivers_normal_reason() {
310 let reg = MonitorRegistry::new();
311 reg.link(1, 2);
312 let (links, _) = reg.process_exited(1, &ExitReason::Normal);
313 assert_eq!(links[0].1, ExitReason::Normal);
314 }
315
316 #[test]
317 fn multiple_monitors_on_same_target() {
318 let reg = MonitorRegistry::new();
319 reg.monitor(10, 1);
320 reg.monitor(20, 1);
321 reg.monitor(30, 1);
322
323 let (_, downs) = reg.process_exited(1, &ExitReason::Timeout);
324 assert_eq!(downs.len(), 3);
325 }
326
327 #[test]
328 fn exit_reason_display() {
329 assert_eq!(ExitReason::Normal.to_string(), "normal");
330 assert_eq!(ExitReason::Crash("oom".into()).to_string(), "crash: oom");
331 assert_eq!(ExitReason::Killed.to_string(), "killed");
332 assert_eq!(ExitReason::Timeout.to_string(), "timeout");
333 }
334
335 #[test]
336 fn exit_reason_serde_roundtrip() {
337 let reasons = vec![
338 ExitReason::Normal,
339 ExitReason::Crash("test".into()),
340 ExitReason::Killed,
341 ExitReason::Timeout,
342 ];
343 for reason in reasons {
344 let json = serde_json::to_string(&reason).unwrap();
345 let restored: ExitReason = serde_json::from_str(&json).unwrap();
346 assert_eq!(restored, reason);
347 }
348 }
349
350 #[test]
351 fn process_down_serde_roundtrip() {
352 let down = ProcessDown {
353 pid: 42,
354 reason: ExitReason::Crash("segfault".into()),
355 ref_id: "mon-1".into(),
356 };
357 let json = serde_json::to_string(&down).unwrap();
358 let restored: ProcessDown = serde_json::from_str(&json).unwrap();
359 assert_eq!(restored.pid, 42);
360 assert_eq!(restored.ref_id, "mon-1");
361 }
362
363 #[test]
364 fn get_links_empty() {
365 let reg = MonitorRegistry::new();
366 assert!(reg.get_links(999).is_empty());
367 }
368
369 #[test]
370 fn get_monitors_empty() {
371 let reg = MonitorRegistry::new();
372 assert!(reg.get_monitors(999).is_empty());
373 }
374}