From e16c28fd400530d54b41fe4fca09941386e354d5 Mon Sep 17 00:00:00 2001 From: Filip Tibell Date: Wed, 11 Oct 2023 16:32:16 -0500 Subject: [PATCH] Refactor process spawn for more granular stdio options --- CHANGELOG.md | 4 + src/lune/builtins/process/mod.rs | 48 +++++------ src/lune/builtins/process/options/kind.rs | 80 +++++++++++++++++++ .../process/{options.rs => options/mod.rs} | 47 ++++------- src/lune/builtins/process/options/stdio.rs | 56 +++++++++++++ src/lune/builtins/process/pipe_inherit.rs | 46 ----------- src/lune/builtins/process/wait_for_child.rs | 74 +++++++++++++++++ types/process.luau | 13 ++- 8 files changed, 263 insertions(+), 105 deletions(-) create mode 100644 src/lune/builtins/process/options/kind.rs rename src/lune/builtins/process/{options.rs => options/mod.rs} (80%) create mode 100644 src/lune/builtins/process/options/stdio.rs delete mode 100644 src/lune/builtins/process/pipe_inherit.rs create mode 100644 src/lune/builtins/process/wait_for_child.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b461f10..0653f15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 part:TestMethod("Hello", "world!") ``` +### Changed + +- Stdio options when using `process.spawn` can now be set with more granularity, allowing stderr & stdout to be disabled individually and completely to improve memory usage when they are not being used. + ## `0.7.8` - October 5th, 2023 ### Added diff --git a/src/lune/builtins/process/mod.rs b/src/lune/builtins/process/mod.rs index 98309c3..154a21f 100644 --- a/src/lune/builtins/process/mod.rs +++ b/src/lune/builtins/process/mod.rs @@ -1,7 +1,7 @@ use std::{ env::{self, consts}, path, - process::{ExitStatus, Stdio}, + process::Stdio, }; use dunce::canonicalize; @@ -13,12 +13,12 @@ use crate::lune::{scheduler::Scheduler, util::TableBuilder}; mod tee_writer; -mod pipe_inherit; -use pipe_inherit::pipe_and_inherit_child_process_stdio; - mod options; use options::ProcessSpawnOptions; +mod wait_for_child; +use wait_for_child::{wait_for_child, WaitForChildResult}; + const PROCESS_EXIT_IMPL_LUA: &str = r#" exit(...) yield() @@ -169,21 +169,26 @@ async fn process_spawn( runtime place it on a different thread if possible / necessary Note that we have to use our scheduler here, we can't - use anything like tokio::task::spawn because our lua - scheduler will not drive those futures to completion + be using tokio::task::spawn directly because our lua + scheduler would not drive those futures to completion */ let sched = lua .app_data_ref::<&Scheduler>() .expect("Lua struct is missing scheduler"); - let (status, stdout, stderr) = sched + let res = sched .spawn(spawn_command(program, args, options)) .await .expect("Failed to receive result of spawned process")?; - // NOTE: If an exit code was not given by the child process, - // we default to 1 if it yielded any error output, otherwise 0 - let code = status.code().unwrap_or(match stderr.is_empty() { + /* + NOTE: If an exit code was not given by the child process, + we default to 1 if it yielded any error output, otherwise 0 + + An exit code may be missing if the process was terminated by + some external signal, which is the only time we use this default + */ + let code = res.status.code().unwrap_or(match res.stderr.is_empty() { true => 0, false => 1, }); @@ -192,8 +197,8 @@ async fn process_spawn( TableBuilder::new(lua)? .with_value("ok", code == 0)? .with_value("code", code)? - .with_value("stdout", lua.create_string(&stdout)?)? - .with_value("stderr", lua.create_string(&stderr)?)? + .with_value("stdout", lua.create_string(&res.stdout)?)? + .with_value("stderr", lua.create_string(&res.stderr)?)? .build_readonly() } @@ -201,9 +206,10 @@ async fn spawn_command( program: String, args: Option>, mut options: ProcessSpawnOptions, -) -> LuaResult<(ExitStatus, Vec, Vec)> { - let inherit_stdio = options.inherit_stdio; - let stdin = options.stdin.take(); +) -> LuaResult { + let stdout = options.stdio.stdout; + let stderr = options.stdio.stderr; + let stdin = options.stdio.stdin.take(); let mut child = options .into_command(program, args) @@ -211,20 +217,14 @@ async fn spawn_command( true => Stdio::piped(), false => Stdio::null(), }) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) + .stdout(stdout.as_stdio()) + .stderr(stderr.as_stdio()) .spawn()?; - // If the stdin option was provided, we write that to the child if let Some(stdin) = stdin { let mut child_stdin = child.stdin.take().unwrap(); child_stdin.write_all(&stdin).await.into_lua_err()?; } - if inherit_stdio { - pipe_and_inherit_child_process_stdio(child).await - } else { - let output = child.wait_with_output().await?; - Ok((output.status, output.stdout, output.stderr)) - } + wait_for_child(child, stdout, stderr).await } diff --git a/src/lune/builtins/process/options/kind.rs b/src/lune/builtins/process/options/kind.rs new file mode 100644 index 0000000..3e0f39c --- /dev/null +++ b/src/lune/builtins/process/options/kind.rs @@ -0,0 +1,80 @@ +use std::{fmt, process::Stdio, str::FromStr}; + +use itertools::Itertools; +use mlua::prelude::*; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum ProcessSpawnOptionsStdioKind { + // TODO: We need better more obvious names + // for these, but that is a breaking change + #[default] + Default, + Forward, + Inherit, + None, +} + +impl ProcessSpawnOptionsStdioKind { + pub fn all() -> &'static [Self] { + &[Self::Default, Self::Forward, Self::Inherit, Self::None] + } + + pub fn as_stdio(self) -> Stdio { + match self { + Self::None => Stdio::null(), + Self::Forward => Stdio::inherit(), + _ => Stdio::piped(), + } + } +} + +impl fmt::Display for ProcessSpawnOptionsStdioKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let s = match *self { + Self::Default => "default", + Self::Forward => "forward", + Self::Inherit => "inherit", + Self::None => "none", + }; + f.write_str(s) + } +} + +impl FromStr for ProcessSpawnOptionsStdioKind { + type Err = LuaError; + fn from_str(s: &str) -> Result { + Ok(match s.trim().to_ascii_lowercase().as_str() { + "default" => Self::Default, + "forward" => Self::Forward, + "inherit" => Self::Inherit, + "none" => Self::None, + _ => { + return Err(LuaError::RuntimeError(format!( + "Invalid spawn options stdio kind - got '{}', expected one of {}", + s, + ProcessSpawnOptionsStdioKind::all() + .iter() + .map(|k| format!("'{k}'")) + .join(", ") + ))) + } + }) + } +} + +impl<'lua> FromLua<'lua> for ProcessSpawnOptionsStdioKind { + fn from_lua(value: LuaValue<'lua>, _: &'lua Lua) -> LuaResult { + match value { + LuaValue::Nil => Ok(Self::default()), + LuaValue::String(s) => s.to_str()?.parse(), + _ => Err(LuaError::FromLuaConversionError { + from: value.type_name(), + to: "ProcessSpawnOptionsStdioKind", + message: Some(format!( + "Invalid spawn options stdio kind - expected string, got {}", + value.type_name() + )), + }), + } + } +} diff --git a/src/lune/builtins/process/options.rs b/src/lune/builtins/process/options/mod.rs similarity index 80% rename from src/lune/builtins/process/options.rs rename to src/lune/builtins/process/options/mod.rs index 4c4fc73..d37f844 100644 --- a/src/lune/builtins/process/options.rs +++ b/src/lune/builtins/process/options/mod.rs @@ -8,13 +8,18 @@ use directories::UserDirs; use mlua::prelude::*; use tokio::process::Command; +mod kind; +mod stdio; + +pub(super) use kind::*; +pub(super) use stdio::*; + #[derive(Debug, Clone, Default)] -pub struct ProcessSpawnOptions { - pub(crate) cwd: Option, - pub(crate) envs: HashMap, - pub(crate) shell: Option, - pub(crate) inherit_stdio: bool, - pub(crate) stdin: Option>, +pub(super) struct ProcessSpawnOptions { + pub cwd: Option, + pub envs: HashMap, + pub shell: Option, + pub stdio: ProcessSpawnOptionsStdio, } impl<'lua> FromLua<'lua> for ProcessSpawnOptions { @@ -112,34 +117,14 @@ impl<'lua> FromLua<'lua> for ProcessSpawnOptions { } /* - If we got options for stdio handling, make sure its one of the constant values - */ - match value.get("stdio")? { - LuaValue::Nil => {} - LuaValue::String(s) => match s.to_str()? { - "inherit" => this.inherit_stdio = true, - "default" => this.inherit_stdio = false, - _ => { - return Err(LuaError::RuntimeError(format!( - "Invalid value for option 'stdio' - expected 'inherit' or 'default', got '{}'", - s.to_string_lossy() - ))) - } - }, - value => { - return Err(LuaError::RuntimeError(format!( - "Invalid type for option 'stdio' - expected 'string', got '{}'", - value.type_name() - ))) - } - } - - /* - If we have stdin contents, we need to pass those to the child process + If we got options for stdio handling, parse those as well - note that + we accept a separate "stdin" value here for compatibility with older + scripts, but the user should preferrably pass it in the stdio table */ + this.stdio = value.get("stdio")?; match value.get("stdin")? { LuaValue::Nil => {} - LuaValue::String(s) => this.stdin = Some(s.as_bytes().to_vec()), + LuaValue::String(s) => this.stdio.stdin = Some(s.as_bytes().to_vec()), value => { return Err(LuaError::RuntimeError(format!( "Invalid type for option 'stdin' - expected 'string', got '{}'", diff --git a/src/lune/builtins/process/options/stdio.rs b/src/lune/builtins/process/options/stdio.rs new file mode 100644 index 0000000..4c12ff4 --- /dev/null +++ b/src/lune/builtins/process/options/stdio.rs @@ -0,0 +1,56 @@ +use mlua::prelude::*; + +use super::kind::ProcessSpawnOptionsStdioKind; + +#[derive(Debug, Clone, Default)] +pub struct ProcessSpawnOptionsStdio { + pub stdout: ProcessSpawnOptionsStdioKind, + pub stderr: ProcessSpawnOptionsStdioKind, + pub stdin: Option>, +} + +impl From for ProcessSpawnOptionsStdio { + fn from(value: ProcessSpawnOptionsStdioKind) -> Self { + Self { + stdout: value, + stderr: value, + ..Default::default() + } + } +} + +impl<'lua> FromLua<'lua> for ProcessSpawnOptionsStdio { + fn from_lua(value: LuaValue<'lua>, lua: &'lua Lua) -> LuaResult { + match value { + LuaValue::Nil => Ok(Self::default()), + LuaValue::String(s) => { + Ok(ProcessSpawnOptionsStdioKind::from_lua(LuaValue::String(s), lua)?.into()) + } + LuaValue::Table(t) => { + let mut this = Self::default(); + + if let Some(stdin) = t.get("stdin")? { + this.stdin = stdin; + } + + if let Some(stdout) = t.get("stdout")? { + this.stdout = stdout; + } + + if let Some(stderr) = t.get("stderr")? { + this.stderr = stderr; + } + + Ok(this) + } + _ => Err(LuaError::FromLuaConversionError { + from: value.type_name(), + to: "ProcessSpawnOptionsStdio", + message: Some(format!( + "Invalid spawn options stdio - expected string or table, got {}", + value.type_name() + )), + }), + } + } +} diff --git a/src/lune/builtins/process/pipe_inherit.rs b/src/lune/builtins/process/pipe_inherit.rs deleted file mode 100644 index 0e4b9a3..0000000 --- a/src/lune/builtins/process/pipe_inherit.rs +++ /dev/null @@ -1,46 +0,0 @@ -use std::process::ExitStatus; - -use mlua::prelude::*; -use tokio::{io, process::Child, task}; - -use super::tee_writer::AsyncTeeWriter; - -pub async fn pipe_and_inherit_child_process_stdio( - mut child: Child, -) -> LuaResult<(ExitStatus, Vec, Vec)> { - let mut child_stdout = child.stdout.take().unwrap(); - let mut child_stderr = child.stderr.take().unwrap(); - - /* - NOTE: We do not need to register these - independent tasks spawning in the scheduler - - This function is only used by `process.spawn` which in - turn registers a task with the scheduler that awaits this - */ - - let stdout_thread = task::spawn(async move { - let mut stdout = io::stdout(); - let mut tee = AsyncTeeWriter::new(&mut stdout); - - io::copy(&mut child_stdout, &mut tee).await.into_lua_err()?; - - Ok::<_, LuaError>(tee.into_vec()) - }); - - let stderr_thread = task::spawn(async move { - let mut stderr = io::stderr(); - let mut tee = AsyncTeeWriter::new(&mut stderr); - - io::copy(&mut child_stderr, &mut tee).await.into_lua_err()?; - - Ok::<_, LuaError>(tee.into_vec()) - }); - - let status = child.wait().await.expect("Child process failed to start"); - - let stdout_buffer = stdout_thread.await.expect("Tee writer for stdout errored"); - let stderr_buffer = stderr_thread.await.expect("Tee writer for stderr errored"); - - Ok::<_, LuaError>((status, stdout_buffer?, stderr_buffer?)) -} diff --git a/src/lune/builtins/process/wait_for_child.rs b/src/lune/builtins/process/wait_for_child.rs new file mode 100644 index 0000000..f126efe --- /dev/null +++ b/src/lune/builtins/process/wait_for_child.rs @@ -0,0 +1,74 @@ +use std::process::ExitStatus; + +use mlua::prelude::*; +use tokio::{ + io::{self, AsyncRead, AsyncReadExt}, + process::Child, + task, +}; + +use super::{options::ProcessSpawnOptionsStdioKind, tee_writer::AsyncTeeWriter}; + +#[derive(Debug, Clone)] +pub(super) struct WaitForChildResult { + pub status: ExitStatus, + pub stdout: Vec, + pub stderr: Vec, +} + +async fn read_with_stdio_kind( + read_from: Option, + kind: ProcessSpawnOptionsStdioKind, +) -> LuaResult> +where + R: AsyncRead + Unpin, +{ + Ok(match kind { + ProcessSpawnOptionsStdioKind::None => Vec::new(), + ProcessSpawnOptionsStdioKind::Forward => Vec::new(), + ProcessSpawnOptionsStdioKind::Default => { + let mut read_from = + read_from.expect("read_from must be Some when stdio kind is Default"); + + let mut buffer = Vec::new(); + + read_from.read_to_end(&mut buffer).await.into_lua_err()?; + + buffer + } + ProcessSpawnOptionsStdioKind::Inherit => { + let mut read_from = + read_from.expect("read_from must be Some when stdio kind is Inherit"); + + let mut stdout = io::stdout(); + let mut tee = AsyncTeeWriter::new(&mut stdout); + + io::copy(&mut read_from, &mut tee).await.into_lua_err()?; + + tee.into_vec() + } + }) +} + +pub(super) async fn wait_for_child( + mut child: Child, + stdout_kind: ProcessSpawnOptionsStdioKind, + stderr_kind: ProcessSpawnOptionsStdioKind, +) -> LuaResult { + let stdout_opt = child.stdout.take(); + let stderr_opt = child.stderr.take(); + + let stdout_task = task::spawn(read_with_stdio_kind(stdout_opt, stdout_kind)); + let stderr_task = task::spawn(read_with_stdio_kind(stderr_opt, stderr_kind)); + + let status = child.wait().await.expect("Child process failed to start"); + + let stdout_buffer = stdout_task.await.into_lua_err()??; + let stderr_buffer = stderr_task.await.into_lua_err()??; + + Ok(WaitForChildResult { + status, + stdout: stdout_buffer, + stderr: stderr_buffer, + }) +} diff --git a/types/process.luau b/types/process.luau index 131b802..7b82052 100644 --- a/types/process.luau +++ b/types/process.luau @@ -1,7 +1,12 @@ export type OS = "linux" | "macos" | "windows" export type Arch = "x86_64" | "aarch64" -export type SpawnOptionsStdio = "inherit" | "default" +export type SpawnOptionsStdioKind = "default" | "inherit" | "forward" | "none" +export type SpawnOptionsStdio = { + stdout: SpawnOptionsStdioKind?, + stderr: SpawnOptionsStdioKind?, + stdin: string?, +} --[=[ @interface SpawnOptions @@ -12,15 +17,15 @@ export type SpawnOptionsStdio = "inherit" | "default" * `cwd` - The current working directory for the process * `env` - Extra environment variables to give to the process * `shell` - Whether to run in a shell or not - set to `true` to run using the default shell, or a string to run using a specific shell - * `stdio` - How to treat output and error streams from the child process - set to "inherit" to pass output and error streams to the current process + * `stdio` - How to treat output and error streams from the child process - see `SpawnOptionsStdioKind` and `SpawnOptionsStdio` for more info * `stdin` - Optional standard input to pass to spawned child process ]=] export type SpawnOptions = { cwd: string?, env: { [string]: string }?, shell: (boolean | string)?, - stdio: SpawnOptionsStdio?, - stdin: string?, + stdio: (SpawnOptionsStdioKind | SpawnOptionsStdio)?, + stdin: string?, -- TODO: Remove this since it is now available in stdio above, breaking change } --[=[