From e57fccc4818d3e3a6e300df70c56a6de170a368c Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Fri, 11 Mar 2016 14:26:07 -0800
Subject: [PATCH] Wait for stdio to close

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 containerd-shim/main.go    |  1 +
 containerd-shim/process.go | 23 ++++++++++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/containerd-shim/main.go b/containerd-shim/main.go
index a648317..c1b78ca 100644
--- a/containerd-shim/main.go
+++ b/containerd-shim/main.go
@@ -117,6 +117,7 @@ func main() {
 		}
 		// runtime has exited so the shim can also exit
 		if exitShim {
+			p.Wait()
 			return
 		}
 	}
diff --git a/containerd-shim/process.go b/containerd-shim/process.go
index c87d2ed..5f73a3f 100644
--- a/containerd-shim/process.go
+++ b/containerd-shim/process.go
@@ -9,6 +9,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strconv"
+	"sync"
 	"syscall"
 
 	"github.com/docker/containerd/runtime"
@@ -16,6 +17,7 @@ import (
 )
 
 type process struct {
+	sync.WaitGroup
 	id           string
 	bundle       string
 	stdio        *stdio
@@ -131,7 +133,12 @@ func (p *process) start() error {
 	cmd.SysProcAttr = &syscall.SysProcAttr{
 		Pdeathsig: syscall.SIGKILL,
 	}
-	if err := cmd.Run(); err != nil {
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+	p.stdio.stdout.Close()
+	p.stdio.stderr.Close()
+	if err := cmd.Wait(); err != nil {
 		return err
 	}
 	data, err := ioutil.ReadFile("pid")
@@ -184,9 +191,11 @@ func (p *process) openIO() error {
 		if err != nil {
 			return err
 		}
+		p.Add(1)
 		go func() {
 			io.Copy(stdout, console)
 			console.Close()
+			p.Done()
 		}()
 		return nil
 	}
@@ -201,10 +210,18 @@ func (p *process) openIO() error {
 			go io.Copy(i.Stdin, f)
 		},
 		p.state.Stdout: func(f *os.File) {
-			go io.Copy(f, i.Stdout)
+			p.Add(1)
+			go func() {
+				io.Copy(f, i.Stdout)
+				p.Done()
+			}()
 		},
 		p.state.Stderr: func(f *os.File) {
-			go io.Copy(f, i.Stderr)
+			p.Add(1)
+			go func() {
+				io.Copy(f, i.Stderr)
+				p.Done()
+			}()
 		},
 	} {
 		f, err := os.OpenFile(name, syscall.O_RDWR, 0)