runsc: Enable waiting on individual containers within a sandbox.

PiperOrigin-RevId: 201742160
Change-Id: Ia9fa1442287c5f9e1196fb117c41536a80f6bb31

6 files changed, 172 insertions, 20 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)
 	}
 
diff --git a/runsc/container/container.go b/runsc/container/container.go
index 604708e2c..9c0169ca8 100644
--- a/runsc/container/container.go
+++ b/runsc/container/container.go
@@ -293,7 +293,7 @@ func (c *Container) Start(conf *boot.Config) error {
 			return err
 		}
 	} else {
-		if err := c.Sandbox.Start(c.Spec, conf); err != nil {
+		if err := c.Sandbox.Start(c.Spec, conf, c.ID); err != nil {
 			c.Destroy()
 			return err
 		}
diff --git a/runsc/container/container_test.go b/runsc/container/container_test.go
index a8320f614..de487ea97 100644
--- a/runsc/container/container_test.go
+++ b/runsc/container/container_test.go
@@ -1020,3 +1020,92 @@ func TestMultiContainerSanity(t *testing.T) {
 		t.Errorf("failed to wait for sleep to start: %v", err)
 	}
 }
+
+func TestMultiContainerWait(t *testing.T) {
+	containerIDs := []string{
+		testutil.UniqueContainerID(),
+		testutil.UniqueContainerID(),
+	}
+	containerAnnotations := []map[string]string{
+		// The first container creates a sandbox.
+		map[string]string{
+			specutils.ContainerdContainerTypeAnnotation: specutils.ContainerdContainerTypeSandbox,
+		},
+		// The second container creates a container within the first
+		// container's sandbox.
+		map[string]string{
+			specutils.ContainerdContainerTypeAnnotation: specutils.ContainerdContainerTypeContainer,
+			specutils.ContainerdSandboxIDAnnotation:     containerIDs[0],
+		},
+	}
+	args := [][]string{
+		// The first container should run the entire duration of the
+		// test.
+		{"sleep", "100"},
+		// We'll wait on the second container, which is much shorter
+		// lived.
+		{"sleep", "1"},
+	}
+
+	rootDir, err := testutil.SetupRootDir()
+	if err != nil {
+		t.Fatalf("error creating root dir: %v", err)
+	}
+	defer os.RemoveAll(rootDir)
+
+	// Setup the containers.
+	containers := make([]*container.Container, 0, len(containerIDs))
+	for i, annotations := range containerAnnotations {
+		spec := testutil.NewSpecWithArgs(args[i][0], args[i][1])
+		spec.Annotations = annotations
+		bundleDir, conf, err := testutil.SetupContainerInRoot(rootDir, spec)
+		if err != nil {
+			t.Fatalf("error setting up container: %v", err)
+		}
+		defer os.RemoveAll(bundleDir)
+		cont, err := container.Create(containerIDs[i], spec, conf, bundleDir, "", "", "")
+		if err != nil {
+			t.Fatalf("error creating container: %v", err)
+		}
+		defer cont.Destroy()
+		if err := cont.Start(conf); err != nil {
+			t.Fatalf("error starting container: %v", err)
+		}
+		containers = append(containers, cont)
+	}
+
+	expectedPL := []*control.Process{
+		{
+			UID:  0,
+			PID:  1,
+			PPID: 0,
+			C:    0,
+			Cmd:  "sleep",
+		},
+		{
+			UID:  0,
+			PID:  2,
+			PPID: 0,
+			C:    0,
+			Cmd:  "sleep",
+		},
+	}
+
+	// Check via ps that multiple processes are running.
+	if err := waitForProcessList(containers[0], expectedPL); err != nil {
+		t.Errorf("failed to wait for sleep to start: %v", err)
+	}
+
+	// Wait on the short lived container.
+	if ws, err := containers[1].Wait(); err != nil {
+		t.Fatalf("failed to wait for process %q: %v", strings.Join(containers[1].Spec.Process.Args, " "), err)
+	} else if es := ws.ExitStatus(); es != 0 {
+		t.Fatalf("process %q exited with non-zero status %d", strings.Join(containers[1].Spec.Process.Args, " "), es)
+	}
+
+	// After Wait returns, ensure that the root container is running and
+	// the child has finished.
+	if err := waitForProcessList(containers[0], expectedPL[:1]); err != nil {
+		t.Errorf("failed to wait for %q to start: %v", strings.Join(containers[0].Spec.Process.Args, " "), err)
+	}
+}
diff --git a/runsc/sandbox/sandbox.go b/runsc/sandbox/sandbox.go
index 870a0ccd3..ed2c40e57 100644
--- a/runsc/sandbox/sandbox.go
+++ b/runsc/sandbox/sandbox.go
@@ -99,7 +99,7 @@ func (s *Sandbox) StartRoot(spec *specs.Spec, conf *boot.Config) error {
 
 	// Send a message to the sandbox control server to start the root
 	// container.
-	if err := conn.Call(boot.RootContainerStart, nil, nil); err != nil {
+	if err := conn.Call(boot.RootContainerStart, &s.ID, nil); err != nil {
 		return fmt.Errorf("error starting root container %v: %v", spec.Process.Args, err)
 	}
 
@@ -107,7 +107,7 @@ func (s *Sandbox) StartRoot(spec *specs.Spec, conf *boot.Config) error {
 }
 
 // Start starts running a non-root container inside the sandbox.
-func (s *Sandbox) Start(spec *specs.Spec, conf *boot.Config) error {
+func (s *Sandbox) Start(spec *specs.Spec, conf *boot.Config, cid string) error {
 	log.Debugf("Start non-root container sandbox %q, pid: %d", s.ID, s.Pid)
 	conn, err := s.connect()
 	if err != nil {
@@ -118,6 +118,7 @@ func (s *Sandbox) Start(spec *specs.Spec, conf *boot.Config) error {
 	args := boot.StartArgs{
 		Spec: spec,
 		Conf: conf,
+		CID:  cid,
 	}
 	if err := conn.Call(boot.ContainerStart, args, nil); err != nil {
 		return fmt.Errorf("error starting non-root container %v: %v", spec.Process.Args, err)