refactor: make TrackedChild easier to use (#20)

# refactor: make `TrackedChild` easier to use

This PR refactors the `TrackedChild` struct to make it easier to use by:

- Adding a `ChildTermination` struct that contains both the exit status and path accesses
- Providing a `wait_handle` future that resolves to `ChildTermination` when the process exits
- Exposing stdin, stdout, and stderr handles directly on `TrackedChild`
- Automatically locking the IPC channel when the child process exits
- Updating all usages of `TrackedChild` throughout the codebase

Author: branchseer

crates/fspy/examples/cli.rs

@@ -1,6 +1,6 @@
 use std::{env::args_os, ffi::OsStr, path::PathBuf, pin::Pin};
 
-use fspy::{AccessMode, TrackedChild};
+use fspy::AccessMode;
 use tokio::{
     fs::File,
     io::{AsyncWrite, stdout},
@@ -21,19 +21,16 @@ async fn main() -> anyhow::Result<()> {
     let mut command = spy.new_command(program);
     command.envs(std::env::vars_os()).args(args);
 
-    let TrackedChild { mut tokio_child, accesses_future } = command.spawn().await?;
-
-    let output = tokio_child.wait().await?;
-
-    let accesses = accesses_future.await?;
+    let child = command.spawn().await?;
+    let termination = child.wait_handle.await?;
 
     let mut path_count = 0usize;
     let out_file: Pin<Box<dyn AsyncWrite>> =
         if out_path == "-" { Box::pin(stdout()) } else { Box::pin(File::create(out_path).await?) };
 
     let mut csv_writer = csv_async::AsyncWriter::from_writer(out_file);
 
-    for acc in accesses.iter() {
+    for acc in termination.path_accesses.iter() {
         path_count += 1;
         csv_writer
             .write_record(&[
@@ -49,6 +46,6 @@ async fn main() -> anyhow::Result<()> {
     }
     csv_writer.flush().await?;
 
-    eprintln!("\nfspy: {path_count} paths accessed. {output}");
+    eprintln!("\nfspy: {path_count} paths accessed. status: {}", termination.status);
     Ok(())
 }

crates/fspy/src/lib.rs

@@ -20,20 +20,38 @@ mod os_impl;
 mod arena;
 mod command;
 
-use std::{env::temp_dir, ffi::OsStr, fs::create_dir, io, sync::OnceLock};
+use std::{env::temp_dir, ffi::OsStr, fs::create_dir, io, process::ExitStatus, sync::OnceLock};
 
 pub use command::Command;
 pub use fspy_shared::ipc::{AccessMode, PathAccess};
 use futures_util::future::BoxFuture;
 pub use os_impl::PathAccessIterable;
 use os_impl::SpyInner;
-use tokio::process::Child;
+use tokio::process::{ChildStderr, ChildStdin, ChildStdout};
+
+/// The result of a tracked child process upon its termination.
+pub struct ChildTermination {
+    /// The exit status of the child process.
+    pub status: ExitStatus,
+    /// The path accesses captured from the child process.
+    pub path_accesses: PathAccessIterable,
+}
 
 pub struct TrackedChild {
-    pub tokio_child: Child,
-    /// This future lazily locks the IPC channel when it's polled.
-    /// Do not `await` it until the child process has exited.
-    pub accesses_future: BoxFuture<'static, io::Result<PathAccessIterable>>,
+    /// The handle for writing to the child's standard input (stdin), if it has
+    /// been captured.
+    pub stdin: Option<ChildStdin>,
+
+    /// The handle for reading from the child's standard output (stdout), if it
+    /// has been captured.
+    pub stdout: Option<ChildStdout>,
+
+    /// The handle for reading from the child's standard error (stderr), if it
+    /// has been captured.
+    pub stderr: Option<ChildStderr>,
+
+    /// The future that resolves to exit status and path accesses when the process exits.
+    pub wait_handle: BoxFuture<'static, io::Result<ChildTermination>>,
 }
 
 pub struct Spy(SpyInner);

crates/fspy/src/unix/mod.rs

@@ -22,7 +22,7 @@ use syscall_handler::SyscallHandler;
 use tokio::task::spawn_blocking;
 
 use crate::{
-    Command, TrackedChild,
+    ChildTermination, Command, TrackedChild,
     arena::PathAccessArena,
     error::SpawnError,
     ipc::{OwnedReceiverLockGuard, SHM_CAPACITY},
@@ -131,26 +131,35 @@ pub(crate) async fn spawn_impl(mut command: Command) -> Result<TrackedChild, Spa
     // tokio_command.spawn blocks while executing the `pre_exec` closure.
     // Run it inside spawn_blocking to avoid blocking the tokio runtime, especially the supervisor loop,
     // which needs to accept incoming connections while `pre_exec` is connecting to it.
-    let child = spawn_blocking(move || tokio_command.spawn())
+    let mut child = spawn_blocking(move || tokio_command.spawn())
         .await
         .map_err(|err| SpawnError::OsSpawnError(err.into()))?
         .map_err(SpawnError::OsSpawnError)?;
 
-    let arenas_future = async move {
-        let arenas = std::iter::once(exec_resolve_accesses);
-        #[cfg(target_os = "linux")]
-        let arenas =
-            arenas.chain(supervisor.stop().await?.into_iter().map(|handler| handler.into_arena()));
-        io::Result::Ok(arenas.collect::<Vec<_>>())
-    };
-
-    let accesses_future = async move {
-        let arenas = arenas_future.await?;
-        // `receiver.lock()` blocks. Run it inside `spawn_blocking` to avoid blocking the tokio runtime.
-        let ipc_receiver_lock_guard = OwnedReceiverLockGuard::lock_async(ipc_receiver).await?;
-        Ok(PathAccessIterable { arenas, ipc_receiver_lock_guard })
-    }
-    .boxed();
-
-    Ok(TrackedChild { tokio_child: child, accesses_future })
+    Ok(TrackedChild {
+        stdin: child.stdin.take(),
+        stdout: child.stdout.take(),
+        stderr: child.stderr.take(),
+        // Keep polling for the child to exit in the background even if `wait_handle` is not awaited,
+        // because we need to stop the supervisor and lock the channel as soon as the child exits.
+        wait_handle: tokio::spawn(async move {
+            let status = child.wait().await?;
+
+            let arenas = std::iter::once(exec_resolve_accesses);
+            // Stop the supervisor and collect path accesses from it.
+            #[cfg(target_os = "linux")]
+            let arenas = arenas
+                .chain(supervisor.stop().await?.into_iter().map(|handler| handler.into_arena()));
+            let arenas = arenas.collect::<Vec<_>>();
+
+            // Lock the ipc channel after the child has exited.
+            // We are not interested in path accesses from descendants after the main child has exited.
+            let ipc_receiver_lock_guard = OwnedReceiverLockGuard::lock_async(ipc_receiver).await?;
+            let path_accesses = PathAccessIterable { arenas, ipc_receiver_lock_guard };
+
+            io::Result::Ok(ChildTermination { status, path_accesses })
+        })
+        .map(|f| io::Result::Ok(f??)) // flatten JoinError and io::Result
+        .boxed(),
+    })
 }

crates/fspy/src/windows/mod.rs

@@ -21,7 +21,7 @@ use winsafe::co::{CP, WC};
 use xxhash_rust::const_xxh3::xxh3_128;
 
 use crate::{
-    TrackedChild,
+    ChildTermination, TrackedChild,
     command::Command,
     error::SpawnError,
     fixture::Fixture,
@@ -81,17 +81,9 @@ pub(crate) async fn spawn_impl(command: Command) -> Result<TrackedChild, SpawnEr
     let (channel_conf, receiver) =
         channel(SHM_CAPACITY).map_err(SpawnError::ChannelCreationError)?;
 
-    let accesses_future = async move {
-        let ipc_receiver_lock_guard = OwnedReceiverLockGuard::lock_async(receiver).await?;
-        io::Result::Ok(PathAccessIterable { ipc_receiver_lock_guard })
-    }
-    .boxed();
-
-    // let path_access_stream = PathAccessIterable { pipe_receiver };
-
     let mut spawn_success = false;
     let spawn_success = &mut spawn_success;
-    let child = command
+    let mut child = command
         .spawn_with(|std_command| {
             let std_child = std_command.spawn()?;
             *spawn_success = true;
@@ -138,5 +130,22 @@ pub(crate) async fn spawn_impl(command: Command) -> Result<TrackedChild, SpawnEr
             }
         })?;
 
-    Ok(TrackedChild { tokio_child: child, accesses_future })
+    Ok(TrackedChild {
+        stdin: child.stdin.take(),
+        stdout: child.stdout.take(),
+        stderr: child.stderr.take(),
+        // Keep polling for the child to exit in the background even if `wait_handle` is not awaited,
+        // because we need to stop the supervisor and lock the channel as soon as the child exits.
+        wait_handle: tokio::spawn(async move {
+            let status = child.wait().await?;
+            // Lock the ipc channel after the child has exited.
+            // We are not interested in path accesses from descendants after the main child has exited.
+            let ipc_receiver_lock_guard = OwnedReceiverLockGuard::lock_async(receiver).await?;
+            let path_accesses = PathAccessIterable { ipc_receiver_lock_guard };
+
+            io::Result::Ok(ChildTermination { status, path_accesses })
+        })
+        .map(|f| io::Result::Ok(f??)) // flatten JoinError and io::Result
+        .boxed(),
+    })
 }

crates/fspy/tests/node_fs.rs

@@ -2,7 +2,7 @@ mod test_utils;
 
 use std::env::{current_dir, vars_os};
 
-use fspy::{AccessMode, PathAccessIterable, TrackedChild};
+use fspy::{AccessMode, PathAccessIterable};
 use test_utils::assert_contains;
 
 async fn track_node_script(script: &str) -> anyhow::Result<PathAccessIterable> {
@@ -11,11 +11,10 @@ async fn track_node_script(script: &str) -> anyhow::Result<PathAccessIterable> {
         .arg("-e")
         .envs(vars_os()) // https://github.com/jdx/mise/discussions/5968
         .arg(script);
-    let TrackedChild { mut tokio_child, accesses_future } = command.spawn().await?;
-    let status = tokio_child.wait().await?;
-    let accesses = accesses_future.await?;
-    assert!(status.success());
-    Ok(accesses)
+    let child = command.spawn().await?;
+    let termination = child.wait_handle.await?;
+    assert!(termination.status.success());
+    Ok(termination.path_accesses)
 }
 
 #[tokio::test]

crates/fspy/tests/static_executable.rs

@@ -41,12 +41,12 @@ async fn track_test_bin(args: &[&str], cwd: Option<&str>) -> PathAccessIterable
         cmd.current_dir(cwd);
     };
     cmd.args(args);
-    let mut tracked_child = cmd.spawn().await.unwrap();
+    let tracked_child = cmd.spawn().await.unwrap();
 
-    let output = tracked_child.tokio_child.wait().await.unwrap();
-    assert!(output.success());
+    let termination = tracked_child.wait_handle.await.unwrap();
+    assert!(termination.status.success());
 
-    tracked_child.accesses_future.await.unwrap()
+    termination.path_accesses
 }
 
 #[tokio::test]

crates/fspy/tests/test_utils.rs

@@ -1,6 +1,6 @@
 use std::path::{Path, PathBuf, StripPrefixError};
 
-use fspy::{AccessMode, PathAccessIterable, TrackedChild};
+use fspy::{AccessMode, PathAccessIterable};
 
 #[track_caller]
 pub fn assert_contains(
@@ -55,10 +55,7 @@ macro_rules! track_child {
 pub async fn _spawn_with_id(id: &str) -> anyhow::Result<PathAccessIterable> {
     let mut command = fspy::Spy::global()?.new_command(::std::env::current_exe()?);
     command.arg(id);
-    let TrackedChild { mut tokio_child, accesses_future } = command.spawn().await?;
-
-    let status = tokio_child.wait().await?;
-    let accesses = accesses_future.await?;
-    assert!(status.success());
-    Ok(accesses)
+    let termination = command.spawn().await?.wait_handle.await?;
+    assert!(termination.status.success());
+    Ok(termination.path_accesses)
 }

crates/fspy_e2e/src/main.rs

@@ -9,6 +9,7 @@ use std::{
 
 use fspy::{AccessMode, PathAccess};
 use serde::{Deserialize, Serialize};
+use tokio::io::AsyncReadExt;
 
 #[derive(Serialize, Deserialize)]
 struct Config {
@@ -86,25 +87,30 @@ async fn main() {
             .stderr(Stdio::piped())
             .current_dir(&dir);
 
-        let tracked_child = cmd.spawn().await.unwrap();
+        let mut tracked_child = cmd.spawn().await.unwrap();
 
-        let output = tracked_child.tokio_child.wait_with_output().await.unwrap();
-        let accesses = tracked_child.accesses_future.await.unwrap();
+        let mut stdout_bytes = Vec::<u8>::new();
+        tracked_child.stdout.take().unwrap().read_to_end(&mut stdout_bytes).await.unwrap();
 
-        if !output.status.success() {
+        let mut stderr_bytes = Vec::<u8>::new();
+        tracked_child.stderr.take().unwrap().read_to_end(&mut stderr_bytes).await.unwrap();
+
+        let termination = tracked_child.wait_handle.await.unwrap();
+
+        if !termination.status.success() {
             eprintln!("----- stdout begin -----");
-            stderr().write_all(&output.stdout).unwrap();
+            stderr().write_all(&stdout_bytes).unwrap();
             eprintln!("----- stdout end -----");
             eprintln!("----- stderr begin-----");
-            stderr().write_all(&output.stderr).unwrap();
+            stderr().write_all(&stderr_bytes).unwrap();
             eprintln!("----- stderr end -----");
 
-            eprintln!("Case `{}` failed with status: {}", name, output.status);
+            eprintln!("Case `{}` failed with status: {}", name, termination.status);
             process::exit(1);
         }
 
         let mut collector = AccessCollector::new(dir);
-        for access in accesses.iter() {
+        for access in termination.path_accesses.iter() {
             collector.add(access);
         }
         let snap_file = File::create(manifest_dir.join(format!("snaps/{name}.txt"))).unwrap();