diff --git a/Cargo.lock b/Cargo.lock index ed8ddfd73..43c21f16f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4258,6 +4258,16 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -4758,6 +4768,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index 1cec84dc1..4416ad621 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] } diff --git a/src/executor/helpers/run_command_with_log_pipe.rs b/src/executor/helpers/run_command_with_log_pipe.rs index b20f06834..59c5df7b1 100644 --- a/src/executor/helpers/run_command_with_log_pipe.rs +++ b/src/executor/helpers/run_command_with_log_pipe.rs @@ -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 /// @@ -23,7 +33,7 @@ pub async fn run_command_with_log_pipe_and_callback( cb: F, ) -> Result where - F: FnOnce(std::process::Child) -> Fut, + F: FnOnce(tokio::process::Child) -> Fut, Fut: Future>, { /// Write text to the rolling buffer if active, otherwise write raw bytes to the writer. @@ -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 { - 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 } diff --git a/src/executor/memory/executor.rs b/src/executor/memory/executor.rs index e021f482f..1cfbe4775 100644 --- a/src/executor/memory/executor.rs +++ b/src/executor/memory/executor.rs @@ -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); @@ -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() { @@ -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, diff --git a/src/executor/shared/fifo.rs b/src/executor/shared/fifo.rs index e63e4c658..bdff3cd61 100644 --- a/src/executor/shared/fifo.rs +++ b/src/executor/shared/fifo.rs @@ -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>(path: P) -> anyhow::Result<()> { @@ -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>, ) -> anyhow::Result<( ExecutionTimestamps, @@ -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() => { + 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)), } } } diff --git a/src/executor/valgrind/build_from_source.rs b/src/executor/valgrind/build_from_source.rs index 4622493ce..c53f7a52b 100644 --- a/src/executor/valgrind/build_from_source.rs +++ b/src/executor/valgrind/build_from_source.rs @@ -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}`"))?; diff --git a/src/executor/valgrind/measure.rs b/src/executor/valgrind/measure.rs index 36fbdc0d6..fca89e776 100644 --- a/src/executor/valgrind/measure.rs +++ b/src/executor/valgrind/measure.rs @@ -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 { @@ -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(); diff --git a/src/executor/wall_time/executor.rs b/src/executor/wall_time/executor.rs index a419f7bc3..d6bc1d487 100644 --- a/src/executor/wall_time/executor.rs +++ b/src/executor/wall_time/executor.rs @@ -167,7 +167,7 @@ impl Executor for WallTimeExecutor { } else { cmd_builder }; - let cmd = cmd_builder.build(); + let cmd = cmd_builder.build().into(); debug!("cmd: {cmd:?}"); run_command_with_log_pipe(cmd).await } @@ -217,7 +217,7 @@ async fn run_with_profiler( let wrapped = profiler .wrap_command(cmd_builder, config, profile_folder, requires_sudo) .await?; - let cmd = wrapped.build(); + let cmd = wrapped.build().into(); debug!("cmd: {cmd:?}"); let mut runner_fifo = RunnerFifo::new()?; diff --git a/src/instruments/mongo_tracer.rs b/src/instruments/mongo_tracer.rs index 7f10d21db..afd27c4f9 100644 --- a/src/instruments/mongo_tracer.rs +++ b/src/instruments/mongo_tracer.rs @@ -2,17 +2,21 @@ use std::{ env, io::Read, path::{Path, PathBuf}, - process::{Child, Command, Stdio}, + process::Stdio, str::FromStr, thread, }; use reqwest::Client; -use tokio::fs; +use tokio::{ + fs, + process::{Child, Command}, +}; use url::Url; use crate::binary_pins::PinnedBinary; use crate::cli::run::helpers::{download_pinned_file, get_env_variable}; +use crate::executor::helpers::run_command_with_log_pipe::take_blocking_stdio; use crate::prelude::*; use super::MongoDBConfig; @@ -152,8 +156,7 @@ impl MongoTracer { .stderr(Stdio::piped()) .spawn()?; - let process_stdout = process.stdout.take().expect("error taking child stdout"); - let process_stderr = process.stderr.take().expect("error taking child stderr"); + let (process_stdout, process_stderr) = take_blocking_stdio(&mut process)?; thread::spawn(move || { dump_tracer_log(process_stdout).expect("error communicating with child stdout") }); @@ -178,6 +181,7 @@ impl MongoTracer { if let Some(process) = self.process.as_mut() { process .kill() + .await .context("Failed to kill the MongoDB tracer")?; } @@ -235,6 +239,7 @@ pub async fn install_mongodb_tracer() -> Result<()> { .arg(installer_path.to_str().unwrap()) .stdout(Stdio::piped()) .output() + .await .map_err(|_| anyhow!("Failed to install mongo-tracer"))?; if !output.status.success() { @@ -377,7 +382,7 @@ mod tests { .expect("Failed to apply the transformations"); assert_eq!( - command.get_envs().collect_vec(), + command.as_std().get_envs().collect_vec(), vec![ ( OsStr::new("CODSPEED_MONGO_INSTR_SERVER_ADDRESS"), @@ -414,7 +419,7 @@ mod tests { .expect("Failed to apply the transformations"); assert_eq!( - command.get_envs().collect_vec(), + command.as_std().get_envs().collect_vec(), vec![ ( OsStr::new("CODSPEED_MONGO_INSTR_SERVER_ADDRESS"), @@ -451,7 +456,7 @@ mod tests { .expect("Failed to apply the transformations"); assert_eq!( - command.get_envs().collect_vec(), + command.as_std().get_envs().collect_vec(), vec![ ( OsStr::new("CODSPEED_MONGO_INSTR_SERVER_ADDRESS"), @@ -485,7 +490,7 @@ mod tests { .expect("Failed to apply the transformations"); assert_eq!( - command.get_envs().collect_vec(), + command.as_std().get_envs().collect_vec(), vec![( OsStr::new("CODSPEED_MONGO_INSTR_SERVER_ADDRESS"), Some(OsStr::new("http://0.0.0.0:55581"))