Merge pull request #90 from github/add-pipe-package

Add an internal `pipe` package for managing process pipelines

Author: mhagger

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/cli/safeexec v1.0.0
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/spf13/pflag v1.0.5
-	github.com/stretchr/testify v1.4.0
-	gopkg.in/yaml.v2 v2.2.7 // indirect
+	github.com/stretchr/testify v1.7.0
+	go.uber.org/goleak v1.1.12
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 )

go.sum

@@ -1,19 +1,53 @@
 github.com/cli/safeexec v1.0.0 h1:0VngyaIyqACHdcMNWfo6+KdUYnqEr2Sg+bSP1pdF+dI=
 github.com/cli/safeexec v1.0.0/go.mod h1:Z/D4tTN8Vs5gXYHDCbaM1S/anmEDnJb1iW0+EJ5zx3Q=
-github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
+github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
-github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
-gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
+go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
+go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs=
+golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=
+golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
+golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
+golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.1.5 h1:ouewzE6p+/VEB31YYnTbEJdi8pFqKp4P4n85vwo3DHA=
+golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
+golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
-gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
-gopkg.in/yaml.v2 v2.2.7 h1:VUgggvou5XRW9mHwD/yXxIYSMtY0zoKQf/v226p2nyo=
-gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/pipe/command.go

@@ -0,0 +1,175 @@
+package pipe
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"os"
+	"os/exec"
+	"sync/atomic"
+	"syscall"
+
+	"golang.org/x/sync/errgroup"
+)
+
+// commandStage is a pipeline `Stage` based on running an external
+// command and piping the data through its stdin and stdout.
+type commandStage struct {
+	name   string
+	stdin  io.Closer
+	cmd    *exec.Cmd
+	done   chan struct{}
+	wg     errgroup.Group
+	stderr bytes.Buffer
+
+	// If the context expired and we attempted to kill the command,
+	// `ctx.Err()` is stored here.
+	ctxErr atomic.Value
+}
+
+// Command returns a pipeline `Stage` based on the specified external
+// `command`, run with the given command-line `args`. Its stdin and
+// stdout are handled as usual, and its stderr is collected and
+// included in any `*exec.ExitError` that the command might emit.
+func Command(command string, args ...string) Stage {
+	if len(command) == 0 {
+		panic("attempt to create command with empty command")
+	}
+
+	cmd := exec.Command(command, args...)
+	return CommandStage(command, cmd)
+}
+
+// Command returns a pipeline `Stage` with the name `name`, based on
+// the specified `cmd`. Its stdin and stdout are handled as usual, and
+// its stderr is collected and included in any `*exec.ExitError` that
+// the command might emit.
+func CommandStage(name string, cmd *exec.Cmd) Stage {
+	return &commandStage{
+		name: name,
+		cmd:  cmd,
+		done: make(chan struct{}),
+	}
+}
+
+func (s *commandStage) Name() string {
+	return s.name
+}
+
+func (s *commandStage) Start(
+	ctx context.Context, env Env, stdin io.ReadCloser,
+) (io.ReadCloser, error) {
+	if s.cmd.Dir == "" {
+		s.cmd.Dir = env.Dir
+	}
+
+	if stdin != nil {
+		s.cmd.Stdin = stdin
+		// Also keep a copy so that we can close it when the command exits:
+		s.stdin = stdin
+	}
+
+	stdout, err := s.cmd.StdoutPipe()
+	if err != nil {
+		return nil, err
+	}
+
+	// If the caller hasn't arranged otherwise, read the command's
+	// standard error into our `stderr` field:
+	if s.cmd.Stderr == nil {
+		// We can't just set `s.cmd.Stderr = &s.stderr`, because if we
+		// do then `s.cmd.Wait()` doesn't wait to be sure that all
+		// error output has been captured. By doing this ourselves, we
+		// can be sure.
+		p, err := s.cmd.StderrPipe()
+		if err != nil {
+			return nil, err
+		}
+		s.wg.Go(func() error {
+			_, err := io.Copy(&s.stderr, p)
+			// We don't consider `ErrClosed` an error (FIXME: is this
+			// correct?):
+			if err != nil && !errors.Is(err, os.ErrClosed) {
+				return err
+			}
+			return nil
+		})
+	}
+
+	// Put the command in its own process group, if possible:
+	s.runInOwnProcessGroup()
+
+	if err := s.cmd.Start(); err != nil {
+		return nil, err
+	}
+
+	// Arrange for the process to be killed (gently) if the context
+	// expires before the command exits normally:
+	go func() {
+		select {
+		case <-ctx.Done():
+			s.kill(ctx.Err())
+		case <-s.done:
+			// Process already done; no need to kill anything.
+		}
+	}()
+
+	return stdout, nil
+}
+
+// filterCmdError interprets `err`, which was returned by `Cmd.Wait()`
+// (possibly `nil`), possibly modifying it or ignoring it. It returns
+// the error that should actually be returned to the caller (possibly
+// `nil`).
+func (s *commandStage) filterCmdError(err error) error {
+	if err == nil {
+		return nil
+	}
+
+	eErr, ok := err.(*exec.ExitError)
+	if !ok {
+		return err
+	}
+
+	ctxErr, ok := s.ctxErr.Load().(error)
+	if ok {
+		// If the process looks like it was killed by us, substitute
+		// `ctxErr` for the process's own exit error. Note that this
+		// doesn't do anything on Windows, where the `Signaled()`
+		// method isn't implemented (it is hardcoded to return
+		// `false`).
+		ps, ok := eErr.ProcessState.Sys().(syscall.WaitStatus)
+		if ok && ps.Signaled() &&
+			(ps.Signal() == syscall.SIGTERM || ps.Signal() == syscall.SIGKILL) {
+			return ctxErr
+		}
+	}
+
+	eErr.Stderr = s.stderr.Bytes()
+	return eErr
+}
+
+func (s *commandStage) Wait() error {
+	defer close(s.done)
+
+	// Make sure that any stderr is copied before `s.cmd.Wait()`
+	// closes the read end of the pipe:
+	wErr := s.wg.Wait()
+
+	err := s.cmd.Wait()
+	err = s.filterCmdError(err)
+
+	if err == nil && wErr != nil {
+		err = wErr
+	}
+
+	if s.stdin != nil {
+		cErr := s.stdin.Close()
+		if cErr != nil && err == nil {
+			return cErr
+		}
+	}
+
+	return err
+}

