Skip to content

Commit bd2bbe6

Browse files
drewelezar
authored andcommitted
fix(supervisor): serialize child registration with reaping
Signed-off-by: Drew Newberry <anewberry@nvidia.com> Signed-off-by: Evan Lezar <elezar@nvidia.com>
1 parent 57b08d3 commit bd2bbe6

5 files changed

Lines changed: 121 additions & 45 deletions

File tree

‎crates/openshell-supervisor-process/src/managed_children.rs‎

Lines changed: 71 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -11,25 +11,42 @@
1111
#![cfg(target_os = "linux")]
1212

1313
use std::collections::HashSet;
14-
use std::sync::{LazyLock, Mutex};
14+
use std::sync::{LazyLock, Mutex, MutexGuard, PoisonError};
1515

1616
static MANAGED_CHILDREN: LazyLock<Mutex<HashSet<i32>>> =
1717
LazyLock::new(|| Mutex::new(HashSet::new()));
1818

19-
/// Add `pid` to the supervised-child set. Non-positive or out-of-range values
20-
/// are silently ignored.
21-
pub fn register(pid: u32) {
22-
let Ok(pid) = i32::try_from(pid) else {
19+
fn lock_children() -> MutexGuard<'static, HashSet<i32>> {
20+
MANAGED_CHILDREN
21+
.lock()
22+
.unwrap_or_else(PoisonError::into_inner)
23+
}
24+
25+
fn insert(children: &mut HashSet<i32>, pid: Option<u32>) {
26+
let Some(pid) = pid.and_then(|pid| i32::try_from(pid).ok()) else {
2327
return;
2428
};
25-
if pid <= 0 {
26-
return;
27-
}
28-
if let Ok(mut children) = MANAGED_CHILDREN.lock() {
29+
if pid > 0 {
2930
children.insert(pid);
3031
}
3132
}
3233

34+
/// Spawn a supervised child and register its PID before the orphan reaper can
35+
/// inspect it.
36+
///
37+
/// The registry lock intentionally spans `spawn`. Otherwise a fast child can
38+
/// exit and be reaped after `spawn` returns but before its PID is registered,
39+
/// causing the child's explicit waiter to fail with `ECHILD`.
40+
pub fn spawn_registered<T, E>(
41+
spawn: impl FnOnce() -> Result<T, E>,
42+
pid: impl FnOnce(&T) -> Option<u32>,
43+
) -> Result<T, E> {
44+
let mut children = lock_children();
45+
let child = spawn()?;
46+
insert(&mut children, pid(&child));
47+
Ok(child)
48+
}
49+
3350
/// Remove `pid` from the supervised-child set. Non-positive or out-of-range
3451
/// values are silently ignored.
3552
pub fn unregister(pid: u32) {
@@ -39,32 +56,58 @@ pub fn unregister(pid: u32) {
3956
if pid <= 0 {
4057
return;
4158
}
42-
if let Ok(mut children) = MANAGED_CHILDREN.lock() {
43-
children.remove(&pid);
44-
}
59+
lock_children().remove(&pid);
4560
}
4661

47-
/// Return `true` if `pid` is currently in the supervised-child set.
48-
#[must_use]
49-
pub fn is_managed(pid: i32) -> bool {
50-
MANAGED_CHILDREN
51-
.lock()
52-
.is_ok_and(|children| children.contains(&pid))
62+
/// Run `reap` only when `pid` is not a supervised child.
63+
///
64+
/// The registry lock remains held while `reap` runs so child creation cannot
65+
/// open a spawn-to-registration window between the membership check and reap.
66+
pub fn reap_if_unmanaged<T>(pid: i32, reap: impl FnOnce() -> T) -> Option<T> {
67+
let children = lock_children();
68+
if children.contains(&pid) {
69+
None
70+
} else {
71+
Some(reap())
72+
}
5373
}
5474

5575
#[cfg(test)]
5676
mod tests {
57-
use super::{is_managed, register, unregister};
77+
use super::{reap_if_unmanaged, spawn_registered, unregister};
5878
use nix::sys::wait::{Id, WaitPidFlag, WaitStatus, waitid, waitpid};
5979
use nix::unistd::Pid;
6080
use std::process::Command;
6181

82+
#[test]
83+
fn spawned_child_is_registered_before_returning() {
84+
let pid = 1_000_000_u32;
85+
let child = spawn_registered(|| Ok::<_, ()>(pid), |pid| Some(*pid)).unwrap();
86+
let mut reaped = false;
87+
88+
let result = reap_if_unmanaged(i32::try_from(child).unwrap(), || {
89+
reaped = true;
90+
});
91+
92+
assert!(result.is_none());
93+
assert!(!reaped);
94+
unregister(pid);
95+
}
96+
97+
#[test]
98+
fn unmanaged_child_can_be_reaped() {
99+
let result = reap_if_unmanaged(1_000_001, || "reaped");
100+
101+
assert_eq!(result, Some("reaped"));
102+
}
103+
62104
#[test]
63105
fn fast_child_remains_waitable_after_orphan_reap_attempt() {
64-
let mut child = Command::new("sh")
65-
.args(["-c", "exit 11"])
66-
.spawn()
67-
.expect("spawn fast child");
106+
let mut child = spawn_registered(
107+
|| Command::new("sh").args(["-c", "exit 11"]).spawn(),
108+
|child| Some(child.id()),
109+
)
110+
.expect("spawn and register fast child");
68111
let child_pid = child.id();
69112
let pid = Pid::from_raw(i32::try_from(child_pid).unwrap());
70113

@@ -81,10 +124,11 @@ mod tests {
81124
}
82125
}
83126

84-
if !is_managed(pid.as_raw()) {
85-
waitpid(pid, Some(WaitPidFlag::WNOHANG)).expect("orphan reaper consumes child");
86-
}
87-
register(child_pid);
127+
let reaped = reap_if_unmanaged(pid.as_raw(), || waitpid(pid, Some(WaitPidFlag::WNOHANG)));
128+
assert!(
129+
reaped.is_none(),
130+
"the orphan reaper must leave a registered child to its explicit waiter"
131+
);
88132

89133
let wait = child.wait();
90134
unregister(child_pid);

‎crates/openshell-supervisor-process/src/process.rs‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -901,17 +901,18 @@ impl ProcessHandle {
901901
}
902902

903903
#[cfg(target_os = "linux")]
904-
let mut child = spawn_command_with_supervisor_identity_namespace(cmd)
905-
.into_diagnostic()
906-
.wrap_err("failed to spawn sandbox entrypoint process")?;
904+
let mut child = managed_children::spawn_registered(
905+
|| spawn_command_with_supervisor_identity_namespace(cmd),
906+
Child::id,
907+
)
908+
.into_diagnostic()
909+
.wrap_err("failed to spawn sandbox entrypoint process")?;
907910
#[cfg(not(target_os = "linux"))]
908911
let mut child = cmd
909912
.spawn()
910913
.into_diagnostic()
911914
.wrap_err("failed to spawn sandbox entrypoint process")?;
912915
let pid = child.id().unwrap_or(0);
913-
managed_children::register(pid);
914-
915916
let io = if let Some(master) = pty_master {
916917
ProcessIo::Pty(master)
917918
} else {
@@ -1055,8 +1056,6 @@ impl ProcessHandle {
10551056

10561057
let mut child = cmd.spawn().into_diagnostic()?;
10571058
let pid = child.id().unwrap_or(0);
1058-
#[cfg(target_os = "linux")]
1059-
managed_children::register(pid);
10601059

10611060
debug!(pid, program, "Process spawned");
10621061

‎crates/openshell-supervisor-process/src/run.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -212,12 +212,15 @@ pub async fn run_process(
212212
break;
213213
};
214214

215-
if managed_children::is_managed(pid.as_raw()) {
216-
// Let the explicit waiter own this child status.
215+
let Some(reap_result) = managed_children::reap_if_unmanaged(pid.as_raw(), || {
216+
waitpid(pid, Some(WaitPidFlag::WNOHANG))
217+
}) else {
218+
// Let the explicit waiter own this child status. Stop here
219+
// because WNOWAIT will keep returning this same child.
217220
break;
218-
}
221+
};
219222

220-
match waitpid(pid, Some(WaitPidFlag::WNOHANG)) {
223+
match reap_result {
221224
Ok(WaitStatus::StillAlive)
222225
| Err(nix::errno::Errno::ECHILD | nix::errno::Errno::EINTR) => {}
223226
Ok(reaped) => {

‎crates/openshell-supervisor-process/src/ssh.rs‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1342,13 +1342,14 @@ fn spawn_pty_shell(
13421342
}
13431343

13441344
#[cfg(target_os = "linux")]
1345-
let mut child = crate::process::spawn_std_command_with_supervisor_identity_namespace(cmd)?;
1345+
let mut child = managed_children::spawn_registered(
1346+
|| crate::process::spawn_std_command_with_supervisor_identity_namespace(cmd),
1347+
|child| Some(child.id()),
1348+
)?;
13461349
#[cfg(not(target_os = "linux"))]
13471350
let mut child = cmd.spawn()?;
13481351
#[cfg(target_os = "linux")]
13491352
let child_pid = child.id();
1350-
#[cfg(target_os = "linux")]
1351-
managed_children::register(child_pid);
13521353
let master_file = master;
13531354

13541355
let (sender, receiver) = mpsc::channel::<Vec<u8>>();
@@ -1501,14 +1502,14 @@ fn spawn_pipe_exec(
15011502
}
15021503

15031504
#[cfg(target_os = "linux")]
1504-
let mut child = crate::process::spawn_std_command_with_supervisor_identity_namespace(cmd)?;
1505+
let mut child = managed_children::spawn_registered(
1506+
|| crate::process::spawn_std_command_with_supervisor_identity_namespace(cmd),
1507+
|child| Some(child.id()),
1508+
)?;
15051509
#[cfg(not(target_os = "linux"))]
15061510
let mut child = cmd.spawn()?;
15071511
#[cfg(target_os = "linux")]
15081512
let child_pid = child.id();
1509-
#[cfg(target_os = "linux")]
1510-
managed_children::register(child_pid);
1511-
15121513
let child_stdin = child.stdin.take();
15131514
let child_stdout = child.stdout.take().expect("stdout must be piped");
15141515
let child_stderr = child.stderr.take().expect("stderr must be piped");

‎e2e/rust/tests/sandbox_lifecycle.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,35 @@ async fn canonical_main_nonzero_exit_preserves_status() {
320320
delete_sandbox(&sandbox_name).await;
321321
}
322322

323+
#[tokio::test]
324+
async fn detached_main_exit_during_provisioning_is_classified_as_workload_result() {
325+
let mut sandbox = SandboxGuard::create_detached_main(&["sh", "-c", "exit 11"])
326+
.await
327+
.expect("fast detached main exit should not be reported as a provisioning failure");
328+
329+
let mut get_cmd = openshell_cmd();
330+
get_cmd
331+
.args(["sandbox", "get", &sandbox.name])
332+
.stdout(Stdio::piped())
333+
.stderr(Stdio::piped());
334+
let get_output = get_cmd.output().await.expect("spawn openshell sandbox get");
335+
let details = normalize_output(&format!(
336+
"{}{}",
337+
String::from_utf8_lossy(&get_output.stdout),
338+
String::from_utf8_lossy(&get_output.stderr),
339+
));
340+
assert!(
341+
get_output.status.success(),
342+
"sandbox get failed:\n{details}"
343+
);
344+
assert!(
345+
details.contains("Phase: Error") && details.contains("Exit Code: 11"),
346+
"fast detached main should retain its workload result:\n{details}"
347+
);
348+
349+
sandbox.cleanup().await;
350+
}
351+
323352
#[tokio::test]
324353
async fn canonical_tty_main_uses_sandbox_environment() {
325354
let script = r#"printf 'canonical_env home=%s user=%s term=%s\n' "$HOME" "$USER" "$TERM"; while true; do sleep 1; done"#;

0 commit comments

Comments
 (0)