refactor: consolidate ForkReader-like implementations to IOCommand

util/iocommand.go

@@ -0,0 +1,107 @@
+package util
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"os/exec"
+)
+
+// An IOCommand exposes a forked process's std(in|out|err) through the io.ReadWriteCloser interface.
+type IOCommand struct {
+	Cmd        *exec.Cmd
+	CmdContext context.Context
+	CmdCancel  context.CancelFunc
+	Stdin      io.Writer
+	Stdout     io.Reader
+	StderrBuf  *bytes.Buffer
+}
+
+const IOCommandStderrBufSize = 1024
+
+type IOCommandError struct {
+	WaitErr error
+	Stderr  []byte
+}
+
+func (e IOCommandError) Error() string {
+	return fmt.Sprintf("underlying process exited with error: %s\nstderr: %s\n", e.WaitErr, e.Stderr)
+}
+
+func RunIOCommand(command string, args ...string) (c *IOCommand, err error) {
+	c, err = NewIOCommand(command, args, IOCommandStderrBufSize)
+	if err != nil {
+		return
+	}
+	err = c.Start()
+	return
+}
+
+func NewIOCommand(command string, args []string, stderrBufSize int) (c *IOCommand, err error) {
+
+	if stderrBufSize == 0 {
+		stderrBufSize = IOCommandStderrBufSize
+	}
+
+	c = &IOCommand{}
+
+	c.CmdContext, c.CmdCancel = context.WithCancel(context.Background())
+	c.Cmd = exec.CommandContext(c.CmdContext, command, args...)
+
+	if c.Stdout, err = c.Cmd.StdoutPipe(); err != nil {
+		return
+	}
+
+	if c.Stdin, err = c.Cmd.StdinPipe(); err != nil {
+		return
+	}
+
+	c.StderrBuf = bytes.NewBuffer(make([]byte, 0, stderrBufSize))
+	c.Cmd.Stderr = c.StderrBuf
+
+	return
+
+}
+
+func (c *IOCommand) Start() (err error) {
+	if err = c.Cmd.Start(); err != nil {
+		return
+	}
+	return
+}
+
+// Read from process's stdout.
+// The behavior after Close()ing is undefined
+func (c *IOCommand) Read(buf []byte) (n int, err error) {
+	n, err = c.Stdout.Read(buf)
+	if err == io.EOF {
+		if waitErr := c.doWait(); waitErr != nil {
+			err = waitErr
+		}
+	}
+	return
+}
+
+func (c *IOCommand) doWait() (err error) {
+	waitErr := c.Cmd.Wait()
+	if waitErr != nil {
+		err = IOCommandError{
+			WaitErr: waitErr,
+			Stderr:  c.StderrBuf.Bytes(),
+		}
+	}
+	return
+}
+
+// Write to process's stdin.
+// The behavior after Close()ing is undefined
+func (c *IOCommand) Write(buf []byte) (n int, err error) {
+	return c.Stdin.Write(buf)
+}
+
+// Kill the child process and collect its exit status
+func (c *IOCommand) Close() error {
+	c.CmdCancel()
+	return c.doWait()
+}