Process piping (#200)

Resolve #187 

This PR adds a way to pipe processes concisely and conveniently.
Previously, creating pipelines behaving similarly to Unix's was
impossible (or hard). `Broken pipe` error was handled by the parent
process, so pipeline like `yes | head -n 5` would run forever, and the
app itself crash/throw error. Additionally, there was no way to pass
arguments like `pipefail`.

All this is done with custom handler threads and not the JVM 9 API.
Java's API is also incapable of handling `Broken pipe`.

On Windows, the broken pipe handling is not supported - there is no
`broken pipe` error.

Pull request: #200

Author: szymon-rd

.github/workflows/build.yml

@@ -25,7 +25,6 @@ jobs:
         with:
           distribution: 'temurin'
           java-version: ${{ matrix.java-version }}
-
       - name: Fetch millw launcher (Windows)
         run: curl -Lo mill.bat "https://raw.githubusercontent.com/lefou/millw/main/millw.bat"
         if: matrix.os == 'windows-latest'

Readme.adoc

@@ -1654,6 +1654,48 @@ val sha = os.proc("shasum", "-a", "256").spawn(stdin = gzip.stdout)
 sha.stdout.trim ==> "acc142175fa520a1cb2be5b97cbbe9bea092e8bba3fe2e95afa645615908229e  -"
 ----
 
+== Spawning Pipelines of Subprocesses
+
+After constructing a subprocess with `os.proc`, you can use the `pipeTo` method
+to pipe its output to another subprocess:
+
+[source,scala]
+----
+val wc = os.proc("ls", "-l")
+  .pipeTo(os.proc("wc", "-l"))
+  .call()
+  .out.text()
+----
+
+This is equivalent to the shell command `ls -l | wc -l`. You can chain together
+as many subprocesses as you like. Note that by using this API you can utilize
+the broken pipe behaviour of Unix systems. For example, you can take 10 first elements
+of output from the `yes` command, and after the `head` command terminates, the `yes`
+command will be terminated as well:
+
+[source,scala]
+----
+val yes10 = os.proc("yes")
+  .pipeTo(os.proc("head", "-n", "10"))
+  .call()
+  .out.text()
+----
+
+This feature is implemented inside the library and will terminate any process reading the
+stdin of other process in pipeline on every IO error. This behavior can be disabled via the 
+`handleBrokenPipe` flag on `call` and `spawn` methods. Note that Windows does not support 
+broken pipe behaviour, so a command like`yes` would run forever. `handleBrokenPipe` is set 
+to false by default on Windows.
+
+Both `call` and `spawn` correspond in their behavior to their counterparts in the `os.proc`,
+but `spawn` returns the `os.ProcessPipeline` instance instead. It offers the same 
+`API` as `SubProcess`, but will operate on the set of processes instead of a single one.
+
+`Pipefail` is enabled by default, so if any of the processes in the pipeline fails, the whole
+pipeline will have a non-zero exit code. This behavior can be disabled via the `pipefail` flag
+on `call` and `spawn` methods. Note that the pipefail does not kill the processes in the pipeline,
+it just sets the exit code of the pipeline to the exit code of the failed process.
+
 === Watching for Changes
 
 ==== `os.watch.watch`

build.sc

@@ -58,6 +58,10 @@ trait MiMaChecks extends Mima {
   )
 }
 
+object testJarWriter extends JavaModule
+object testJarReader extends JavaModule
+object testJarExit extends JavaModule
+
 trait OsLibModule
     extends CrossScalaModule
     with PublishModule
@@ -84,7 +88,12 @@ trait OsLibModule
     def ivyDeps = Agg(Deps.utest, Deps.sourcecode)
 
     // we check the textual output of system commands and expect it in english
-    def forkEnv = super.forkEnv() ++ Map("LC_ALL" -> "C")
+    def forkEnv = super.forkEnv() ++ Map(
+      "LC_ALL" -> "C",
+      "TEST_JAR_WRITER_ASSEMBLY" -> testJarWriter.assembly().path.toString,
+      "TEST_JAR_READER_ASSEMBLY" -> testJarReader.assembly().path.toString,
+      "TEST_JAR_EXIT_ASSEMBLY" -> testJarExit.assembly().path.toString
+      )
   }
 }
 

os/src-jvm/ProcessOps.scala

@@ -1,8 +1,14 @@
 package os
 
 import java.util.concurrent.{ArrayBlockingQueue, Semaphore, TimeUnit}
