os/exec: start checking for context cancelation in Start

Previously we started checking for context cancelation in Wait, but
that meant that when using StdoutPipe context cancelation never took
effect.

Fixes #16222.

Change-Id: I89cd26d3499a6080bf1a07718ce38d825561899e
Reviewed-on: https://go-review.googlesource.com/24650
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

1 files changed, 23 insertions, 13 deletions

diff --git a/src/os/exec/exec.go b/src/os/exec/exec.go
index 10300ce234..d2c1b17e50 100644
--- a/src/os/exec/exec.go
+++ b/src/os/exec/exec.go
@@ -111,6 +111,7 @@ type Cmd struct {
 	closeAfterWait  []io.Closer
 	goroutine       []func() error
 	errch           chan error // one send per goroutine
+	waitDone        chan struct{}
 }
 
 // Command returns the Cmd struct to execute the named program with
@@ -326,6 +327,15 @@ func (c *Cmd) Start() error {
 	if c.Process != nil {
 		return errors.New("exec: already started")
 	}
+	if c.ctx != nil {
+		select {
+		case <-c.ctx.Done():
+			c.closeDescriptors(c.closeAfterStart)
+			c.closeDescriptors(c.closeAfterWait)
+			return c.ctx.Err()
+		default:
+		}
+	}
 
 	type F func(*Cmd) (*os.File, error)
 	for _, setupFd := range []F{(*Cmd).stdin, (*Cmd).stdout, (*Cmd).stderr} {
@@ -361,6 +371,17 @@ func (c *Cmd) Start() error {
 		}(fn)
 	}
 
+	if c.ctx != nil {
+		c.waitDone = make(chan struct{})
+		go func() {
+			select {
+			case <-c.ctx.Done():
+				c.Process.Kill()
+			case <-c.waitDone:
+			}
+		}()
+	}
+
 	return nil
 }
 
@@ -410,20 +431,9 @@ func (c *Cmd) Wait() error {
 	}
 	c.finished = true
 
-	var waitDone chan struct{}
-	if c.ctx != nil {
-		waitDone = make(chan struct{})
-		go func() {
-			select {
-			case <-c.ctx.Done():
-				c.Process.Kill()
-			case <-waitDone:
-			}
-		}()
-	}
 	state, err := c.Process.Wait()
-	if waitDone != nil {
-		close(waitDone)
+	if c.waitDone != nil {
+		close(c.waitDone)
 	}
 	c.ProcessState = state