refactor: use sourcegraph/run for command execution

blob.go

@@ -19,21 +19,20 @@ type Blob struct {
 // can be very slow and memory consuming for huge content.
 func (b *Blob) Bytes(ctx context.Context) ([]byte, error) {
 	stdout := new(bytes.Buffer)
-	stderr := new(bytes.Buffer)
 
 	// Preallocate memory to save ~50% memory usage on big files.
 	if size := b.Size(ctx); size > 0 && size < int64(^uint(0)>>1) {
 		stdout.Grow(int(size))
 	}
 
-	if err := b.Pipeline(ctx, stdout, stderr); err != nil {
-		return nil, concatenateError(err, stderr.String())
+	if err := b.Pipeline(ctx, stdout, nil); err != nil {
+		return nil, err
 	}
 	return stdout.Bytes(), nil
 }
 
 // Pipeline reads the content of the blob and pipes stdout and stderr to
 // supplied io.Writer.
 func (b *Blob) Pipeline(ctx context.Context, stdout, stderr io.Writer) error {
-	return NewCommand(ctx, "show", b.id.String()).RunInDirPipeline(stdout, stderr, b.parent.repo.path)
+	return gitPipeline(ctx, b.parent.repo.path, []string{"show", b.id.String()}, nil, stdout, stderr)
 }

command.go

@@ -10,78 +10,159 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"os/exec"
+	"strconv"
 	"strings"
 	"time"
-)
 
-// Command contains the name, arguments and environment variables of a command.
-type Command struct {
-	name string
-	args []string
-	envs []string
-	ctx  context.Context
-}
+	"github.com/sourcegraph/run"
+)
 
 // CommandOptions contains options for running a command.
 type CommandOptions struct {
 	Args []string
 	Envs []string
 }
 