internal/pipe/command_unix.go

@@ -0,0 +1,66 @@
+//go:build !windows
+// +build !windows
+
+package pipe
+
+import (
+	"syscall"
+	"time"
+)
+
+// runInOwnProcessGroup arranges for `cmd` to be run in its own
+// process group.
+func (s *commandStage) runInOwnProcessGroup() {
+	// Put the command in its own process group:
+	if s.cmd.SysProcAttr == nil {
+		s.cmd.SysProcAttr = &syscall.SysProcAttr{}
+	}
+	s.cmd.SysProcAttr.Setpgid = true
+}
+
+// kill is called to kill the process if the context expires. `err` is
+// the corresponding value of `Context.Err()`.
+func (s *commandStage) kill(err error) {
+	// I believe that the calls to `syscall.Kill()` in this method are
+	// racy. It could be that s.cmd.Wait() succeeds immediately before
+	// this call, in which case the process group wouldn't exist
+	// anymore. But I don't see any way to avoid this without
+	// duplicating a lot of code from `exec.Cmd`. (`os.Cmd.Kill()` and
+	// `os.Cmd.Signal()` appear to be race-free, but only because they
+	// use internal synchronization. But those methods only kill the
+	// process, not the process group, so they are not suitable here.
+
+	// We started the process with PGID == PID:
+	pid := s.cmd.Process.Pid
+	select {
+	case <-s.done:
+		// Process has ended; no need to kill it again.
+		return
+	default:
+	}
+
+	// Record the `ctx.Err()`, which will be used as the error result
+	// for this stage.
+	s.ctxErr.Store(err)
+
+	// First try to kill using a relatively gentle signal so that
+	// the processes have a chance to clean up after themselves:
+	_ = syscall.Kill(-pid, syscall.SIGTERM)
+
+	// Well-behaved processes should commit suicide after the above,
+	// but if they don't exit within 2s, murder the whole lot of them:
+	go func() {
+		// Use an explicit `time.Timer` rather than `time.After()` so
+		// that we can stop it (freeing resources) promptly if the
+		// command exits before the timer triggers.
+		timer := time.NewTimer(2 * time.Second)
+		defer timer.Stop()
+
+		select {
+		case <-s.done:
+			// Process has ended; no need to kill it again.
+		case <-timer.C:
+			_ = syscall.Kill(-pid, syscall.SIGKILL)
+		}
+	}()
+}

internal/pipe/command_windows.go

@@ -0,0 +1,24 @@
+//go:build windows
+// +build windows
+
+package pipe
+
+// runInOwnProcessGroup is not supported on Windows.
+func (s *commandStage) runInOwnProcessGroup() {}
+
+// kill is called to kill the process if the context expires. `err` is
+// the corresponding value of `Context.Err()`.
+func (s *commandStage) kill(err error) {
+	select {
+	case <-s.done:
+		// Process has ended; no need to kill it again.
+		return
+	default:
+	}
+
+	// Record the `ctx.Err()`, which will be used as the error result
+	// for this stage.
+	s.ctxErr.Store(err)
+
+	s.cmd.Process.Kill()
+}