Refactor process spawn for more granular stdio options

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

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
+        be using tokio::task::spawn directly because our lua
-        scheduler will not drive those futures to completion
+        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
+        NOTE: If an exit code was not given by the child process,
-    let code = status.code().unwrap_or(match stderr.is_empty() {
+        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("stdout", lua.create_string(&res.stdout)?)?
-        .with_value("stderr", lua.create_string(&stderr)?)?
+        .with_value("stderr", lua.create_string(&res.stderr)?)?
         .build_readonly()
 }
 
@@ -201,9 +206,10 @@ async fn spawn_command(
     program: String,
     args: Option<Vec<String>>,
     mut options: ProcessSpawnOptions,
-) -> LuaResult<(ExitStatus, Vec<u8>, Vec<u8>)> {
+) -> LuaResult<WaitForChildResult> {
-    let inherit_stdio = options.inherit_stdio;
+    let stdout = options.stdio.stdout;
-    let stdin = options.stdin.take();
+    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())
+        .stdout(stdout.as_stdio())
-        .stderr(Stdio::piped())
+        .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 {
+    wait_for_child(child, stdout, stderr).await
-        pipe_and_inherit_child_process_stdio(child).await
-    } else {
-        let output = child.wait_with_output().await?;
-        Ok((output.status, output.stdout, output.stderr))
-    }
 }

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<Self, Self::Err> {
+        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<Self> {
+        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()
+                )),
+            }),
+        }
+    }
+}

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(super) struct ProcessSpawnOptions {
-    pub(crate) cwd: Option<PathBuf>,
+    pub cwd: Option<PathBuf>,
-    pub(crate) envs: HashMap<String, String>,
+    pub envs: HashMap<String, String>,
-    pub(crate) shell: Option<String>,
+    pub shell: Option<String>,
-    pub(crate) inherit_stdio: bool,
+    pub stdio: ProcessSpawnOptionsStdio,
-    pub(crate) stdin: Option<Vec<u8>>,
 }
 
 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
+            If we got options for stdio handling, parse those as well - note that
-        */
+            we accept a separate "stdin" value here for compatibility with older
-        match value.get("stdio")? {
+            scripts, but the user should preferrably pass it in the stdio table
-            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
         */
+        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 '{}'",

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<Vec<u8>>,
+}
+
+impl From<ProcessSpawnOptionsStdioKind> 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<Self> {
+        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()
+                )),
+            }),
+        }
+    }
+}

src/lune/builtins/process/pipe_inherit.rs

@@ -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<u8>, Vec<u8>)> {
-    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?))
-}

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<u8>,
+    pub stderr: Vec<u8>,
+}
+
+async fn read_with_stdio_kind<R>(
+    read_from: Option<R>,
+    kind: ProcessSpawnOptionsStdioKind,
+) -> LuaResult<Vec<u8>>
+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<WaitForChildResult> {
+    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,
+    })
+}

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?,
+	stdio: (SpawnOptionsStdioKind | SpawnOptionsStdio)?,
-	stdin: string?,
+	stdin: string?, -- TODO: Remove this since it is now available in stdio above, breaking change
 }
 
 --[=[