-// String returns the string representation of the command.
-func (c *Command) String() string {
-	if len(c.args) == 0 {
-		return c.name
+// DefaultTimeout is the default timeout duration for all commands. It is
+// applied when the context does not already have a deadline.
+const DefaultTimeout = time.Minute
+
+// gitCmd builds a *run.Command for "git" with the given arguments, environment
+// variables and working directory. If the context does not already have a
+// deadline, DefaultTimeout will be applied automatically.
+func gitCmd(ctx context.Context, dir string, args []string, envs []string) (*run.Command, context.CancelFunc) {
+	cancel := func() {}
+
+	// Apply default timeout if the context doesn't already have a deadline.
+	if _, ok := ctx.Deadline(); !ok {
+		var timeoutCancel context.CancelFunc
+		ctx, timeoutCancel = context.WithTimeout(ctx, DefaultTimeout)
+		cancel = timeoutCancel
 	}
-	return fmt.Sprintf("%s %s", c.name, strings.Join(c.args, " "))
-}
 
-// NewCommand creates and returns a new Command with given arguments for "git".
-func NewCommand(ctx context.Context, args ...string) *Command {
-	return &Command{
-		name: "git",
-		args: args,
-		ctx:  ctx,
+	// run.Cmd joins all parts into a single string and then shell-parses it.
+	// We must quote each argument so that special characters (spaces, quotes,
+	// angle brackets, etc.) are preserved correctly.
+	parts := make([]string, 0, 1+len(args))
+	parts = append(parts, "git")
+	for _, arg := range args {
+		parts = append(parts, run.Arg(arg))
 	}
-}
 
-// AddArgs appends given arguments to the command.
-func (c *Command) AddArgs(args ...string) *Command {
-	c.args = append(c.args, args...)
-	return c
+	cmd := run.Cmd(ctx, parts...)
+	if dir != "" {
+		cmd = cmd.Dir(dir)
+	}
+	if len(envs) > 0 {
+		cmd = cmd.Environ(append(os.Environ(), envs...))
+	}
+	return cmd, cancel
 }
 
-// AddEnvs appends given environment variables to the command.
-func (c *Command) AddEnvs(envs ...string) *Command {
-	c.envs = append(c.envs, envs...)
-	return c
-}
+// gitRun executes a git command in the given directory and returns stdout as
+// bytes. Stderr is included in the error message on failure. If the command's
+// context does not have a deadline, DefaultTimeout will be applied
+// automatically. It returns an ErrExecTimeout if the execution was timed out.
+func gitRun(ctx context.Context, dir string, args []string, envs []string) ([]byte, error) {
+	cmd, cancel := gitCmd(ctx, dir, args, envs)
+	defer cancel()
+
+	var logBuf *bytes.Buffer
+	if logOutput != nil {
+		logBuf = new(bytes.Buffer)
+		logBuf.Grow(512)
+		defer func() {
+			logf(dir, args, logBuf.Bytes())
+		}()
+	}
+
+	// Use Stream to a buffer to preserve raw bytes (including NUL bytes from
+	// commands like "ls-tree -z"). The String/Lines methods process output
+	// line-by-line which corrupts binary-ish output.
+	stdout := new(bytes.Buffer)
+	err := cmd.StdOut().Run().Stream(stdout)
+
+	// Capture (partial) stdout for logging even on error, so failed commands
+	// produce a useful log entry rather than an empty one.
+	if logOutput != nil {
+		data := stdout.Bytes()
+		limit := len(data)
+		if limit > 512 {
+			limit = 512
+		}
+		logBuf.Write(data[:limit])
+		if len(data) > 512 {
+			logBuf.WriteString("... (more omitted)")
+		}
+	}
 
-// WithContext returns a new Command with the given context.
-func (c Command) WithContext(ctx context.Context) *Command {
-	c.ctx = ctx
-	return &c
+	if err != nil {
+		return nil, mapContextError(err, ctx)
+	}
+	return stdout.Bytes(), nil
 }
 
-// AddOptions adds options to the command.
-func (c *Command) AddOptions(opts ...CommandOptions) *Command {
-	for _, opt := range opts {
-		c.AddArgs(opt.Args...)
-		c.AddEnvs(opt.Envs...)
+// gitPipeline executes a git command in the given directory, streaming stdout
+// to the given writer. If stderr writer is provided and the command fails,
+// stderr content extracted from the error is written to it.
+func gitPipeline(ctx context.Context, dir string, args []string, envs []string, stdout, stderr io.Writer) error {
+	cmd, cancel := gitCmd(ctx, dir, args, envs)
+	defer cancel()
+
+	var buf *bytes.Buffer
+	w := stdout
+	if logOutput != nil {
+		buf = new(bytes.Buffer)
+		buf.Grow(512)
+		w = &limitDualWriter{
+			W: buf,
+			N: int64(buf.Cap()),
+			w: stdout,
+		}
+
+		defer func() {
+			logf(dir, args, buf.Bytes())
+		}()
+	}
+
+	streamErr := cmd.StdOut().Run().Stream(w)
+	if streamErr != nil {
+		if stderr != nil {
+			_, _ = fmt.Fprint(stderr, extractStderr(streamErr))
+		}
+		return mapContextError(streamErr, ctx)
 	}
-	return c
+	return nil
 }
 
-// AddCommitter appends given committer to the command.
-func (c *Command) AddCommitter(committer *Signature) *Command {
-	c.AddEnvs("GIT_COMMITTER_NAME="+committer.Name, "GIT_COMMITTER_EMAIL="+committer.Email)
-	return c
+// committerEnvs returns environment variables for setting the Git committer.
+func committerEnvs(committer *Signature) []string {
+	return []string{
+		"GIT_COMMITTER_NAME=" + committer.Name,
+		"GIT_COMMITTER_EMAIL=" + committer.Email,
+	}
 }
 
-// DefaultTimeout is the default timeout duration for all commands. It is
-// applied when the context does not already have a deadline.
-const DefaultTimeout = time.Minute
+// logf logs a git command execution with optional output.
+func logf(dir string, args []string, output []byte) {
+	cmdStr := "git"
+	if len(args) > 0 {
+		quoted := make([]string, len(args))
+		for i, a := range args {
+			if strings.ContainsAny(a, " \t\n\"'\\<>") {
+				quoted[i] = strconv.Quote(a)
+			} else {
+				quoted[i] = a
+			}
+		}
+		cmdStr = "git " + strings.Join(quoted, " ")
+	}
+	if len(dir) == 0 {
+		log("%s\n%s", cmdStr, output)
+	} else {
+		log("%s: %s\n%s", dir, cmdStr, output)
+	}
+}
 
 // A limitDualWriter writes to W but limits the amount of data written to just N
 // bytes. On the other hand, it passes everything to w.
@@ -111,134 +192,38 @@ func (w *limitDualWriter) Write(p []byte) (int, error) {
 	return w.w.Write(p)
 }
 
-// RunInDirOptions contains options for running a command in a directory.
-type RunInDirOptions struct {
-	// Stdin is the input to the command.
-	Stdin io.Reader
-	// Stdout is the outputs from the command.
-	Stdout io.Writer
-	// Stderr is the error output from the command.
-	Stderr io.Writer
-}
-
-// RunInDirWithOptions executes the command in given directory and options. It
-// pipes stdin from supplied io.Reader, and pipes stdout and stderr to supplied
-// io.Writer. If the command's context does not have a deadline, DefaultTimeout
-// will be applied automatically. It returns an ErrExecTimeout if the execution
-// was timed out.
-func (c *Command) RunInDirWithOptions(dir string, opts ...RunInDirOptions) (err error) {
-	var opt RunInDirOptions
-	if len(opts) > 0 {
-		opt = opts[0]
-	}
-
-	buf := new(bytes.Buffer)
-	w := opt.Stdout
-	if logOutput != nil {
-		buf.Grow(512)
-		w = &limitDualWriter{
-			W: buf,
-			N: int64(buf.Cap()),
-			w: opt.Stdout,
-		}
-	}
-
-	defer func() {
-		if len(dir) == 0 {
-			log("%s\n%s", c, buf.Bytes())
-		} else {
-			log("%s: %s\n%s", dir, c, buf.Bytes())
-		}
-	}()
-
-	ctx := c.ctx
+// mapContextError maps context errors to the appropriate sentinel errors used
+// by this package.
+func mapContextError(err error, ctx context.Context) error {
 	if ctx == nil {
-		ctx = context.Background()
-	}
-
-	// Apply default timeout if the context doesn't already have a deadline.
-	if _, ok := ctx.Deadline(); !ok {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, DefaultTimeout)
-		defer cancel()
-	}
-
-	cmd := exec.CommandContext(ctx, c.name, c.args...)
-	if len(c.envs) > 0 {
-		cmd.Env = append(os.Environ(), c.envs...)
-	}
-	cmd.Dir = dir
-	cmd.Stdin = opt.Stdin
-	cmd.Stdout = w
-	cmd.Stderr = opt.Stderr
-	if err = cmd.Start(); err != nil {
-		if ctx.Err() == context.DeadlineExceeded {
-			return ErrExecTimeout
-		} else if ctx.Err() != nil {
-			return ctx.Err()
-		}
 		return err
 	}
-
-	result := make(chan error)
-	go func() {
-		result <- cmd.Wait()
-	}()
-
-	select {
-	case <-ctx.Done():
-		// Kill the process before waiting so cancellation is enforced promptly.
-		if cmd.Process != nil {
-			_ = cmd.Process.Kill()
-		}
-		<-result
-
-		if ctx.Err() == context.DeadlineExceeded {
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		if ctxErr == context.DeadlineExceeded {
 			return ErrExecTimeout
 		}
-		return ctx.Err()
-	case err = <-result:
-		// Normalize errors when the context may have expired around the same time.
-		if err != nil {
-			if ctxErr := ctx.Err(); ctxErr != nil {
-				if ctxErr == context.DeadlineExceeded {
-					return ErrExecTimeout
-				}
-				return ctxErr
-			}
-		}
-		return err
+		return ctxErr
 	}
-
+	return err
 }
 
-// RunInDirPipeline executes the command in given directory. It pipes stdout and
-// stderr to supplied io.Writer.
-func (c *Command) RunInDirPipeline(stdout, stderr io.Writer, dir string) error {
-	return c.RunInDirWithOptions(dir, RunInDirOptions{
-		Stdin:  nil,
-		Stdout: stdout,
-		Stderr: stderr,
-	})
+// isExitStatus reports whether err represents a specific process exit status
+// code, using the run.ExitCoder interface provided by sourcegraph/run.
+func isExitStatus(err error, code int) bool {
+	exitCoder, ok := err.(run.ExitCoder)
+	return ok && exitCoder.ExitCode() == code
 }
 
-// RunInDir executes the command in given directory. It returns stdout and error
-// (combined with stderr).
-func (c *Command) RunInDir(dir string) ([]byte, error) {
-	stdout := new(bytes.Buffer)
-	stderr := new(bytes.Buffer)
-	if err := c.RunInDirPipeline(stdout, stderr, dir); err != nil {
-		return nil, concatenateError(err, stderr.String())
+// extractStderr attempts to extract the stderr portion from a sourcegraph/run
+// error. The error format is typically "exit status N: <stderr content>".
+func extractStderr(err error) string {
+	if err == nil {
+		return ""
 	}
-	return stdout.Bytes(), nil
-}
-
-// Run executes the command in working directory. It returns stdout and
-// error (combined with stderr).
-func (c *Command) Run() ([]byte, error) {
-	stdout, err := c.RunInDir("")
-	if err != nil {
-		return nil, err
+	msg := err.Error()
+	// sourcegraph/run error format: "exit status N: <stderr>"
+	if idx := strings.Index(msg, ": "); idx >= 0 && strings.HasPrefix(msg, "exit status") {
+		return msg[idx+2:]
 	}
-	return stdout, nil
+	return msg
 }