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
@@ -137,6 +145,7 @@ type wgCloser struct {
 	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

@@ -15,7 +15,7 @@ func copyIO(fifos *FIFOSet, ioset *ioSet, tty bool) (_ *wgCloser, err error) {
 	var (
 		f           io.ReadWriteCloser
 		set         []io.Closer
-		ctx = context.Background()
+		ctx, cancel = context.WithCancel(context.Background())
 		wg          = &sync.WaitGroup{}
 	)
 	defer func() {
@@ -23,6 +23,7 @@ func copyIO(fifos *FIFOSet, ioset *ioSet, tty bool) (_ *wgCloser, err error) {
 			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,
+		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
 	}