From 397ed9e8a8a45e0faaa51713d7ba11e7cc728c9f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fe=CC=81lix=20Saparelli?= Date: Sat, 9 Mar 2024 18:33:06 +1300 Subject: [PATCH] wip: add with_new for ergonomics --- README.md | 22 ++-- src/tokio/core.rs | 21 ++- tests/tokio_unix.rs | 310 ++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 327 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 36f4b75..7a4807e 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ process-wrap = "6.0.0" use std::process::Command; use process_wrap::std::*; -let mut child = StdCommandWrap::new(Command::new("watch").arg("ls")) +let mut child = StdCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessGroup::leader()) .spawn()?; let status = child.wait()?; @@ -51,7 +51,7 @@ dbg!(status); use tokio::process::Command; use process_wrap::tokio::*; -let mut child = TokioCommandWrap::new(Command::new("watch").arg("ls")) +let mut child = TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessGroup::leader()) .spawn()?; let status = child.wait().await?; @@ -64,7 +64,7 @@ dbg!(status); use tokio::process::Command; use process_wrap::tokio::*; -let mut child = TokioCommandWrap::new(Command::new("watch").arg("ls")) +let mut child = TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(JobObject::new()) .spawn()?; let status = child.wait().await?; @@ -77,7 +77,7 @@ dbg!(status); use tokio::process::Command; use process_wrap::tokio::*; -let mut child = TokioCommandWrap::new(Command::new("watch").arg("ls")) +let mut child = TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessSession::leader()) .spawn()?; let status = child.wait().await?; @@ -90,7 +90,7 @@ dbg!(status); use tokio::process::Command; use process_wrap::tokio::*; -let mut child = TokioCommandWrap::new(Command::new("watch").arg("ls")) +let mut child = TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessSession) .wrap(KillOnDrop) .spawn()?; @@ -113,7 +113,7 @@ dbg!(status); - Feature: `process-group` (default) ```rust -TokioCommandWrap::new(Command::new("watch").arg("ls")) +TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessGroup::leader()) .spawn()?; ``` @@ -121,7 +121,7 @@ TokioCommandWrap::new(Command::new("watch").arg("ls")) Or join a different group instead: ```rust -TokioCommandWrap::new(Command::new("watch").arg("ls")) +TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessGroup::attach_to(pgid)) .spawn()?; ``` @@ -138,7 +138,7 @@ This combines creating a new session and a new group, and setting this process a To join the session from another process, use `ProcessGroup::attach_to()` instead. ```rust -TokioCommandWrap::new(Command::new("watch").arg("ls")) +TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(ProcessSession) .spawn()?; ``` @@ -159,7 +159,7 @@ This is a shim to allow setting Windows process creation flags with this API, as Note the `CREATE_SUSPENDED` will always be set, as it is required for the crate to function. ```rust -TokioCommandWrap::new(Command::new("watch").arg("ls")) +TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(CreationFlags::NO_WINDOW | CreationFlags::DETACHED) .wrap(JobObject::new()) .spawn()?; @@ -168,7 +168,7 @@ TokioCommandWrap::new(Command::new("watch").arg("ls")) Or with custom flags: ```rust -TokioCommandWrap::new(Command::new("watch").arg("ls")) +TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(CreationFlags::custom(CREATE_NEW_CONSOLE | CREATE_PROTECTED_PROCESS)) .wrap(JobObject::new()) .spawn()?; @@ -183,7 +183,7 @@ TokioCommandWrap::new(Command::new("watch").arg("ls")) This is a shim to allow wrappers to handle the kill-on-drop flag, as it can't be read from Command. ```rust -let child = TokioCommandWrap::new(Command::new("watch").arg("ls")) +let child = TokioCommandWrap::with_new("watch", |command| { command.arg("ls"); }) .wrap(KillOnDrop) .wrap(ProcessGroup::leader()) .spawn()?; diff --git a/src/tokio/core.rs b/src/tokio/core.rs index b437022..b48aeba 100644 --- a/src/tokio/core.rs +++ b/src/tokio/core.rs @@ -1,5 +1,6 @@ use std::{ any::TypeId, + ffi::OsStr, future::Future, io::{Error, Result}, process::ExitStatus, @@ -20,7 +21,9 @@ pub struct TokioCommandWrap { } impl TokioCommandWrap { - pub fn new(command: Command) -> Self { + pub fn with_new(program: impl AsRef, init: impl FnOnce(&mut Command)) -> Self { + let mut command = Command::new(program); + init(&mut command); Self { command, wrappers: IndexMap::new(), @@ -41,14 +44,13 @@ impl TokioCommandWrap { self } - pub fn spawn(mut self) -> Result> { - let mut command = self.command; + pub fn spawn(&mut self) -> Result> { for (id, wrapper) in &mut self.wrappers { debug!(?id, "pre_spawn"); - wrapper.pre_spawn(&mut command)?; + wrapper.pre_spawn(&mut self.command)?; } - let mut child = command.spawn()?; + let mut child = self.command.spawn()?; for (id, wrapper) in &mut self.wrappers { debug!(?id, "post_spawn"); wrapper.post_spawn(&mut child)?; @@ -64,6 +66,15 @@ impl TokioCommandWrap { } } +impl From for TokioCommandWrap { + fn from(command: Command) -> Self { + Self { + command, + wrappers: IndexMap::new(), + } + } +} + pub trait TokioCommandWrapper: std::fmt::Debug { // process-wrap guarantees that `other` will be of the same type as `self` // note that other crates that may use this trait should guarantee this, but diff --git a/tests/tokio_unix.rs b/tests/tokio_unix.rs index ecc8fe1..1dce420 100644 --- a/tests/tokio_unix.rs +++ b/tests/tokio_unix.rs @@ -2,22 +2,30 @@ use std::{ io::Result, + // os::unix::process::ExitStatusExt, process::Stdio, + // time::Duration, }; use process_wrap::tokio::*; use tokio::{ - io::AsyncReadExt, - process::Command, + io::{ + AsyncReadExt, + // AsyncWriteExt + }, + // time::sleep, }; +// const DIE_TIME: Duration = Duration::from_millis(100); + // each test has a _nowrap variant that uses the process-wrap API but doesn't apply any Wrappers for comparison/debugging. #[tokio::test] async fn inner_read_stdout_nowrap() -> Result<()> { - let mut command = Command::new("echo"); - command.arg("hello").stdout(Stdio::piped()); - let mut child = TokioCommandWrap::new(command).spawn()?; + let mut child = TokioCommandWrap::with_new("echo", |command| { + command.arg("hello").stdout(Stdio::piped()); + }) + .spawn()?; let mut output = String::new(); if let Some(mut out) = child.stdout().take() { @@ -30,12 +38,29 @@ async fn inner_read_stdout_nowrap() -> Result<()> { #[tokio::test] async fn inner_read_stdout_process_group() -> Result<()> { - let mut command = Command::new("echo"); - command.arg("hello").stdout(Stdio::piped()); - let mut command = TokioCommandWrap::new(command); - command.wrap(ProcessGroup::leader()); + let mut child = TokioCommandWrap::with_new("echo", |command| { + command.arg("hello").stdout(Stdio::piped()); + }) + .wrap(ProcessGroup::leader()) + .spawn()?; + + let mut output = String::new(); + if let Some(mut out) = child.stdout().take() { + out.read_to_string(&mut output).await?; + } + + assert_eq!(output.as_str(), "hello\n"); + Ok(()) +} + +#[tokio::test] +async fn inner_read_stdout_process_session() -> Result<()> { + let mut child = TokioCommandWrap::with_new("echo", |command| { + command.arg("hello").stdout(Stdio::piped()); + }) + .wrap(ProcessSession) + .spawn()?; - let mut child = command.spawn()?; let mut output = String::new(); if let Some(mut out) = child.stdout().take() { out.read_to_string(&mut output).await?; @@ -44,3 +69,268 @@ async fn inner_read_stdout_process_group() -> Result<()> { assert_eq!(output.as_str(), "hello\n"); Ok(()) } + +/* +#[tokio::test] +async fn into_inner_write_stdin_nowrap() -> Result<()> { + let mut child = Command::new("cat") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn()?; + + if let Some(mut din) = child.stdin.take() { + din.write_all(b"hello").await?; + } + + let mut output = String::new(); + if let Some(mut out) = child.stdout.take() { + out.read_to_string(&mut output).await?; + } + + assert_eq!(output.as_str(), "hello"); + Ok(()) +} + +#[tokio::test] +async fn into_inner_write_stdin_group() -> Result<()> { + let mut child = Command::new("cat") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .group_spawn()? + .into_inner(); + + if let Some(mut din) = child.stdin.take() { + din.write_all(b"hello").await?; + } + + let mut output = String::new(); + if let Some(mut out) = child.stdout.take() { + out.read_to_string(&mut output).await?; + } + + assert_eq!(output.as_str(), "hello"); + Ok(()) +} + +#[tokio::test] +async fn kill_and_try_wait_nowrap() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).spawn()?; + assert!(child.try_wait()?.is_none()); + child.kill().await?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some()); + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some()); + Ok(()) +} + +#[tokio::test] +async fn kill_and_try_wait_group() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).group_spawn()?; + assert!(child.try_wait()?.is_none()); + child.kill().await?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some()); + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some()); + Ok(()) +} + +#[tokio::test] +async fn try_wait_twice_after_sigterm_nowrap() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).spawn()?; + assert!(child.try_wait()?.is_none(), "pre try_wait"); + child.signal(Signal::SIGTERM)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "first try_wait"); + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "second try_wait"); + Ok(()) +} + +#[tokio::test] +async fn try_wait_twice_after_sigterm_group() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).group_spawn()?; + assert!(child.try_wait()?.is_none(), "pre try_wait"); + child.signal(Signal::SIGTERM)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "first try_wait"); + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "second try_wait"); + Ok(()) +} + +#[tokio::test] +async fn wait_twice_after_sigterm_nowrap() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).spawn()?; + assert!(child.try_wait()?.is_none(), "pre try_wait"); + child.signal(Signal::SIGTERM)?; + let status = child.wait().await?; + assert_eq!( + status.signal(), + Some(Signal::SIGTERM as i32), + "first wait status" + ); + let status = child.wait().await?; + assert_eq!( + status.signal(), + Some(Signal::SIGTERM as i32), + "second wait status" + ); + Ok(()) +} + +#[tokio::test] +async fn wait_twice_after_sigterm_group() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).group_spawn()?; + assert!(child.try_wait()?.is_none(), "pre try_wait"); + child.signal(Signal::SIGTERM)?; + let status = child.wait().await?; + assert_eq!( + status.signal(), + Some(Signal::SIGTERM as i32), + "first wait status" + ); + let status = child.wait().await?; + assert_eq!( + status.signal(), + Some(Signal::SIGTERM as i32), + "second wait status" + ); + Ok(()) +} + +#[tokio::test] +async fn wait_after_die_nowrap() -> Result<()> { + let mut child = Command::new("echo").stdout(Stdio::null()).spawn()?; + sleep(DIE_TIME).await; + + let status = child.wait().await?; + assert!(status.success()); + + Ok(()) +} + +#[tokio::test] +async fn wait_after_die_group() -> Result<()> { + let mut child = Command::new("echo").stdout(Stdio::null()).group_spawn()?; + sleep(DIE_TIME).await; + + let status = child.wait().await?; + assert!(status.success()); + + Ok(()) +} + +#[tokio::test] +async fn try_wait_after_die_nowrap() -> Result<()> { + let mut child = Command::new("echo").stdout(Stdio::null()).spawn()?; + sleep(DIE_TIME).await; + + let status = child.try_wait()?; + assert!(status.is_some()); + assert!(status.unwrap().success()); + + Ok(()) +} + +#[tokio::test] +async fn try_wait_after_die_group() -> Result<()> { + let mut child = Command::new("echo").stdout(Stdio::null()).group_spawn()?; + sleep(DIE_TIME).await; + + let status = child.try_wait()?; + assert!(status.is_some()); + assert!(status.unwrap().success()); + + Ok(()) +} + +#[tokio::test] +async fn wait_nowrap() -> Result<()> { + let mut command = Command::new("echo"); + let mut child = command.spawn()?; + let status = child.wait().await?; + assert!(status.success()); + let status = child.wait().await?; + assert!(status.success()); + Ok(()) +} + +#[tokio::test] +async fn wait_group() -> Result<()> { + let mut command = Command::new("echo"); + let mut child = command.group_spawn()?; + let status = child.wait().await?; + assert!(status.success()); + let status = child.wait().await?; + assert!(status.success()); + Ok(()) +} + +#[tokio::test] +async fn wait_with_output_nowrap() -> Result<()> { + let child = Command::new("echo") + .arg("hello") + .stdout(Stdio::piped()) + .spawn()?; + + let output = child.wait_with_output().await?; + assert!(output.status.success()); + assert_eq!(output.stdout, b"hello\n".to_vec()); + assert_eq!(output.stderr, Vec::new()); + Ok(()) +} + +#[tokio::test] +async fn wait_with_output_group() -> Result<()> { + let child = Command::new("echo") + .arg("hello") + .stdout(Stdio::piped()) + .group_spawn()?; + + let output = child.wait_with_output().await?; + assert!(output.status.success()); + assert_eq!(output.stdout, b"hello\n".to_vec()); + assert_eq!(output.stderr, Vec::new()); + Ok(()) +} + +#[tokio::test] +async fn id_same_as_inner_group() -> Result<()> { + let mut command = Command::new("echo"); + let mut child = command.group_spawn()?; + assert_eq!(child.id(), child.inner().id()); + Ok(()) +} + +#[tokio::test] +async fn signal_nowrap() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).spawn()?; + + child.signal(Signal::SIGCONT)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_none(), "not exited with sigcont"); + + child.signal(Signal::SIGTERM)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "exited with sigterm"); + + Ok(()) +} + +#[tokio::test] +async fn signal_group() -> Result<()> { + let mut child = Command::new("yes").stdout(Stdio::null()).group_spawn()?; + + child.signal(Signal::SIGCONT)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_none(), "not exited with sigcont"); + + child.signal(Signal::SIGTERM)?; + sleep(DIE_TIME).await; + assert!(child.try_wait()?.is_some(), "exited with sigterm"); + + Ok(()) +} +*/