Merge pull request #593 from crosbymichael/wait-pipes

Wait for pipes to write all data before exit

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
 		}
 	}

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()
 	}

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)
 }