2 files changed, 58 insertions, 12 deletions

diff --git a/runsc/boot/controller.go b/runsc/boot/controller.go
index 56829c605..ff75a382e 100644
--- a/runsc/boot/controller.go
+++ b/runsc/boot/controller.go
@@ -61,6 +61,10 @@ const (
 	// and return its ExitStatus.
 	ContainerWait = "containerManager.Wait"
 
+	// ContainerWaitPID is used to wait on a process with a certain PID in
+	// the sandbox and return its ExitStatus.
+	ContainerWaitPID = "containerManager.WaitPID"
+
 	// NetworkCreateLinksAndRoutes is the URPC endpoint for creating links
 	// and routes in a network stack.
 	NetworkCreateLinksAndRoutes = "Network.CreateLinksAndRoutes"
@@ -234,7 +238,22 @@ func (cm *containerManager) Resume(_, _ *struct{}) error {
 // Wait waits for the init process in the given container.
 func (cm *containerManager) Wait(cid *string, waitStatus *uint32) error {
 	log.Debugf("containerManager.Wait")
-	return cm.l.wait(cid, waitStatus)
+	return cm.l.waitContainer(*cid, waitStatus)
+}
+
+// WaitPIDArgs are arguments to the WaitPID method.
+type WaitPIDArgs struct {
+	// PID is the PID in the container's PID namespace.
+	PID int32
+
+	// CID is the container ID.
+	CID string
+}
+
+// WaitPID waits for the process with PID 'pid' in the sandbox.
+func (cm *containerManager) WaitPID(args *WaitPIDArgs, waitStatus *uint32) error {
+	log.Debugf("containerManager.Wait")
+	return cm.l.waitPID(kernel.ThreadID(args.PID), args.CID, waitStatus)
 }
 
 // SignalArgs are arguments to the Signal method.
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index f359a0eb0..014908179 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -83,6 +83,9 @@ type Loader struct {
 	// rootProcArgs refers to the root sandbox init task.
 	rootProcArgs kernel.CreateProcessArgs
 
+	// sandboxID is the ID for the whole sandbox.
+	sandboxID string
+
 	// mu guards containerRootTGIDs.
 	mu sync.Mutex
 
@@ -452,23 +455,46 @@ func (l *Loader) startContainer(args *StartArgs, k *kernel.Kernel) (kernel.Threa
 	return tgid, nil
 }
 
-// wait waits for the init process in the given container.
-func (l *Loader) wait(cid *string, waitStatus *uint32) error {
+// TODO: Per-container namespaces must be supported
+// for -pid.
+
+// waitContainer waits for the root process of a container to exit.
+func (l *Loader) waitContainer(cid string, waitStatus *uint32) error {
+	// Don't defer unlock, as doing so would make it impossible for
+	// multiple clients to wait on the same container.
 	l.mu.Lock()
-	defer l.mu.Unlock()
-	tgid, ok := l.containerRootTGIDs[*cid]
+	tgid, ok := l.containerRootTGIDs[cid]
+	l.mu.Unlock()
 	if !ok {
-		return fmt.Errorf("can't find process for container %q in %v", *cid, l.containerRootTGIDs)
+		return fmt.Errorf("can't find process for container %q in %v", cid, l.containerRootTGIDs)
 	}
-
-	// TODO: Containers don't map 1:1 with their root
-	// processes. Container exits should be managed explicitly
-	// rather than via PID.
 	// If the thread either has already exited or exits during waiting,
 	// consider the container exited.
-	defer delete(l.containerRootTGIDs, *cid)
+	defer func() {
+		l.mu.Lock()
+		defer l.mu.Unlock()
+		// TODO: Containers don't map 1:1 with their root
+		// processes. Container exits should be managed explicitly
+		// rather than via PID.
+		delete(l.containerRootTGIDs, cid)
+	}()
+	return l.wait(tgid, cid, waitStatus)
+}
+
+func (l *Loader) waitPID(tgid kernel.ThreadID, cid string, waitStatus *uint32) error {
+	// TODO: Containers all currently share a PID namespace.
+	// When per-container PID namespaces are supported, wait should use cid
+	// to find the appropriate PID namespace.
+	if cid != l.sandboxID {
+		return errors.New("non-sandbox PID namespaces are not yet implemented")
+	}
+	return l.wait(tgid, cid, waitStatus)
+}
 
-	tg := l.k.TaskSet().Root.ThreadGroupWithID(tgid)
+// wait waits for the process with TGID 'tgid' in a container's PID namespace
+// to exit.
+func (l *Loader) wait(tgid kernel.ThreadID, cid string, waitStatus *uint32) error {
+	tg := l.k.TaskSet().Root.ThreadGroupWithID(kernel.ThreadID(tgid))
 	if tg == nil {
 		return fmt.Errorf("no thread group with ID %d", tgid)
 	}
@@ -482,6 +508,7 @@ func (l *Loader) setRootContainerID(cid string) {
 	defer l.mu.Unlock()
 	// The root container has PID 1.
 	l.containerRootTGIDs = map[string]kernel.ThreadID{cid: 1}
+	l.sandboxID = cid
 }
 
 // WaitForStartSignal waits for a start signal from the control server.