Spawn commands into process groups.

Author: twittner

Cargo.lock

@@ -79,6 +79,7 @@ quickcheck = "1.0.3"
 quick_cache = "0.6.16"
 rand = "0.9"
 rayon = "1.10"
+rustix = { version = "1.1.2", features = ["process"] }
 reqwest = { version = "0.12", features = ["json"] }
 secp256k1 = { version = "0.31.0", features = ["global-context", "hashes", "rand", "serde"] }
 serde = { version = "1", features = ["derive", "rc"] }

Cargo.toml

@@ -5,23 +5,26 @@ edition.workspace = true
 rust-version.workspace = true
 
 [features]
-ports = ["dep:bincode", "dep:tokio-util"]
-run-committee = ["dep:tokio-util"]
-verifier = ["dep:alloy", "dep:timeboost-utils"]
-maker = [
+default = [
+    "dep:alloy",
+    "dep:bincode",
     "dep:bytes",
     "dep:prost",
     "dep:quick_cache",
+    "dep:tokio-util",
+    "dep:tonic",
     "dep:timeboost-config",
     "dep:timeboost-proto",
     "dep:timeboost-utils",
-    "dep:tonic"
+    "ports"
 ]
+ports = ["dep:bincode", "dep:tokio-util"]
 
 [dependencies]
 anyhow = { workspace = true }
 clap = { workspace = true }
 futures = { workspace = true }
+rustix = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
 # optional
@@ -39,12 +42,10 @@ timeboost-utils = { path = "../timeboost-utils", optional = true }
 [[bin]]
 name = "block-maker"
 path = "src/binaries/block-maker.rs"
-required-features = ["maker"]
 
 [[bin]]
 name = "block-verifier"
 path = "src/binaries/block-verifier.rs"
-required-features = ["verifier"]
 
 [[bin]]
 name = "port-alloc"
@@ -58,4 +59,3 @@ path = "src/binaries/run.rs"
 [[bin]]
 name = "run-committee"
 path = "src/binaries/run-committee.rs"
-required-features = ["run-committee"]

test-utils/Cargo.toml

@@ -1,8 +1,8 @@
-use std::{ffi::OsStr, path::PathBuf, time::Duration};
+use std::{ffi::OsStr, path::PathBuf};
 
 use anyhow::{Result, bail};
 use clap::Parser;
-use tokio::{fs::read_dir, process::Command, time::timeout};
+use tokio::{fs::read_dir, process::Command};
 use tokio_util::task::TaskTracker;
 
 #[derive(Parser, Debug)]
@@ -18,9 +18,6 @@ struct Args {
 
     #[clap(long, short, default_value = "/tmp")]
     tmp: PathBuf,
-
-    #[clap(long)]
-    max_duration: Option<u64>,
 }
 
 #[tokio::main]
@@ -68,12 +65,7 @@ async fn main() -> Result<()> {
     }
 
     tasks.close();
-
-    if let Some(d) = args.max_duration {
-        let _ = timeout(Duration::from_secs(d), tasks.wait()).await;
-    } else {
-        tasks.wait().await
-    }
+    tasks.wait().await;
 
     Ok(())
 }

test-utils/src/binaries/run-committee.rs

@@ -1,8 +1,15 @@
+use std::future::pending;
+use std::ops::{Deref, DerefMut};
+use std::time::Duration;
 use std::{ffi::OsStr, process::ExitStatus};
 
 use anyhow::{Result, anyhow, bail};
 use clap::Parser;
+use futures::FutureExt;
+use rustix::process::{Pid, Signal, kill_process_group};
 use tokio::select;
+use tokio::signal::unix::{SignalKind, signal};
+use tokio::time::sleep;
 use tokio::{
     process::{Child, Command},
     task::JoinSet,
@@ -18,6 +25,10 @@ struct Args {
     #[clap(long, short)]
     with: Vec<String>,
 
+    /// Optional timeout in seconds.
+    #[clap(long, short)]
+    timeout: Option<u64>,
+
     /// Main command to execute.
     main: Vec<String>,
 }
@@ -26,8 +37,16 @@ struct Args {
 async fn main() -> Result<()> {
     let args = Args::parse();
 
+    let mut term = signal(SignalKind::terminate())?;
+    let mut intr = signal(SignalKind::interrupt())?;
+
     for exe in &args.exec {
-        let status = command(exe.split_whitespace())?.status().await?;
+        let mut pg = ProcessGroup::spawn(exe.split_whitespace())?;
+        let status = select! {
+            s = pg.wait()   => s?,
+            _ = term.recv() => return Ok(()),
+            _ = intr.recv() => return Ok(()),
+        };
         if !status.success() {
             bail!("{exe:?} failed with {:?}", status.code());
         }
@@ -36,13 +55,19 @@ async fn main() -> Result<()> {
     let mut helpers = JoinSet::<Result<ExitStatus>>::new();
     for w in args.with {
         helpers.spawn(async move {
-            let mut c = spawn_command(w.split_whitespace())?;
-            let status = c.wait().await?;
+            let mut pg = ProcessGroup::spawn(w.split_whitespace())?;
+            let status = pg.wait().await?;
             Ok(status)
         });
     }
 
-    let mut main = spawn_command(args.main)?;
+    let mut main = ProcessGroup::spawn(args.main)?;
+
+    let timeout = if let Some(d) = args.timeout {
+        sleep(Duration::from_secs(d)).boxed()
+    } else {
+        pending().boxed()
+    };
 
     select! {
         status = main.wait() => {
@@ -55,32 +80,55 @@ async fn main() -> Result<()> {
             let status = status??;
             bail!("helper command exited before main with status {:?}", status)
         }
+        _ = term.recv() => {}
+        _ = intr.recv() => {}
+        _ = timeout     => eprintln!("timeout")
     }
 
     Ok(())
 }
 
-fn command<I, S>(it: I) -> Result<Command>
-where
-    I: IntoIterator<Item = S>,
-    S: AsRef<OsStr>,
-{
-    let mut args = it.into_iter();
-    let exe = args
-        .next()
-        .ok_or_else(|| anyhow!("invalid command-line args"))?;
-    let mut cmd = Command::new(exe);
-    for a in args {
-        cmd.arg(a);
+/// Every command is spawned into its own, newly created process group.
+struct ProcessGroup(Child, Pid);
+
+impl ProcessGroup {
+    fn spawn<I, S>(it: I) -> Result<Self>
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<OsStr>,
+    {
+        let mut args = it.into_iter();
+        let exe = args
+            .next()
+            .ok_or_else(|| anyhow!("invalid command-line args"))?;
+        let mut cmd = Command::new(exe);
+        for a in args {
+            cmd.arg(a);
+        }
+        cmd.process_group(0);
+        let child = cmd.spawn()?;
+        let id = child.id().ok_or_else(|| anyhow!("child already exited"))?;
+        let pid = Pid::from_raw(id.try_into()?).ok_or_else(|| anyhow!("invalid pid"))?;
+        Ok(Self(child, pid))
+    }
+}
+
+impl Drop for ProcessGroup {
+    fn drop(&mut self) {
+        let _ = kill_process_group(self.1, Signal::KILL);
     }
-    cmd.kill_on_drop(true);
-    Ok(cmd)
 }
 
-fn spawn_command<I, S>(it: I) -> Result<Child>
-where
-    I: IntoIterator<Item = S>,
-    S: AsRef<OsStr>,
-{
-    command(it)?.spawn().map_err(|e| e.into())
+impl Deref for ProcessGroup {
+    type Target = Child;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for ProcessGroup {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
 }