From 4641ab4101395c6189cc033d8887d147f7f35d3a Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Thu, 6 Oct 2016 11:46:03 -0700
Subject: [PATCH] Fix process wait with monitor

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 shim/process.go |  4 +---
 shim/shim.go    | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/shim/process.go b/shim/process.go
index 3098fc7..2e2a212 100644
--- a/shim/process.go
+++ b/shim/process.go
@@ -200,9 +200,7 @@ func (p *process) FD() int {
 }
 
 func (p *process) Wait() (rst uint32, rerr error) {
-	if _, err := ioutil.ReadAll(p.exit); err != nil {
-		return 255, err
-	}
+	<-p.done
 	data, err := ioutil.ReadFile(filepath.Join(p.root, "exitStatus"))
 	defer func() {
 		if rerr != nil {
diff --git a/shim/shim.go b/shim/shim.go
index 8c9395a..8d771e0 100644
--- a/shim/shim.go
+++ b/shim/shim.go
@@ -12,6 +12,8 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/Sirupsen/logrus"
+	"github.com/docker/containerd/monitor"
 	"github.com/docker/containerd/oci"
 	"github.com/docker/containerkit"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
@@ -62,13 +64,19 @@ func New(opts Opts) (*Shim, error) {
 	if err != nil {
 		return nil, err
 	}
+	m, err := monitor.New()
+	if err != nil {
+		return nil, err
+	}
 	s := &Shim{
 		root:      opts.Root,
 		name:      opts.Name,
 		timeout:   opts.Timeout,
 		runtime:   r,
 		processes: make(map[string]*process),
+		m:         m,
 	}
+	go s.startMonitor()
 	f, err := os.Create(filepath.Join(opts.Root, "state.json"))
 	if err != nil {
 		return nil, err
@@ -131,6 +139,7 @@ type Shim struct {
 	processes   map[string]*process
 	bundle      string
 	checkpoint  string
+	m           *monitor.Monitor
 }
 
 type state struct {
@@ -198,6 +207,9 @@ func (s *Shim) Create(c *containerkit.Container) (containerkit.ProcessDelegate,
 	if err != nil {
 		return nil, err
 	}
+	if err := s.m.Add(p); err != nil {
+		return nil, err
+	}
 	s.pmu.Lock()
 	s.processes["init"] = p
 	s.pmu.Unlock()
@@ -302,6 +314,16 @@ func (s *Shim) command(args ...string) *exec.Cmd {
 	return exec.Command(s.name, args...)
 }
 
+func (s *Shim) startMonitor() {
+	for m := range s.m.Events() {
+		p := m.(*process)
+		close(p.done)
+		if err := s.m.Remove(p); err != nil {
+			logrus.Error(err)
+		}
+	}
+}
+
 // checkShimNotFound checks the error returned from a exec call to see if the binary
 // that was called exists on the system and returns true if the shim binary does not exist
 func checkShimNotFound(err error) bool {