-
+import collection.JavaConverters._
 import scala.annotation.tailrec
+import java.lang.ProcessBuilder.Redirect
+import os.SubProcess.InputStream
+import java.io.IOException
+import java.util.concurrent.LinkedBlockingQueue
+import ProcessOps._
+import scala.util.Try
 
 /**
  * Convenience APIs around [[java.lang.Process]] and [[java.lang.ProcessBuilder]]:
@@ -21,6 +27,7 @@ import scala.annotation.tailrec
  *   the standard stdin/stdout/stderr streams, using whatever protocol you
  *   want
  */
+
 case class proc(command: Shellable*) {
   def commandChunks: Seq[String] = command.flatMap(_.value)
 
@@ -79,7 +86,6 @@ case class proc(command: Shellable*) {
       mergeErrIntoOut,
       propagateEnv
     )
-    import collection.JavaConverters._
 
     sub.join(timeout)
 
@@ -94,9 +100,6 @@ case class proc(command: Shellable*) {
    * and starts a subprocess, and returns it as a `java.lang.Process` for you to
    * interact with however you like.
    *
-   * To implement pipes, you can spawn a process, take it's stdout, and pass it
-   * as the stdin of a second spawned process.
-   *
    * Note that if you provide `ProcessOutput` callbacks to `stdout`/`stderr`,
    * the calls to those callbacks take place on newly spawned threads that
    * execute in parallel with the main thread. Thus make sure any data
@@ -111,28 +114,13 @@ case class proc(command: Shellable*) {
       mergeErrIntoOut: Boolean = false,
       propagateEnv: Boolean = true
   ): SubProcess = {
-    val builder = new java.lang.ProcessBuilder()
-
-    val baseEnv =
-      if (propagateEnv) sys.env
-      else Map()
-    for ((k, v) <- baseEnv ++ Option(env).getOrElse(Map())) {
-      if (v != null) builder.environment().put(k, v)
-      else builder.environment().remove(k)
-    }
-
-    builder.directory(Option(cwd).getOrElse(os.pwd).toIO)
+    val builder =
+      buildProcess(commandChunks, cwd, env, stdin, stdout, stderr, mergeErrIntoOut, propagateEnv)
 
     val cmdChunks = commandChunks
     val commandStr = cmdChunks.mkString(" ")
     lazy val proc: SubProcess = new SubProcess(
-      builder
-        .command(cmdChunks: _*)
-        .redirectInput(stdin.redirectFrom)
-        .redirectOutput(stdout.redirectTo)
-        .redirectError(stderr.redirectTo)
-        .redirectErrorStream(mergeErrIntoOut)
-        .start(),
+      builder.start(),
       stdin.processInput(proc.stdin).map(new Thread(_, commandStr + " stdin thread")),
       stdout.processOutput(proc.stdout).map(new Thread(_, commandStr + " stdout thread")),
       stderr.processOutput(proc.stderr).map(new Thread(_, commandStr + " stderr thread"))
@@ -143,4 +131,228 @@ case class proc(command: Shellable*) {
     proc.errorPumperThread.foreach(_.start())
     proc
   }
+
+  /**
+   * Pipes the output of this process into the input of the [[next]] process. Returns a
+   * [[ProcGroup]] containing both processes, which you can then either execute or
+   * pipe further.
+   */
+  def pipeTo(next: proc): ProcGroup = ProcGroup(Seq(this, next))
+}
+
+/**
+ * A group of processes that are piped together, corresponding to e.g. `ls -l | grep .scala`.
+ * You can create a `ProcGroup` by calling `.pipeTo` on a [[proc]] multiple times.
+ * Contains methods corresponding to the methods on [[proc]], but defined for pipelines
+ * of processes.
+ */
+case class ProcGroup private[os] (commands: Seq[proc]) {
+  assert(commands.size >= 2)
+
+  private lazy val isWindows = sys.props("os.name").toLowerCase().contains("windows")
+
+  /**
+   * Invokes the given pipeline like a function, passing in input and returning a
+   * [[CommandResult]]. You can then call `result.exitCode` to see how it exited, or
+   * `result.out.bytes` or `result.err.string` to access the aggregated stdout and
+   * stderr of the subprocess in a number of convenient ways. If a non-zero exit code
+   * is returned, this throws a [[os.SubprocessException]] containing the
+   * [[CommandResult]], unless you pass in `check = false`.
+   *
+   * For each process in pipeline, the output will be forwarded to the input of the next
+   * process. Input of the first process is set to provided [[stdin]] The output of the last
+   * process will be returned as the output of the pipeline. [[stderr]] is set for all processes.
+   *
+   * `call` provides a number of parameters that let you configure how the pipeline
+   * is run:
+   *
+   * @param cwd              the working directory of the pipeline
+   * @param env              any additional environment variables you wish to set in the pipeline
+   * @param stdin            any data you wish to pass to the pipelines's standard input (to the first process)
+   * @param stdout           How the pipelines's output stream is configured (the last process stdout)
+   * @param stderr           How the process's error stream is configured (set for all processes)
+   * @param mergeErrIntoOut  merges the pipeline's stderr stream into it's stdout. Note that then the
+   *                         stderr will be forwarded with stdout to subsequent processes in the pipeline.
+   * @param timeout          how long to wait in milliseconds for the pipeline to complete
+   * @param check            disable this to avoid throwing an exception if the pipeline
+   *                         fails with a non-zero exit code
+   * @param propagateEnv     disable this to avoid passing in this parent process's
+   *                         environment variables to the pipeline
+   * @param pipefail         if true, the pipeline's exitCode will be the exit code of the first
+   *                         failing process. If no process fails, the exit code will be 0.
+   * @param handleBrokenPipe if true, every [[java.io.IOException]] when redirecting output of a process
+   *                         will be caught and handled by killing the writing process. This behaviour
+   *                         is consistent with handlers of SIGPIPE signals in most programs
+   *                         supporting interruptable piping. Disabled by default on Windows.
+   */
+  def call(
+      cwd: Path = null,
+      env: Map[String, String] = null,
+      stdin: ProcessInput = Pipe,
+      stdout: ProcessOutput = Pipe,
+      stderr: ProcessOutput = os.Inherit,
+      mergeErrIntoOut: Boolean = false,
+      timeout: Long = -1,
+      check: Boolean = true,
+      propagateEnv: Boolean = true,
+      pipefail: Boolean = true,
+      handleBrokenPipe: Boolean = !isWindows
+  ): CommandResult = {
+    val chunks = new java.util.concurrent.ConcurrentLinkedQueue[Either[geny.Bytes, geny.Bytes]]
+
+    val sub = spawn(
+      cwd,
+      env,
+      stdin,
+      if (stdout ne os.Pipe) stdout
+      else os.ProcessOutput.ReadBytes((buf, n) =>
+        chunks.add(Left(new geny.Bytes(java.util.Arrays.copyOf(buf, n))))
+      ),
+      if (stderr ne os.Pipe) stderr
+      else os.ProcessOutput.ReadBytes((buf, n) =>
+        chunks.add(Right(new geny.Bytes(java.util.Arrays.copyOf(buf, n))))
+      ),
+      mergeErrIntoOut,
+      propagateEnv,
+      pipefail
+    )
+
+    sub.join(timeout)
+
+    val chunksSeq = chunks.iterator.asScala.toIndexedSeq
+    val res =
+      CommandResult(commands.flatMap(_.commandChunks :+ "|").init, sub.exitCode(), chunksSeq)
+    if (res.exitCode == 0 || !check) res
+    else throw SubprocessException(res)
+  }
+
+  /**
+   * The most flexible of the [[os.ProcGroup]] calls. It sets-up a pipeline of processes,
+   * and returns a [[ProcessPipeline]] for you to interact with however you like.
+   *
+   * Note that if you provide `ProcessOutput` callbacks to `stdout`/`stderr`,
+   * the calls to those callbacks take place on newly spawned threads that
+   * execute in parallel with the main thread. Thus make sure any data
+   * processing you do in those callbacks is thread safe!
+   * @param cwd              the working directory of the pipeline
+   * @param env              any additional environment variables you wish to set in the pipeline
+   * @param stdin            any data you wish to pass to the pipelines's standard input (to the first process)
+   * @param stdout           How the pipelines's output stream is configured (the last process stdout)
+   * @param stderr           How the process's error stream is configured (set for all processes)
+   * @param mergeErrIntoOut  merges the pipeline's stderr stream into it's stdout. Note that then the
+   *                         stderr will be forwarded with stdout to subsequent processes in the pipeline.
+   * @param propagateEnv     disable this to avoid passing in this parent process's
+   *                         environment variables to the pipeline
+   * @param pipefail         if true, the pipeline's exitCode will be the exit code of the first
+   *                         failing process. If no process fails, the exit code will be 0.
+   * @param handleBrokenPipe if true, every [[java.io.IOException]] when redirecting output of a process
+   *                         will be caught and handled by killing the writing process. This behaviour
+   *                         is consistent with handlers of SIGPIPE signals in most programs
+   *                         supporting interruptable piping. Disabled by default on Windows.
+   */
+  def spawn(
+      cwd: Path = null,
+      env: Map[String, String] = null,
+      stdin: ProcessInput = Pipe,
+      stdout: ProcessOutput = Pipe,
+      stderr: ProcessOutput = os.Inherit,
+      mergeErrIntoOut: Boolean = false,
+      propagateEnv: Boolean = true,
+      pipefail: Boolean = true,
+      handleBrokenPipe: Boolean = !isWindows
+  ): ProcessPipeline = {
+    val brokenPipeQueue = new LinkedBlockingQueue[Int]()
+    val (_, procs) =
+      commands.zipWithIndex.foldLeft((Option.empty[ProcessInput], Seq.empty[SubProcess])) {
+        case ((None, _), (proc, _)) =>
+          val spawned = proc.spawn(cwd, env, stdin, Pipe, stderr, mergeErrIntoOut, propagateEnv)
+          (Some(spawned.stdout), Seq(spawned))
+        case ((Some(input), acc), (proc, index)) if index == commands.length - 1 =>
+          val spawned = proc.spawn(
+            cwd,
+            env,
+            wrapWithBrokenPipeHandler(input, index - 1, brokenPipeQueue),
+            stdout,
+            stderr,
+            mergeErrIntoOut,
+            propagateEnv
+          )
+          (None, acc :+ spawned)
+        case ((Some(input), acc), (proc, index)) =>
+          val spawned = proc.spawn(
+            cwd,
+            env,
+            wrapWithBrokenPipeHandler(input, index - 1, brokenPipeQueue),
+            Pipe,
+            stderr,
+            mergeErrIntoOut,
+            propagateEnv
+          )
+          (Some(spawned.stdout), acc :+ spawned)
+      }
+    val pipeline =
+      new ProcessPipeline(procs, pipefail, if (handleBrokenPipe) Some(brokenPipeQueue) else None)
+    pipeline.brokenPipeHandler.foreach(_.start())
+    pipeline
+  }
+
+  private def wrapWithBrokenPipeHandler(
+      wrapped: ProcessInput,
+      index: Int,
+      queue: LinkedBlockingQueue[Int]
+  ) =
+    new ProcessInput {
+      override def redirectFrom: Redirect = wrapped.redirectFrom
+      override def processInput(stdin: => InputStream): Option[Runnable] =
+        wrapped.processInput(stdin).map { runnable =>
+          new Runnable {
+            def run() = {
+              try {
+                runnable.run()
+              } catch {
+                case e: IOException =>
+                  println(s"Broken pipe in process $index")
+                  queue.put(index)
+              }
+            }
+          }
+        }
+    }
+
+  /**
+   * Pipes the output of this pipeline into the input of the [[next]] process.
+   */
+  def pipeTo(next: proc) = ProcGroup(commands :+ next)
+}
+
+private[os] object ProcessOps {
+  def buildProcess(
+      command: Seq[String],
+      cwd: Path = null,
+      env: Map[String, String] = null,
+      stdin: ProcessInput = Pipe,
+      stdout: ProcessOutput = Pipe,
+      stderr: ProcessOutput = os.Inherit,
+      mergeErrIntoOut: Boolean = false,
+      propagateEnv: Boolean = true
+  ): ProcessBuilder = {
+    val builder = new java.lang.ProcessBuilder()
+
+    val baseEnv =
+      if (propagateEnv) sys.env
+      else Map()
+    for ((k, v) <- baseEnv ++ Option(env).getOrElse(Map())) {
+      if (v != null) builder.environment().put(k, v)
+      else builder.environment().remove(k)
+    }
+
+    builder.directory(Option(cwd).getOrElse(os.pwd).toIO)
+
+    builder
+      .command(command: _*)
+      .redirectInput(stdin.redirectFrom)
+      .redirectOutput(stdout.redirectTo)
+      .redirectError(stderr.redirectTo)
+      .redirectErrorStream(mergeErrIntoOut)
+  }
 }