Skip to content

Commit 98a7639

Browse files
committed
refactor(supervisor): make managed child ownership structural
Closes #3154 Signed-off-by: Evan Lezar <elezar@nvidia.com>
1 parent 473d1e9 commit 98a7639

3 files changed

Lines changed: 199 additions & 61 deletions

File tree

‎crates/openshell-sandbox/src/boundary_exec.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,9 @@ impl LocalExecProcess {
384384
runtime: Arc<crate::boundary_io::BoundaryRuntimeState>,
385385
terminal: Arc<std::sync::atomic::AtomicBool>,
386386
signal_lock: Arc<std::sync::Mutex<()>>,
387-
#[cfg(target_os = "linux")] managed_child: Option<crate::managed_children::ManagedChild>,
387+
#[cfg(target_os = "linux")] managed_child: Option<
388+
crate::managed_children::ManagedChildRegistration,
389+
>,
388390
) -> Self {
389391
let result = Arc::new(std::sync::Mutex::new(None));
390392
let exited = Arc::new(tokio::sync::Notify::new());

‎crates/openshell-sandbox/src/managed_children.rs‎

Lines changed: 170 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,22 +8,28 @@
88
//! the orchestrator's `SIGCHLD` reaper can distinguish supervised processes
99
//! from incidental zombies.
1010
11-
#![cfg(target_os = "linux")]
12-
11+
#[cfg(target_os = "linux")]
1312
use std::collections::{HashMap, HashSet};
13+
#[cfg(target_os = "linux")]
1414
use std::io;
15+
#[cfg(target_os = "linux")]
1516
use std::sync::atomic::{AtomicU64, Ordering};
17+
#[cfg(target_os = "linux")]
1618
use std::sync::{LazyLock, Mutex, MutexGuard};
19+
#[cfg(target_os = "linux")]
1720
use std::time::Duration;
1821

22+
#[cfg(target_os = "linux")]
1923
static MANAGED_CHILDREN: LazyLock<Mutex<HashMap<i32, u64>>> =
2024
LazyLock::new(|| Mutex::new(HashMap::new()));
25+
#[cfg(target_os = "linux")]
2126
static NEXT_GENERATION: AtomicU64 = AtomicU64::new(1);
2227

2328
/// Identity of one registry entry. The generation prevents an old waiter from
2429
/// removing a newer child that reused the same numeric PID after reap.
30+
#[cfg(target_os = "linux")]
2531
#[derive(Clone, Copy)]
26-
pub struct ManagedChild {
32+
pub struct ManagedChildRegistration {
2733
pid: i32,
2834
generation: u64,
2935
}
@@ -35,11 +41,13 @@ pub struct ManagedChild {
3541
/// same guard while deciding whether to reap an exited child. This closes the
3642
/// otherwise unavoidable window in which a fast-exiting managed child exists
3743
/// but its PID has not yet been published.
44+
#[cfg(target_os = "linux")]
3845
pub struct RegistryGuard(MutexGuard<'static, HashMap<i32, u64>>);
3946

47+
#[cfg(target_os = "linux")]
4048
impl RegistryGuard {
4149
/// Add a newly spawned managed child.
42-
pub fn register(&mut self, pid: u32) -> Option<ManagedChild> {
50+
pub fn register(&mut self, pid: u32) -> Option<ManagedChildRegistration> {
4351
let Ok(pid) = i32::try_from(pid) else {
4452
return None;
4553
};
@@ -48,7 +56,7 @@ impl RegistryGuard {
4856
}
4957
let generation = NEXT_GENERATION.fetch_add(1, Ordering::Relaxed);
5058
self.0.insert(pid, generation);
51-
Some(ManagedChild { pid, generation })
59+
Some(ManagedChildRegistration { pid, generation })
5260
}
5361

5462
/// Return whether the PID belongs to an explicit waiter.
@@ -60,6 +68,7 @@ impl RegistryGuard {
6068

6169
/// Lock the registry for an atomic spawn-and-register or inspect-and-reap
6270
/// operation.
71+
#[cfg(target_os = "linux")]
6372
pub fn lock() -> RegistryGuard {
6473
RegistryGuard(
6574
MANAGED_CHILDREN
@@ -69,13 +78,15 @@ pub fn lock() -> RegistryGuard {
6978
}
7079

7180
/// Register a child and return the generation-bearing removal token.
72-
pub fn register(pid: u32) -> Option<ManagedChild> {
81+
#[cfg(target_os = "linux")]
82+
pub fn register(pid: u32) -> Option<ManagedChildRegistration> {
7383
lock().register(pid)
7484
}
7585

7686
/// Remove exactly this supervised-child registration. A newer registration
7787
/// for a reused PID is preserved.
78-
pub fn unregister(child: ManagedChild) {
88+
#[cfg(target_os = "linux")]
89+
pub fn unregister(child: ManagedChildRegistration) {
7990
if let Ok(mut children) = MANAGED_CHILDREN.lock()
8091
&& children.get(&child.pid) == Some(&child.generation)
8192
{
@@ -84,15 +95,164 @@ pub fn unregister(child: ManagedChild) {
8495
}
8596

8697
/// Return `true` if `pid` is currently in the supervised-child set.
98+
#[cfg(target_os = "linux")]
8799
#[must_use]
88100
pub fn is_managed(pid: i32) -> bool {
89101
lock().contains(pid)
90102
}
91103

104+
fn valid_pid(pid: Option<u32>) -> Option<u32> {
105+
let pid = pid.and_then(|pid| i32::try_from(pid).ok())?;
106+
(pid > 0).then(|| u32::try_from(pid).expect("positive i32 fits in u32"))
107+
}
108+
109+
/// A child whose exit status is owned by the sandbox runtime.
110+
///
111+
/// Construction holds the Linux registry lock across the real spawn and PID
112+
/// insertion, so the orphan reaper cannot consume a fast child's status
113+
/// between those operations. Other platforms retain the same ownership API
114+
/// without enabling Linux-only reaper bookkeeping.
115+
pub struct ManagedChild<C> {
116+
child: C,
117+
pid: Option<u32>,
118+
#[cfg(target_os = "linux")]
119+
registration: Option<ManagedChildRegistration>,
120+
}
121+
122+
impl<C> ManagedChild<C> {
123+
/// Spawn a child and register its PID before returning it.
124+
pub fn spawn<E>(
125+
spawn: impl FnOnce() -> Result<C, E>,
126+
pid: impl FnOnce(&C) -> Option<u32>,
127+
) -> Result<Self, E> {
128+
#[cfg(target_os = "linux")]
129+
{
130+
let mut registry = lock();
131+
let child = spawn()?;
132+
let pid = valid_pid(pid(&child));
133+
let registration = pid.and_then(|pid| registry.register(pid));
134+
Ok(Self {
135+
child,
136+
pid,
137+
registration,
138+
})
139+
}
140+
141+
#[cfg(not(target_os = "linux"))]
142+
{
143+
let child = spawn()?;
144+
let pid = valid_pid(pid(&child));
145+
Ok(Self { child, pid })
146+
}
147+
}
148+
149+
/// Return the child's PID when it was available at spawn time.
150+
#[must_use]
151+
pub const fn id(&self) -> Option<u32> {
152+
self.pid
153+
}
154+
155+
fn unregister(&mut self) {
156+
#[cfg(target_os = "linux")]
157+
if let Some(registration) = self.registration.take() {
158+
unregister(registration);
159+
}
160+
self.pid = None;
161+
}
162+
}
163+
164+
impl<C> Drop for ManagedChild<C> {
165+
fn drop(&mut self) {
166+
self.unregister();
167+
}
168+
}
169+
170+
fn log_wait_error(pid: Option<u32>, error: &std::io::Error) {
171+
if error.raw_os_error() == Some(libc::ECHILD) {
172+
tracing::error!(
173+
pid = ?pid,
174+
error = %error,
175+
"managed child status was reaped before its explicit waiter"
176+
);
177+
}
178+
}
179+
180+
impl ManagedChild<tokio::process::Child> {
181+
/// Wait for a Tokio child and release its managed PID.
182+
pub async fn wait(&mut self) -> std::io::Result<std::process::ExitStatus> {
183+
let status = self.child.wait().await;
184+
if let Err(error) = &status {
185+
log_wait_error(self.pid, error);
186+
}
187+
self.unregister();
188+
status
189+
}
190+
191+
/// Observe a Tokio child without blocking.
192+
pub fn try_wait(&mut self) -> std::io::Result<Option<std::process::ExitStatus>> {
193+
match self.child.try_wait() {
194+
Ok(status) => {
195+
if status.is_some() {
196+
self.unregister();
197+
}
198+
Ok(status)
199+
}
200+
Err(error) => {
201+
log_wait_error(self.pid, &error);
202+
self.unregister();
203+
Err(error)
204+
}
205+
}
206+
}
207+
208+
/// Take the child's stdin handle without exposing the child itself.
209+
pub fn take_stdin(&mut self) -> Option<tokio::process::ChildStdin> {
210+
self.child.stdin.take()
211+
}
212+
213+
/// Take the child's stdout handle without exposing the child itself.
214+
pub fn take_stdout(&mut self) -> Option<tokio::process::ChildStdout> {
215+
self.child.stdout.take()
216+
}
217+
218+
/// Take the child's stderr handle without exposing the child itself.
219+
pub fn take_stderr(&mut self) -> Option<tokio::process::ChildStderr> {
220+
self.child.stderr.take()
221+
}
222+
}
223+
224+
impl ManagedChild<std::process::Child> {
225+
/// Wait for a standard-library child and release its managed PID.
226+
pub fn wait(&mut self) -> std::io::Result<std::process::ExitStatus> {
227+
let status = self.child.wait();
228+
if let Err(error) = &status {
229+
log_wait_error(self.pid, error);
230+
}
231+
self.unregister();
232+
status
233+
}
234+
235+
/// Take the child's stdin handle without exposing the child itself.
236+
pub fn take_stdin(&mut self) -> Option<std::process::ChildStdin> {
237+
self.child.stdin.take()
238+
}
239+
240+
/// Take the child's stdout handle without exposing the child itself.
241+
pub fn take_stdout(&mut self) -> Option<std::process::ChildStdout> {
242+
self.child.stdout.take()
243+
}
244+
245+
/// Take the child's stderr handle without exposing the child itself.
246+
pub fn take_stderr(&mut self) -> Option<std::process::ChildStderr> {
247+
self.child.stderr.take()
248+
}
249+
}
250+
92251
/// Wait until a managed child is terminal without reaping it.
93252
///
94253
/// Keeping the child as a zombie prevents PID/process-group reuse until the
95254
/// owner publishes terminal state and performs the final wait.
255+
#[cfg(target_os = "linux")]
96256
pub fn wait_until_terminal(pid: u32) -> io::Result<()> {
97257
use nix::sys::wait::{Id, WaitPidFlag, waitid};
98258
let pid = i32::try_from(pid)
@@ -109,6 +269,7 @@ pub fn wait_until_terminal(pid: u32) -> io::Result<()> {
109269
///
110270
/// Explicitly managed children remain owned by their normal waiters. Only
111271
/// unregistered children adopted from the workload process tree are reaped.
272+
#[cfg(target_os = "linux")]
112273
pub fn start_orphan_reaper() -> io::Result<()> {
113274
std::thread::Builder::new()
114275
.name("openshell-orphan-reaper".to_string())
@@ -123,6 +284,7 @@ pub fn start_orphan_reaper() -> io::Result<()> {
123284
.map(|_| ())
124285
}
125286

287+
#[cfg(target_os = "linux")]
126288
fn reap_unmanaged_children_once() -> io::Result<usize> {
127289
use nix::sys::wait::{WaitPidFlag, WaitStatus, waitpid};
128290

@@ -143,6 +305,7 @@ fn reap_unmanaged_children_once() -> io::Result<usize> {
143305
Ok(reaped)
144306
}
145307

308+
#[cfg(target_os = "linux")]
146309
fn direct_child_pids() -> io::Result<HashSet<i32>> {
147310
let mut children = HashSet::new();
148311
for task in std::fs::read_dir("/proc/self/task")? {

0 commit comments

Comments
 (0)