refactor: make TrackedChild easier to use

Author: branchseer

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,33 @@ 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(),
+        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 decendants 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(),
+    })
 }