refactor: replace Timeout options with context.Context (#126)

Author: unknwon

blob.go

@@ -6,6 +6,7 @@ package git
 
 import (
 	"bytes"
+	"context"
 	"io"
 )
 
@@ -16,21 +17,23 @@ type Blob struct {
 
 // Bytes reads and returns the content of the blob all at once in bytes. This
 // can be very slow and memory consuming for huge content.
-func (b *Blob) Bytes() ([]byte, error) {
+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.
-	stdout.Grow(int(b.Size()))
+	if size := b.Size(ctx); size > 0 && size < int64(^uint(0)>>1) {
+		stdout.Grow(int(size))
+	}
 
-	if err := b.Pipeline(stdout, stderr); err != nil {
+	if err := b.Pipeline(ctx, stdout, stderr); err != nil {
 		return nil, concatenateError(err, stderr.String())
 	}
 	return stdout.Bytes(), nil
 }
 
 // Pipeline reads the content of the blob and pipes stdout and stderr to
 // supplied io.Writer.
-func (b *Blob) Pipeline(stdout, stderr io.Writer) error {
-	return NewCommand("show", b.id.String()).RunInDirPipeline(stdout, stderr, b.parent.repo.path)
+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)
 }

blob_test.go

@@ -6,12 +6,14 @@ package git
 
 import (
 	"bytes"
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestBlob(t *testing.T) {
+	ctx := context.Background()
 	expOutput := `This is a sample project students can use during Matthew's Git class.
 
 Here is an addition by me
@@ -41,14 +43,14 @@ This demo also includes an image with changes on a branch for examination of ima
 	}
 
 	t.Run("get data all at once", func(t *testing.T) {
-		p, err := blob.Bytes()
+		p, err := blob.Bytes(ctx)
 		assert.Nil(t, err)
 		assert.Equal(t, expOutput, string(p))
 	})
 
 	t.Run("get data with pipeline", func(t *testing.T) {
 		stdout := new(bytes.Buffer)
-		err := blob.Pipeline(stdout, nil)
+		err := blob.Pipeline(ctx, stdout, nil)
 		assert.Nil(t, err)
 		assert.Equal(t, expOutput, stdout.String())
 	})

command.go

@@ -17,22 +17,16 @@ import (
 
 // Command contains the name, arguments and environment variables of a command.
 type Command struct {
-	name    string
-	args    []string
-	envs    []string
-	timeout time.Duration
-	ctx     context.Context
+	name string
+	args []string
+	envs []string
+	ctx  context.Context
 }
 
 // CommandOptions contains options for running a command.
-// If timeout is zero, DefaultTimeout will be used.
-// If timeout is less than zero, no timeout will be set.
-// If context is nil, context.Background() will be used.
 type CommandOptions struct {
-	Args    []string
-	Envs    []string
-	Timeout time.Duration
-	Context context.Context
+	Args []string
+	Envs []string
 }
 
 // String returns the string representation of the command.
@@ -44,13 +38,7 @@ func (c *Command) String() string {
 }
 
 // NewCommand creates and returns a new Command with given arguments for "git".
-func NewCommand(args ...string) *Command {
-	return NewCommandWithContext(context.Background(), args...)
-}
-
-// NewCommandWithContext creates and returns a new Command with given arguments
-// and context for "git".
-func NewCommandWithContext(ctx context.Context, args ...string) *Command {
+func NewCommand(ctx context.Context, args ...string) *Command {
 	return &Command{
 		name: "git",
 		args: args,
@@ -76,23 +64,9 @@ func (c Command) WithContext(ctx context.Context) *Command {
 	return &c
 }
 
-// WithTimeout returns a new Command with given timeout.
-func (c Command) WithTimeout(timeout time.Duration) *Command {
-	c.timeout = timeout
-	return &c
-}
-
-// SetTimeout sets the timeout for the command.
-func (c *Command) SetTimeout(timeout time.Duration) {
-	c.timeout = timeout
-}
-
 // AddOptions adds options to the command.
-// Note: only the last option will take effect if there are duplicated options.
 func (c *Command) AddOptions(opts ...CommandOptions) *Command {
 	for _, opt := range opts {
-		c.timeout = opt.Timeout
-		c.ctx = opt.Context
 		c.AddArgs(opt.Args...)
 		c.AddEnvs(opt.Envs...)
 	}
@@ -105,7 +79,8 @@ func (c *Command) AddCommitter(committer *Signature) *Command {
 	return c
 }
 
-// DefaultTimeout is the default timeout duration for all commands.
+// 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
 
 // A limitDualWriter writes to W but limits the amount of data written to just N
@@ -144,33 +119,19 @@ type RunInDirOptions struct {
 	Stdout io.Writer
 	// Stderr is the error output from the command.
 	Stderr io.Writer
-	// Timeout is the duration to wait before timing out.
-	//
-	// Deprecated: Use CommandOptions.Timeout or *Command.WithTimeout instead.
-	Timeout time.Duration
 }
 
 // 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. DefaultTimeout will be used if the timeout duration is less than
-// time.Nanosecond (i.e. less than or equal to 0). It returns an ErrExecTimeout
-// if the execution was timed out.
+// 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]
 	}
 
-	timeout := c.timeout
-	// TODO: remove this in newer version
-	if opt.Timeout > 0 {
-		timeout = opt.Timeout
-	}
-
-	if timeout == 0 {
-		timeout = DefaultTimeout
-	}
-
 	buf := new(bytes.Buffer)
 	w := opt.Stdout
 	if logOutput != nil {
@@ -184,26 +145,22 @@ func (c *Command) RunInDirWithOptions(dir string, opts ...RunInDirOptions) (err
 
 	defer func() {
 		if len(dir) == 0 {
-			log("[timeout: %v] %s\n%s", timeout, c, buf.Bytes())
+			log("%s\n%s", c, buf.Bytes())
 		} else {
-			log("[timeout: %v] %s: %s\n%s", timeout, dir, c, buf.Bytes())
+			log("%s: %s\n%s", dir, c, buf.Bytes())
 		}
 	}()
 
-	ctx := context.Background()
-	if c.ctx != nil {
-		ctx = c.ctx
+	ctx := c.ctx
+	if ctx == nil {
+		ctx = context.Background()
 	}
 
-	if timeout > 0 {
+	// 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, timeout)
-		defer func() {
-			cancel()
-			if err == context.DeadlineExceeded {
-				err = ErrExecTimeout
-			}
-		}()
+		ctx, cancel = context.WithTimeout(ctx, DefaultTimeout)
+		defer cancel()
 	}
 
 	cmd := exec.CommandContext(ctx, c.name, c.args...)
@@ -215,6 +172,11 @@ func (c *Command) RunInDirWithOptions(dir string, opts ...RunInDirOptions) (err
 	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
 	}
 
@@ -225,22 +187,33 @@ func (c *Command) RunInDirWithOptions(dir string, opts ...RunInDirOptions) (err
 
 	select {
 	case <-ctx.Done():
-		<-result
-		if cmd.Process != nil && cmd.ProcessState != nil && !cmd.ProcessState.Exited() {
-			if err := cmd.Process.Kill(); err != nil {
-				return fmt.Errorf("kill process: %v", err)
-			}
+		// Kill the process before waiting so cancellation is enforced promptly.
+		if cmd.Process != nil {
+			_ = cmd.Process.Kill()
 		}
+		<-result
 
-		return ErrExecTimeout
+		if ctx.Err() == 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
 	}
 
 }
 
-// RunInDirPipeline executes the command in given directory and default timeout
-// duration. It pipes stdout and stderr to supplied io.Writer.
+// 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,
@@ -249,35 +222,8 @@ func (c *Command) RunInDirPipeline(stdout, stderr io.Writer, dir string) error {
 	})
 }
 
-// RunInDirPipelineWithTimeout executes the command in given directory and
-// timeout duration. It pipes stdout and stderr to supplied io.Writer.
-// DefaultTimeout will be used if the timeout duration is less than
-// time.Nanosecond (i.e. less than or equal to 0). It returns an ErrExecTimeout
-// if the execution was timed out.
-//
-// Deprecated: Use RunInDirPipeline and CommandOptions instead.
-// TODO: remove this in the next major version
-func (c *Command) RunInDirPipelineWithTimeout(timeout time.Duration, stdout, stderr io.Writer, dir string) (err error) {
-	if timeout != 0 {
-		c = c.WithTimeout(timeout)
-	}
-	return c.RunInDirPipeline(stdout, stderr, dir)
-}
-
-// RunInDirWithTimeout executes the command in given directory and timeout
-// duration. It returns stdout in []byte and error (combined with stderr).
-//
-// Deprecated: Use RunInDir and CommandOptions instead.
-// TODO: remove this in the next major version
-func (c *Command) RunInDirWithTimeout(timeout time.Duration, dir string) ([]byte, error) {
-	if timeout != 0 {
-		c = c.WithTimeout(timeout)
-	}
-	return c.RunInDir(dir)
-}
-
-// RunInDir executes the command in given directory and default timeout
-// duration. It returns stdout and error (combined with stderr).
+// 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)
@@ -287,20 +233,8 @@ func (c *Command) RunInDir(dir string) ([]byte, error) {
 	return stdout.Bytes(), nil
 }
 
-// RunWithTimeout executes the command in working directory and given timeout
-// duration. It returns stdout in string and error (combined with stderr).
-//
-// Deprecated: Use RunInDir and CommandOptions instead.
-// TODO: remove this in the next major version
-func (c *Command) RunWithTimeout(timeout time.Duration) ([]byte, error) {
-	if timeout != 0 {
-		c = c.WithTimeout(timeout)
-	}
-	return c.Run()
-}
-
-// Run executes the command in working directory and default timeout duration.
-// It returns stdout in string and error (combined with stderr).
+// 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 {