3 files changed, 79 insertions, 17 deletions
diff --git a/runsc/boot/controller.go b/runsc/boot/controller.go
index ec24c4dad..56829c605 100644
--- a/runsc/boot/controller.go
+++ b/runsc/boot/controller.go
@@ -145,10 +145,11 @@ type containerManager struct {
 }
 
 // StartRoot will start the root container process.
-func (cm *containerManager) StartRoot(_, _ *struct{}) error {
+func (cm *containerManager) StartRoot(cid *string, _ *struct{}) error {
 	log.Debugf("containerManager.StartRoot")
 	// Tell the root container to start and wait for the result.
 	cm.startChan <- struct{}{}
+	cm.l.setRootContainerID(*cid)
 	return <-cm.startResultChan
 }
 
@@ -166,6 +167,9 @@ type StartArgs struct {
 	// TODO: Separate sandbox and container configs.
 	// Config is the runsc-specific configuration for the sandbox.
 	Conf *Config
+
+	// CID is the ID of the container to start.
+	CID string
 }
 
 // Start runs a created container within a sandbox.
@@ -182,8 +186,16 @@ func (cm *containerManager) Start(args *StartArgs, _ *struct{}) error {
 	if args.Conf == nil {
 		return errors.New("start arguments missing config")
 	}
+	if args.CID == "" {
+		return errors.New("start argument missing container ID")
+	}
+
+	tgid, err := cm.l.startContainer(args, cm.k)
+	if err != nil {
+		return err
+	}
+	log.Debugf("Container %q started with root PID of %d", args.CID, tgid)
 
-	cm.l.startContainer(args, cm.k)
 	return nil
 }
 
@@ -222,15 +234,7 @@ 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")
-	// TODO: Use the cid and wait on the init process in that
-	// container. Currently we just wait on PID 1 in the sandbox.
-	tg := cm.k.TaskSet().Root.ThreadGroupWithID(1)
-	if tg == nil {
-		return fmt.Errorf("cannot wait: no thread group with id 1")
-	}
-	tg.WaitExited()
-	*waitStatus = tg.ExitStatus().Status()
-	return nil
+	return cm.l.wait(cid, waitStatus)
 }
 
 // SignalArgs are arguments to the Signal method.
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index a0a28dc43..7097f220b 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -16,10 +16,12 @@
 package boot
 
 import (
+	"errors"
 	"fmt"
 	"math/rand"
 	"os"
 	"runtime"
+	"sync"
 	"sync/atomic"
 	"syscall"
 	gtime "time"
@@ -81,6 +83,16 @@ type Loader struct {
 
 	// rootProcArgs refers to the root sandbox init task.
 	rootProcArgs kernel.CreateProcessArgs
+
+	// mu guards containerRootTGIDs.
+	mu sync.Mutex
+
+	// containerRootTGIDs maps container IDs to their root processes. It
+	// can be used to determine which process to manipulate when clients
+	// call methods on particular containers.
+	//
+	// containerRootTGIDs is guarded by mu.
+	containerRootTGIDs map[string]kernel.ThreadID
 }
 
 func init() {
@@ -377,12 +389,14 @@ func (l *Loader) run() error {
 	return l.k.Start()
 }
 
-func (l *Loader) startContainer(args *StartArgs, k *kernel.Kernel) error {
+// startContainer starts a child container. It returns the thread group ID of
+// the newly created process.
+func (l *Loader) startContainer(args *StartArgs, k *kernel.Kernel) (kernel.ThreadID, error) {
 	spec := args.Spec
 	// Create capabilities.
 	caps, err := specutils.Capabilities(spec.Process.Capabilities)
 	if err != nil {
-		return fmt.Errorf("error creating capabilities: %v", err)
+		return 0, fmt.Errorf("error creating capabilities: %v", err)
 	}
 
 	// Convert the spec's additional GIDs to KGIDs.
@@ -416,19 +430,62 @@ func (l *Loader) startContainer(args *StartArgs, k *kernel.Kernel) error {
 		k.RootIPCNamespace(),
 		k)
 	if err != nil {
-		return fmt.Errorf("failed to create new process: %v", err)
+		return 0, fmt.Errorf("failed to create new process: %v", err)
+	}
+
+	tg, err := l.k.CreateProcess(procArgs)
+	if err != nil {
+		return 0, fmt.Errorf("failed to create process in sentry: %v", err)
 	}
 
-	if _, err := l.k.CreateProcess(procArgs); err != nil {
-		return fmt.Errorf("failed to create process in sentry: %v", err)
+	ts := k.TaskSet()
+	tgid := ts.Root.IDOfThreadGroup(tg)
+	if tgid == 0 {
+		return 0, errors.New("failed to get thread group ID of new process")
 	}
 
 	// CreateProcess takes a reference on FDMap if successful.
 	procArgs.FDMap.DecRef()
 
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.containerRootTGIDs[args.CID] = tgid
+
+	return tgid, nil
+}
+
+// wait waits for the init process in the given container.
+func (l *Loader) wait(cid *string, waitStatus *uint32) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	tgid, ok := l.containerRootTGIDs[*cid]
+	if !ok {
+		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)
+
+	tg := l.k.TaskSet().Root.ThreadGroupWithID(tgid)
+	if tg == nil {
+		return fmt.Errorf("no thread group with ID %d", tgid)
+	}
+	tg.WaitExited()
+	*waitStatus = tg.ExitStatus().Status()
 	return nil
 }
 
+func (l *Loader) setRootContainerID(cid string) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	// The root container has PID 1.
+	l.containerRootTGIDs = map[string]kernel.ThreadID{cid: 1}
+}
+
 // WaitForStartSignal waits for a start signal from the control server.
 func (l *Loader) WaitForStartSignal() {
 	<-l.ctrl.manager.startChan
diff --git a/runsc/boot/loader_test.go b/runsc/boot/loader_test.go
index 5ec1084db..15ced0601 100644
--- a/runsc/boot/loader_test.go
+++ b/runsc/boot/loader_test.go
@@ -133,7 +133,8 @@ func TestStartSignal(t *testing.T) {
 	}
 
 	// Trigger the control server StartRoot method.
-	if err := s.ctrl.manager.StartRoot(nil, nil); err != nil {
+	cid := "foo"
+	if err := s.ctrl.manager.StartRoot(&cid, nil); err != nil {
 		t.Errorf("error calling StartRoot: %v", err)
 	}