Merge pull request #1181 from mlaventure/fix-races

Fix races

client_test.go

@@ -98,7 +98,7 @@ func TestMain(m *testing.M) {
 		if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
 			fmt.Fprintln(os.Stderr, err)
 		}
-		if _, err := cmd.Process.Wait(); err != nil {
+		if err := cmd.Wait(); err != nil {
 			fmt.Fprintln(os.Stderr, err)
 		}
 		if err := os.RemoveAll(defaultRoot); err != nil {

io.go

@@ -1,6 +1,7 @@
 package containerd
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -18,6 +19,13 @@ type IO struct {
 	closer *wgCloser
 }
 
+func (i *IO) Cancel() {
+	if i.closer == nil {
+		return
+	}
+	i.closer.Cancel()
+}
+
 func (i *IO) Wait() {
 	if i.closer == nil {
 		return
@@ -134,9 +142,10 @@ type ioSet struct {
 }
 
 type wgCloser struct {
-	wg  *sync.WaitGroup
-	dir string
-	set []io.Closer
+	wg     *sync.WaitGroup
+	dir    string
+	set    []io.Closer
+	cancel context.CancelFunc
 }
 
 func (g *wgCloser) Wait() {
@@ -152,3 +161,7 @@ func (g *wgCloser) Close() error {
 	}
 	return nil
 }
+
+func (g *wgCloser) Cancel() {
+	g.cancel()
+}

io_unix.go

@@ -13,16 +13,17 @@ import (
 
 func copyIO(fifos *FIFOSet, ioset *ioSet, tty bool) (_ *wgCloser, err error) {
 	var (
-		f   io.ReadWriteCloser
-		set []io.Closer
-		ctx = context.Background()
-		wg  = &sync.WaitGroup{}
+		f           io.ReadWriteCloser
+		set         []io.Closer
+		ctx, cancel = context.WithCancel(context.Background())
+		wg          = &sync.WaitGroup{}
 	)
 	defer func() {
 		if err != nil {
 			for _, f := range set {
 				f.Close()
 			}
+			cancel()
 		}
 	}()
 
@@ -55,13 +56,14 @@ func copyIO(fifos *FIFOSet, ioset *ioSet, tty bool) (_ *wgCloser, err error) {
 		wg.Add(1)
 		go func(r io.ReadCloser) {
 			io.Copy(ioset.err, r)
-			wg.Done()
 			r.Close()
+			wg.Done()
 		}(f)
 	}
 	return &wgCloser{
-		wg:  wg,
-		dir: fifos.Dir,
-		set: set,
+		wg:     wg,
+		dir:    fifos.Dir,
+		set:    set,
+		cancel: cancel,
 	}, nil
 }

process.go

@@ -45,6 +45,8 @@ func (p *process) Start(ctx context.Context) error {
 	}
 	response, err := p.task.client.TaskService().Exec(ctx, request)
 	if err != nil {
+		p.io.Cancel()
+		p.io.Wait()
 		p.io.Close()
 		return err
 	}