Skip to main content

monoloop_loop/transaction/
owned_process_registry.rs

1//! Runtime-owned ProcessIsolated child registry (D-048).
2//!
3//! Every successfully spawned OS child is retained here (via [`ToolKillHandle`])
4//! until exit is observed. Quiesce may kill and poll, but MUST NOT clear an
5//! entry without reap. `Stopped` requires this set empty.
6
7use super::tool_capacity::ToolPermit;
8use super::tool_handler::ToolKillHandle;
9use std::sync::Mutex;
10
11struct RegistryEntry {
12    kill: ToolKillHandle,
13    /// Capacity held until the child is reaped (DispatchGuard mid-drop path).
14    #[allow(dead_code)]
15    permit: Option<ToolPermit>,
16}
17
18/// ProcessIsolated children owned until OS exit is observed.
19#[derive(Default)]
20pub struct OwnedProcessRegistry {
21    entries: Mutex<Vec<RegistryEntry>>,
22}
23
24impl OwnedProcessRegistry {
25    /// Empty registry.
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    /// Park a live ProcessIsolated kill handle (and optional tool permit).
31    ///
32    /// Caller MUST have already requested kill when transferring from a dropping
33    /// dispatch guard. Entries remain until [`Self::shutdown_progress`] or an
34    /// explicit reap observes exit.
35    pub fn park(&self, kill: ToolKillHandle, permit: Option<ToolPermit>) {
36        debug_assert!(kill.is_process_isolated());
37        let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
38        entries.push(RegistryEntry { kill, permit });
39    }
40
41    /// Kill all live children, then drop entries whose exit has been observed.
42    ///
43    /// Returns the number of children still live after this poll (blocks `Stopped`).
44    pub fn shutdown_progress(&self) -> usize {
45        let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
46        for e in entries.iter() {
47            e.kill.kill();
48        }
49        entries.retain(|e| e.kill.has_join());
50        entries.len()
51    }
52
53    /// Number of ProcessIsolated children not yet observed exited.
54    pub fn live_count(&self) -> usize {
55        let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
56        // Opportunistically drop already-exited children (drive may have reaped).
57        entries.retain(|e| e.kill.has_join());
58        entries.len()
59    }
60
61    /// True when no unreaped ProcessIsolated children remain.
62    pub fn is_empty(&self) -> bool {
63        self.live_count() == 0
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::transaction::process_tool::ProcessIsolatedToolHandler;
71    use crate::transaction::tool_handler::ToolHandler;
72    use monoloop_contracts::{
73        ChannelId, SessionId, SessionKey, ToolCall, ToolCallContext, ToolId, ToolName,
74        TransactionId,
75    };
76    use std::time::{Duration, Instant};
77
78    fn call_ctx() -> (ToolCall, ToolCallContext) {
79        let call = ToolCall {
80            tool_name: ToolName::try_new("sleep").unwrap(),
81            tool_id: ToolId::try_new("sleep").unwrap(),
82            provider_tool_call_id: "p".into(),
83            arguments: serde_json::json!({}),
84            request_ordinal: 0,
85        };
86        let ctx = ToolCallContext {
87            transaction_id: TransactionId::generate(),
88            session_key: SessionKey::new(
89                ChannelId::try_new("llm").unwrap(),
90                SessionId::try_new("s").unwrap(),
91            ),
92            exchange_id: None,
93            tool_action_id: monoloop_contracts::ToolActionId::new("a"),
94            tool_id: ToolId::try_new("sleep").unwrap(),
95            deadline: Instant::now() + Duration::from_secs(5),
96        };
97        (call, ctx)
98    }
99
100    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
101    async fn registry_retains_until_reap_then_empties() {
102        let reg = OwnedProcessRegistry::new();
103        let handler = ProcessIsolatedToolHandler::sleep_until_killed(3600);
104        let (call, ctx) = call_ctx();
105        let handle = handler.start(call, ctx).expect("start");
106        let kill = handle.kill.expect("kill");
107        assert!(kill.has_join());
108        reg.park(kill.clone(), None);
109        assert_eq!(reg.live_count(), 1);
110        assert!(!reg.is_empty());
111        // Quiesce poll: kill + try_wait; sleep child may still be live briefly.
112        let _ = reg.shutdown_progress();
113        // Force wait for exit.
114        kill.join_timeout(Duration::from_secs(2))
115            .await
116            .expect("reaped");
117        assert_eq!(reg.shutdown_progress(), 0);
118        assert!(reg.is_empty());
119    }
120}