Skip to content

refactor: spawn child processes with tokio::process #587

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lvaroqui wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor: spawn child processes with tokio::process #587

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ serde = { workspace = true }
serde_json = { workspace = true, features = ["preserve_order"] }
url = "2.5.8"
sha256 = "1.6"
tokio = { version = "1", features = ["macros", "rt"] }
tokio = { version = "1", features = ["macros", "rt", "process"] }
tokio-tar = { package = "astral-tokio-tar", version = "0.6.2" }
tokio-util = "0.7.18"
crc-fast = { version = "1.9.0", default-features = false, features = ["std"] }
Expand Down
30 changes: 20 additions & 10 deletions src/executor/helpers/run_command_with_log_pipe.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,21 @@ use crate::executor::EXECUTOR_TARGET;
use crate::local_logger::rolling_buffer::ROLLING_BUFFER;
use crate::local_logger::suspend_progress_bar;
use crate::prelude::*;
use std::fs::File;
use std::future::Future;
use std::io::{Read, Write};
use std::process::Command;
use std::process::ExitStatus;
use std::thread;
use tokio::process::Command;

/// Take the child's piped stdout and stderr as blocking files, for reading on plain threads.
pub fn take_blocking_stdio(child: &mut tokio::process::Child) -> Result<(File, File)> {
let stdout = child.stdout.take().context("child stdout is not piped")?;
let stderr = child.stderr.take().context("child stderr is not piped")?;
Ok((
File::from(stdout.into_owned_fd()?),
File::from(stderr.into_owned_fd()?),
))
}

