aion/supervision/
tree.rs

1//! Structural supervision tree for engine, workflow-type, workflow, and activity nodes.
2
3use std::collections::{BTreeSet, HashMap};
4use std::sync::{Mutex, MutexGuard};
5
6use crate::{EngineError, Pid};
7
8/// Stable identifier for the single engine supervisor root.
9#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
10pub struct EngineSupervisorId;
11
12/// Stable identifier for a per-workflow-type supervisor.
13#[derive(Clone, Debug, Eq, PartialEq)]
14pub struct TypeSupervisorId {
15    workflow_type: String,
16}
17
18impl TypeSupervisorId {
19    /// Logical workflow type supervised by this node.
20    #[must_use]
21    pub fn workflow_type(&self) -> &str {
22        &self.workflow_type
23    }
24}
25
26/// Snapshot of a per-type supervisor and the workflow processes under it.
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct TypeSupervisorNode {
29    id: TypeSupervisorId,
30    workflow_processes: BTreeSet<Pid>,
31}
32
33impl TypeSupervisorNode {
34    /// Identifier for this per-type supervisor.
35    #[must_use]
36    pub fn id(&self) -> &TypeSupervisorId {
37        &self.id
38    }
39
40    /// Workflow process PIDs directly supervised by this type supervisor.
41    #[must_use]
42    pub fn workflow_processes(&self) -> &BTreeSet<Pid> {
43        &self.workflow_processes
44    }
45}
46
47/// Snapshot of a workflow process and its linked activity children.
48#[derive(Clone, Debug, Eq, PartialEq)]
49pub struct WorkflowNode {
50    workflow_type: String,
51    workflow_pid: Pid,
52    activity_children: BTreeSet<Pid>,
53}
54
55impl WorkflowNode {
56    /// Logical workflow type for this process.
57    #[must_use]
58    pub fn workflow_type(&self) -> &str {
59        &self.workflow_type
60    }
61
62    /// Workflow process PID.
63    #[must_use]
64    pub fn workflow_pid(&self) -> Pid {
65        self.workflow_pid
66    }
67
68    /// Activity child PIDs one level below this workflow process.
69    #[must_use]
70    pub fn activity_children(&self) -> &BTreeSet<Pid> {
71        &self.activity_children
72    }
73
74    /// Returns true when the activity is recorded as a linked child of this workflow.
75    #[must_use]
76    pub fn has_activity_child(&self, activity_pid: Pid) -> bool {
77        self.activity_children.contains(&activity_pid)
78    }
79}
80
81#[derive(Debug, Default)]
82struct TreeState {
83    type_supervisors: HashMap<String, TypeSupervisorNode>,
84    workflows: HashMap<Pid, WorkflowNode>,
85}
86
87/// In-memory model of Aion's three-level supervision structure.
88///
89/// The real cancellation and crash behavior is provided by beamr links through
90/// the runtime boundary. This tree records Aion's intended parent-child shape:
91/// one engine supervisor, one supervisor per distinct workflow type, workflow
92/// processes under their type supervisor, and linked activity children directly
93/// under each workflow process.
94#[derive(Debug, Default)]
95pub struct SupervisionTree {
96    state: Mutex<TreeState>,
97}
98
99impl SupervisionTree {
100    /// Create an empty supervision tree with the single engine root present.
101    #[must_use]
102    pub fn new() -> Self {
103        Self::default()
104    }
105
106    /// Identifier for the single engine supervisor root.
107    #[must_use]
108    pub const fn engine_supervisor(&self) -> EngineSupervisorId {
109        EngineSupervisorId
110    }
111
112    /// Ensure there is exactly one supervisor for `workflow_type`.
113    ///
114    /// # Errors
115    ///
116    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned.
117    pub fn ensure_type_supervisor(
118        &self,
119        workflow_type: impl Into<String>,
120    ) -> Result<TypeSupervisorId, EngineError> {
121        let workflow_type = workflow_type.into();
122        let mut state = self.state()?;
123        let supervisor = state
124            .type_supervisors
125            .entry(workflow_type.clone())
126            .or_insert_with(|| TypeSupervisorNode {
127                id: TypeSupervisorId { workflow_type },
128                workflow_processes: BTreeSet::new(),
129            });
130        Ok(supervisor.id.clone())
131    }
132
133    /// Place a workflow process under its per-type supervisor.
134    ///
135    /// This creates the per-type supervisor on first use and never creates a
136    /// supervisor per workflow execution.
137    ///
138    /// # Errors
139    ///
140    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned.
141    pub fn place_workflow(
142        &self,
143        workflow_type: impl Into<String>,
144        workflow_pid: Pid,
145    ) -> Result<TypeSupervisorId, EngineError> {
146        let workflow_type = workflow_type.into();
147        let mut state = self.state()?;
148
149        if let Some(existing) = state.workflows.get(&workflow_pid) {
150            if existing.workflow_type != workflow_type {
151                return Err(tree_runtime_error(format!(
152                    "workflow process {workflow_pid} is already placed under workflow type `{}`",
153                    existing.workflow_type
154                )));
155            }
156
157            let supervisor = state
158                .type_supervisors
159                .entry(workflow_type.clone())
160                .or_insert_with(|| TypeSupervisorNode {
161                    id: TypeSupervisorId { workflow_type },
162                    workflow_processes: BTreeSet::new(),
163                });
164            supervisor.workflow_processes.insert(workflow_pid);
165            return Ok(supervisor.id.clone());
166        }
167
168        let supervisor = state
169            .type_supervisors
170            .entry(workflow_type.clone())
171            .or_insert_with(|| TypeSupervisorNode {
172                id: TypeSupervisorId {
173                    workflow_type: workflow_type.clone(),
174                },
175                workflow_processes: BTreeSet::new(),
176            });
177        supervisor.workflow_processes.insert(workflow_pid);
178        let id = supervisor.id.clone();
179        state.workflows.insert(
180            workflow_pid,
181            WorkflowNode {
182                workflow_type,
183                workflow_pid,
184                activity_children: BTreeSet::new(),
185            },
186        );
187        Ok(id)
188    }
189
190    /// Record an already-spawned linked activity child under its workflow process.
191    ///
192    /// Call this only after [`crate::RuntimeHandle::spawn_activity`] succeeds;
193    /// that runtime call is the authoritative BEAM link establishment.
194    ///
195    /// # Errors
196    ///
197    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned, or
198    /// [`EngineError::Runtime`] when `workflow_pid` is not present in the tree.
199    pub fn record_activity_child(
200        &self,
201        workflow_pid: Pid,
202        activity_pid: Pid,
203    ) -> Result<(), EngineError> {
204        let mut state = self.state()?;
205        if workflow_pid == activity_pid {
206            return Err(tree_runtime_error(format!(
207                "workflow process {workflow_pid} cannot be its own activity child"
208            )));
209        }
210        if state.workflows.contains_key(&activity_pid) {
211            return Err(tree_runtime_error(format!(
212                "process {activity_pid} is already placed as a workflow process"
213            )));
214        }
215        let workflow = state.workflows.get_mut(&workflow_pid).ok_or_else(|| {
216            tree_runtime_error(format!(
217                "workflow process {workflow_pid} is not in the supervision tree"
218            ))
219        })?;
220        workflow.activity_children.insert(activity_pid);
221        Ok(())
222    }
223
224    /// Number of per-workflow-type supervisors under the engine root.
225    ///
226    /// # Errors
227    ///
228    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned.
229    pub fn type_supervisor_count(&self) -> Result<usize, EngineError> {
230        Ok(self.state()?.type_supervisors.len())
231    }
232
233    /// Snapshot all per-type supervisors without holding the tree lock.
234    ///
235    /// # Errors
236    ///
237    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned.
238    pub fn type_supervisors(&self) -> Result<Vec<TypeSupervisorNode>, EngineError> {
239        Ok(self.state()?.type_supervisors.values().cloned().collect())
240    }
241
242    /// Snapshot a workflow node by PID.
243    ///
244    /// # Errors
245    ///
246    /// Returns [`EngineError::RegistryPoisoned`] if the tree lock was poisoned.
247    pub fn workflow(&self, workflow_pid: Pid) -> Result<Option<WorkflowNode>, EngineError> {
248        Ok(self.state()?.workflows.get(&workflow_pid).cloned())
249    }
250
251    fn state(&self) -> Result<MutexGuard<'_, TreeState>, EngineError> {
252        self.state.lock().map_err(|_| EngineError::RegistryPoisoned)
253    }
254}
255
256fn tree_runtime_error(reason: String) -> EngineError {
257    EngineError::Runtime { reason }
258}
259
260#[cfg(test)]
261mod tests {
262    use std::sync::Arc;
263
264    use crate::EngineError;
265
266    use super::SupervisionTree;
267
268    #[test]
269    fn one_engine_root_has_one_supervisor_per_workflow_type() -> Result<(), EngineError> {
270        let tree = SupervisionTree::new();
271        let root = tree.engine_supervisor();
272        let checkout = tree.ensure_type_supervisor("checkout")?;
273        let billing = tree.ensure_type_supervisor("billing")?;
274        let checkout_again = tree.ensure_type_supervisor("checkout")?;
275
276        assert_eq!(root, tree.engine_supervisor());
277        assert_eq!(checkout.workflow_type(), "checkout");
278        assert_eq!(billing.workflow_type(), "billing");
279        assert_eq!(checkout, checkout_again);
280        assert_eq!(tree.type_supervisor_count()?, 2);
281        Ok(())
282    }
283
284    #[test]
285    fn workflows_sit_under_type_supervisors_not_new_supervisors() -> Result<(), EngineError> {
286        let tree = SupervisionTree::new();
287
288        tree.place_workflow("checkout", 10)?;
289        tree.place_workflow("checkout", 11)?;
290
291        assert_eq!(tree.type_supervisor_count()?, 1);
292        let supervisors = tree.type_supervisors()?;
293        let checkout = supervisors
294            .iter()
295            .find(|node| node.id().workflow_type() == "checkout");
296        assert!(checkout.is_some());
297        if let Some(checkout) = checkout {
298            assert_eq!(checkout.workflow_processes().len(), 2);
299            assert!(checkout.workflow_processes().contains(&10));
300            assert!(checkout.workflow_processes().contains(&11));
301        }
302        Ok(())
303    }
304
305    #[test]
306    fn activity_children_are_one_level_below_workflow_process() -> Result<(), EngineError> {
307        let tree = SupervisionTree::new();
308
309        tree.place_workflow("checkout", 10)?;
310        tree.record_activity_child(10, 20)?;
311        tree.record_activity_child(10, 21)?;
312
313        let workflow = tree.workflow(10)?;
314        assert!(workflow.is_some());
315        if let Some(workflow) = workflow {
316            assert_eq!(workflow.workflow_type(), "checkout");
317            assert_eq!(workflow.workflow_pid(), 10);
318            assert!(workflow.has_activity_child(20));
319            assert!(workflow.has_activity_child(21));
320            assert_eq!(workflow.activity_children().len(), 2);
321        }
322        Ok(())
323    }
324
325    #[test]
326    fn missing_workflow_activity_parent_is_typed_error() {
327        let tree = SupervisionTree::new();
328
329        let error = tree.record_activity_child(99, 20);
330
331        assert!(matches!(error, Err(EngineError::Runtime { .. })));
332    }
333
334    #[test]
335    fn placing_existing_workflow_is_idempotent_and_preserves_children() -> Result<(), EngineError> {
336        let tree = SupervisionTree::new();
337
338        let first = tree.place_workflow("checkout", 10)?;
339        tree.record_activity_child(10, 20)?;
340        let second = tree.place_workflow("checkout", 10)?;
341
342        assert_eq!(first, second);
343        assert_eq!(tree.type_supervisor_count()?, 1);
344        let workflow = tree.workflow(10)?;
345        assert!(workflow.is_some());
346        if let Some(workflow) = workflow {
347            assert!(workflow.has_activity_child(20));
348        }
349        Ok(())
350    }
351
352    #[test]
353    fn workflow_pid_cannot_move_between_type_supervisors() -> Result<(), EngineError> {
354        let tree = SupervisionTree::new();
355
356        tree.place_workflow("checkout", 10)?;
357        let error = tree.place_workflow("billing", 10);
358
359        assert!(matches!(error, Err(EngineError::Runtime { .. })));
360        assert_eq!(tree.type_supervisor_count()?, 1);
361        Ok(())
362    }
363
364    #[test]
365    fn activity_child_cannot_alias_workflow_process() -> Result<(), EngineError> {
366        let tree = SupervisionTree::new();
367
368        tree.place_workflow("checkout", 10)?;
369        assert!(matches!(
370            tree.record_activity_child(10, 10),
371            Err(EngineError::Runtime { .. })
372        ));
373        Ok(())
374    }
375
376    #[test]
377    fn poisoned_tree_lock_returns_typed_registry_error() {
378        let tree = Arc::new(SupervisionTree::new());
379        let poisoner_tree = Arc::clone(&tree);
380        let poisoner = std::thread::spawn(move || {
381            let guard = poisoner_tree.state.lock();
382            assert!(guard.is_ok());
383            std::panic::resume_unwind(Box::new("poison supervision tree lock"));
384        });
385
386        assert!(poisoner.join().is_err());
387        assert!(matches!(
388            tree.type_supervisor_count(),
389            Err(EngineError::RegistryPoisoned)
390        ));
391    }
392}
aion/supervision/tree.rs

aion/supervision/
tree.rs
