From 8d0a05b8dda1e738415ac2bdf978ebbf40059fec Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Fri, 26 Feb 2016 12:14:47 -0800
Subject: [PATCH] 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 | 10 ++++------
 tty.go     | 52 ++++++++++++++++++++++++++++++++++++++++++++--------
 utils.go   | 18 +++++++-----------
 3 files changed, 55 insertions(+), 25 deletions(-)

diff --git a/restore.go b/restore.go
index 3ec0a1c0..5584c785 100644
--- a/restore.go
+++ b/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 {
-			tty.Close()
-		}
+		return -1, err
+	}
+	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
 		}
 	}
diff --git a/tty.go b/tty.go
index b552621b..fd531068 100644
--- a/tty.go
+++ b/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)
-	go io.Copy(os.Stderr, i.Stderr)
+	t.wg.Add(2)
+	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,
 		},
-	}
-	return t, nil
+	}, nil
 }
 
 type tty struct {
-	console libcontainer.Console
-	state   *term.State
-	closers []io.Closer
+	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()
 	}
diff --git a/utils.go b/utils.go
index 96212e36..94952114 100644
--- a/utils.go
+++ b/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)
 }