Wait for pipes to write all data before exit

Add a waitgroup to wait for the io.Copy of stdout/err to finish before
existing runc.  The problem happens more in exec because it is really
fast and the pipe has data buffered but not yet read after the process
has already exited.

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>

restore.go

@@ -140,21 +140,19 @@ func restoreContainer(context *cli.Context, spec *specs.LinuxSpec, config *confi
 	if err != nil {
 		return -1, err
 	}
+	defer tty.Close()
 	handler := newSignalHandler(tty)
 	defer handler.Close()
 	if err := container.Restore(process, options); err != nil {
-		if tty != nil {
+		return -1, err
-			tty.Close()
 	}
+	if err := tty.ClosePostStart(); err != nil {
 		return -1, err
 	}
 	if pidFile := context.String("pid-file"); pidFile != "" {
 		if err := createPidFile(pidFile, process); err != nil {
 			process.Signal(syscall.SIGKILL)
 			process.Wait()
-			if tty != nil {
-				tty.Close()
-			}
 			return -1, err
 		}
 	}

tty.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"sync"
 
 	"github.com/docker/docker/pkg/term"
 	"github.com/opencontainers/runc/libcontainer"
@@ -35,15 +36,32 @@ func createStdioPipes(p *libcontainer.Process, rootuid int) (*tty, error) {
 			i.Stderr,
 		},
 	}
+	// add the process's io to the post start closers if they support close
+	for _, cc := range []interface{}{
+		p.Stdin,
+		p.Stdout,
+		p.Stderr,
+	} {
+		if c, ok := cc.(io.Closer); ok {
+			t.postStart = append(t.postStart, c)
+		}
+	}
 	go func() {
 		io.Copy(i.Stdin, os.Stdin)
 		i.Stdin.Close()
 	}()
-	go io.Copy(os.Stdout, i.Stdout)
+	t.wg.Add(2)
-	go io.Copy(os.Stderr, i.Stderr)
+	go t.copyIO(os.Stdout, i.Stdout)
+	go t.copyIO(os.Stderr, i.Stderr)
 	return t, nil
 }
 
+func (t *tty) copyIO(w io.Writer, r io.ReadCloser) {
+	defer t.wg.Done()
+	io.Copy(w, r)
+	r.Close()
+}
+
 func createTty(p *libcontainer.Process, rootuid int, consolePath string) (*tty, error) {
 	if consolePath != "" {
 		if err := p.ConsoleFromPath(consolePath); err != nil {
@@ -62,23 +80,41 @@ func createTty(p *libcontainer.Process, rootuid int, consolePath string) (*tty,
 	if err != nil {
 		return nil, fmt.Errorf("failed to set the terminal from the stdin: %v", err)
 	}
-	t := &tty{
+	return &tty{
 		console: console,
 		state:   state,
 		closers: []io.Closer{
 			console,
 		},
-	}
+	}, nil
-	return t, nil
 }
 
 type tty struct {
 	console   libcontainer.Console
 	state     *term.State
 	closers   []io.Closer
+	postStart []io.Closer
+	wg        sync.WaitGroup
 }
 
+// ClosePostStart closes any fds that are provided to the container and dup2'd
+// so that we no longer have copy in our process.
+func (t *tty) ClosePostStart() error {
+	for _, c := range t.postStart {
+		c.Close()
+	}
+	return nil
+}
+
+// Close closes all open fds for the tty and/or restores the orignal
+// stdin state to what it was prior to the container execution
 func (t *tty) Close() error {
+	// ensure that our side of the fds are always closed
+	for _, c := range t.postStart {
+		c.Close()
+	}
+	// wait for the copy routines to finish before closing the fds
+	t.wg.Wait()
 	for _, c := range t.closers {
 		c.Close()
 	}

utils.go

@@ -218,6 +218,8 @@ func destroy(container libcontainer.Container) {
 	}
 }
 
+// setupIO sets the proper IO on the process depending on the configuration
+// If there is a nil error then there must be a non nil tty returned
 func setupIO(process *libcontainer.Process, rootuid int, console string, createTTY, detach bool) (*tty, error) {
 	// detach and createTty will not work unless a console path is passed
 	// so error out here before changing any terminal settings
@@ -231,7 +233,7 @@ func setupIO(process *libcontainer.Process, rootuid int, console string, createT
 		if err := dupStdio(process, rootuid); err != nil {
 			return nil, err
 		}
-		return nil, nil
+		return &tty{}, nil
 	}
 	return createStdioPipes(process, rootuid)
 }
@@ -290,30 +292,24 @@ func runProcess(container libcontainer.Container, config *specs.Process, listenF
 	if err != nil {
 		return -1, err
 	}
-
+	defer tty.Close()
 	handler := newSignalHandler(tty)
 	defer handler.Close()
-
 	if err := container.Start(process); err != nil {
-		if tty != nil {
-			tty.Close()
-		}
 		return -1, err
 	}
-
+	if err := tty.ClosePostStart(); err != nil {
+		return -1, err
+	}
 	if pidFile != "" {
 		if err := createPidFile(pidFile, process); err != nil {
 			process.Signal(syscall.SIGKILL)
 			process.Wait()
-			if tty != nil {
-				tty.Close()
-			}
 			return -1, err
 		}
 	}
 	if detach {
 		return 0, nil
 	}
-
 	return handler.forward(process)
 }