/// Run a command and log its output to stdout and stderr
///
Expand All @@ -23,7 +33,7 @@ pub async fn run_command_with_log_pipe_and_callback<F, Fut>(
cb: F,
) -> Result<ExitStatus>
where
F: FnOnce(std::process::Child) -> Fut,
F: FnOnce(tokio::process::Child) -> Fut,
Fut: Future<Output = anyhow::Result<ExitStatus>>,
{
/// Write text to the rolling buffer if active, otherwise write raw bytes to the writer.
Expand Down Expand Up @@ -85,26 +95,26 @@ where
.stderr(std::process::Stdio::piped())
.spawn()
.context("failed to spawn the process")?;
let stdout = process.stdout.take().expect("unable to get stdout");
let stderr = process.stderr.take().expect("unable to get stderr");
let (stdout, stderr) = take_blocking_stdio(&mut process)?;

let stdout_handle = thread::spawn(move || {
let stdout_handle = tokio::task::spawn_blocking(move || {
log_tee(stdout, std::io::stdout(), None).unwrap();
});

let stderr_handle = thread::spawn(move || {
let stderr_handle = tokio::task::spawn_blocking(move || {
log_tee(stderr, std::io::stderr(), Some("[stderr]")).unwrap();
});

let result = cb(process).await;

// Wait for threads to drain remaining output
let _ = stdout_handle.join();
let _ = stderr_handle.join();
let _ = stdout_handle.await;
let _ = stderr_handle.await;

result
}

pub async fn run_command_with_log_pipe(cmd: Command) -> Result<ExitStatus> {
run_command_with_log_pipe_and_callback(cmd, |mut child| async move { Ok(child.wait()?) }).await
run_command_with_log_pipe_and_callback(cmd, |mut child| async move { Ok(child.wait().await?) })
.await
}
6 changes: 3 additions & 3 deletions src/executor/memory/executor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ impl Executor for MemoryExecutor {
let fifo_data = Rc::new(RefCell::new(None));
let on_process_started = {
let fifo_data_cell = fifo_data.clone();
|mut child: std::process::Child| async move {
|mut child: tokio::process::Child| async move {
let (marker_result, data, exit_status) =
Self::handle_fifo(runner_fifo, ipc, &mut child).await?;
*fifo_data_cell.borrow_mut() = Some(data);
Expand All @@ -194,7 +194,7 @@ impl Executor for MemoryExecutor {
}
};

let status = run_command_with_log_pipe_and_callback(cmd, on_process_started).await?;
let status = run_command_with_log_pipe_and_callback(cmd.into(), on_process_started).await?;
debug!("cmd exit status: {status:?}");

if !status.success() {
Expand Down Expand Up @@ -256,7 +256,7 @@ impl MemoryExecutor {
async fn handle_fifo(
mut runner_fifo: RunnerFifo,
ipc: MemtrackIpcServer,
child: &mut std::process::Child,
child: &mut tokio::process::Child,
) -> anyhow::Result<(
ExecutionTimestamps,
crate::executor::shared::fifo::FifoBenchmarkData,
Expand Down
184 changes: 91 additions & 93 deletions src/executor/shared/fifo.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,13 @@ use runner_shared::artifacts::ExecutionTimestamps;
use runner_shared::fifo::{Command as FifoCommand, MarkerType};
use runner_shared::fifo::{RUNNER_ACK_FIFO, RUNNER_CTL_FIFO};
use std::cmp::Ordering;
use std::collections::HashSet;
use std::os::unix::fs::OpenOptionsExt;
use std::path::{Path, PathBuf};
use std::{collections::HashSet, time::Duration};
use tokio::io::AsyncWriteExt;
use tokio::net::unix::pid_t;
use tokio::net::unix::pipe::Receiver as TokioPipeReader;
use tokio::net::unix::pipe::Sender as TokioPipeSender;
use tokio::time::error::Elapsed;
use tokio_util::codec::{FramedRead, LengthDelimitedCodec};

fn create_fifo<P: AsRef<std::path::Path>>(path: P) -> anyhow::Result<()> {
Expand Down Expand Up @@ -162,7 +161,7 @@ impl RunnerFifo {
/// Returns execution timestamps, benchmark data, and the exit status of the child process.
pub async fn handle_fifo_messages(
&mut self,
child: &mut std::process::Child,
child: &mut tokio::process::Child,
mut handle_cmd: impl AsyncFnMut(&FifoCommand) -> anyhow::Result<Option<FifoCommand>>,
) -> anyhow::Result<(
ExecutionTimestamps,
Expand All @@ -181,107 +180,106 @@ impl RunnerFifo {

let mut benchmark_started = false;

// Outer loop: continues until health check fails
// Process commands until the child process exits
loop {
// Inner loop: process commands until timeout/error
loop {
let result: Result<_, Elapsed> =
tokio::time::timeout(Duration::from_secs(1), self.recv_cmd()).await;
let cmd = match result {
Ok(Ok(cmd)) => cmd,
Ok(Err(e)) => {
warn!("Failed to parse FIFO command: {e}");
break;
}
Err(_) => break, // Timeout
};
trace!("Received command: {cmd:?}");

// Try executor-specific handler first
if let Some(response) = handle_cmd(&cmd).await? {
self.send_cmd(response).await?;
continue;
}
tokio::select! {
result = self.recv_cmd() => {

@not-matthias not-matthias Oct 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We must ensure that we don't accidentally discard events inside the FIFO when the process exits. I suggest we add biased; to the recv_cmd branch to ensure we always finish reading the FIFO first (see: https://docs.rs/tokio/latest/tokio/macro.select.html#fairness)

tokio::select! {
+	// Read commands already in the FIFO before acting on the exit.
+	biased;
	result = self.recv_cmd() => {

let cmd = match result {
Ok(cmd) => cmd,
Err(e) => {
warn!("Failed to parse FIFO command: {e}");
continue;
}
};
trace!("Received command: {cmd:?}");

// Fall through to shared implementation for standard commands
match &cmd {
FifoCommand::CurrentBenchmark { pid, uri } => {
bench_order_by_timestamp.push((get_current_time(), uri.to_string()));
bench_pids.insert(*pid);
self.send_cmd(FifoCommand::Ack).await?;
// Try executor-specific handler first
if let Some(response) = handle_cmd(&cmd).await? {
self.send_cmd(response).await?;
continue;
}
FifoCommand::StartProfiler => {
if !benchmark_started {
benchmark_started = true;
markers.push(MarkerType::SampleStart(get_current_time()));
} else {
warn!("Received duplicate StartProfiler command, ignoring");

// Fall through to shared implementation for standard commands
match &cmd {
FifoCommand::CurrentBenchmark { pid, uri } => {
bench_order_by_timestamp.push((get_current_time(), uri.to_string()));
bench_pids.insert(*pid);
self.send_cmd(FifoCommand::Ack).await?;
}
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::StopProfiler => {
if benchmark_started {
benchmark_started = false;
markers.push(MarkerType::SampleEnd(get_current_time()));
} else {
warn!("Received StopProfiler command before StartProfiler, ignoring");
FifoCommand::StartProfiler => {
if !benchmark_started {
benchmark_started = true;
markers.push(MarkerType::SampleStart(get_current_time()));
} else {
warn!("Received duplicate StartProfiler command, ignoring");
}
self.send_cmd(FifoCommand::Ack).await?;
}
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::SetIntegration { name, version } => {
integration = Some((name.into(), version.into()));
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::AddMarker { marker, .. } => {
markers.push(*marker);
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::SetVersion(protocol_version) => {
match protocol_version.cmp(&runner_shared::fifo::CURRENT_PROTOCOL_VERSION) {
Ordering::Less => {
if *protocol_version
< runner_shared::fifo::MINIMAL_SUPPORTED_PROTOCOL_VERSION
{
bail!(
"Integration is using a version of the protocol that is smaller than the minimal supported protocol version ({protocol_version} < {}). \
Please update the integration to a supported version.",
runner_shared::fifo::MINIMAL_SUPPORTED_PROTOCOL_VERSION
);
}
self.send_cmd(FifoCommand::Ack).await?;
FifoCommand::StopProfiler => {
if benchmark_started {
benchmark_started = false;
markers.push(MarkerType::SampleEnd(get_current_time()));
} else {
warn!("Received StopProfiler command before StartProfiler, ignoring");
}
Ordering::Greater => bail!(
"Runner is using an incompatible protocol version ({} < {protocol_version}). Please update the runner to the latest version.",
runner_shared::fifo::CURRENT_PROTOCOL_VERSION
),
Ordering::Equal => {
self.send_cmd(FifoCommand::Ack).await?;
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::SetIntegration { name, version } => {
integration = Some((name.into(), version.into()));
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::AddMarker { marker, .. } => {
markers.push(*marker);
self.send_cmd(FifoCommand::Ack).await?;
}
FifoCommand::SetVersion(protocol_version) => {
match protocol_version.cmp(&runner_shared::fifo::CURRENT_PROTOCOL_VERSION) {
Ordering::Less => {
if *protocol_version
< runner_shared::fifo::MINIMAL_SUPPORTED_PROTOCOL_VERSION
{
bail!(
"Integration is using a version of the protocol that is smaller than the minimal supported protocol version ({protocol_version} < {}). \
Please update the integration to a supported version.",
runner_shared::fifo::MINIMAL_SUPPORTED_PROTOCOL_VERSION
);
}
self.send_cmd(FifoCommand::Ack).await?;
}
Ordering::Greater => bail!(
"Runner is using an incompatible protocol version ({} < {protocol_version}). Please update the runner to the latest version.",
runner_shared::fifo::CURRENT_PROTOCOL_VERSION
),
Ordering::Equal => {
self.send_cmd(FifoCommand::Ack).await?;
}
}
}
}
_ => {
warn!("Unhandled FIFO command: {cmd:?}");
self.send_cmd(FifoCommand::Err).await?;
_ => {
warn!("Unhandled FIFO command: {cmd:?}");
self.send_cmd(FifoCommand::Err).await?;
}
}
}
}

// Check if the process has exited using try_wait (non-blocking)
match child.try_wait() {
Ok(None) => {} // Still running, continue loop
Ok(Some(exit_status)) => {
debug!(
"Process terminated with status: {exit_status}, stopping the command handler"
);
let marker_result =
ExecutionTimestamps::new(&bench_order_by_timestamp, &markers);
let fifo_data = FifoBenchmarkData {
integration,
bench_pids,
};
return Ok((marker_result, fifo_data, exit_status));
// Returning on exit drops any command not yet read. This relies on the
// protocol acking every command: the sender blocks until it gets a
// reply, so it can't exit with a command still in the FIFO.
result = child.wait() => {
match result {
Ok(exit_status) => {
debug!(
"Process terminated with status: {exit_status}, stopping the command handler"
);
let marker_result = ExecutionTimestamps::new(&bench_order_by_timestamp, &markers);
let fifo_data = FifoBenchmarkData {
integration,
bench_pids,
};
return Ok((marker_result, fifo_data, exit_status));
}
Err(e) => return Err(anyhow::Error::from(e)),
}
}
Err(e) => return Err(anyhow::Error::from(e)),
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/executor/valgrind/build_from_source.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ async fn run_build_command(builder: CommandBuilder) -> Result<()> {
let command_line = builder.as_command_line();
debug!("Running: {command_line}");

let status = run_command_with_log_pipe(builder.build())
let status = run_command_with_log_pipe(builder.build().into())
.await
.with_context(|| format!("failed to run `{command_line}`"))?;

Expand Down
5 changes: 4 additions & 1 deletion src/executor/valgrind/measure.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@ use crate::executor::valgrind::helpers::python::is_free_threaded_python;
use crate::instruments::mongo_tracer::MongoTracer;
use crate::prelude::*;
use log::log_enabled;
use std::env::consts::ARCH;
use std::fs::canonicalize;
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::os::unix::process::ExitStatusExt;
use std::path::Path;
use std::{env::consts::ARCH, process::Command};
use tempfile::TempPath;
use tokio::process::Command;

/// Builds the Valgrind argument list for the given simulation tool.
fn get_valgrind_args(tool: &SimulationTool, config: &ExecutorConfig) -> Vec<String> {
Expand Down Expand Up @@ -232,6 +233,8 @@ mod tests {
use super::*;

fn safe_run(to_execute: &str) -> (u32, u32) {
use std::process::Command;

let script_path = create_run_script().unwrap();
let out_status = tempfile::NamedTempFile::new().unwrap().into_temp_path();

Expand Down
Loading
Loading