Merge Loader.containerRootTGs and execProcess into a single map

It's easier to manage a single map with processes that we're interested
to track. This will make the next change to clean up the map on destroy
easier.

PiperOrigin-RevId: 214894210
Change-Id: I099247323a0487cd0767120df47ba786fac0926d

2 files changed, 43 insertions, 46 deletions

diff --git a/runsc/boot/controller.go b/runsc/boot/controller.go
index 116a8369c..362e74df5 100644
--- a/runsc/boot/controller.go
+++ b/runsc/boot/controller.go
@@ -250,7 +250,8 @@ func (cm *containerManager) Destroy(cid *string, _ *struct{}) error {
 	cm.l.mu.Lock()
 	defer cm.l.mu.Unlock()
 
-	if tg, ok := cm.l.containerRootTGs[*cid]; ok {
+	key := execID{cid: *cid}
+	if tg, ok := cm.l.processes[key]; ok {
 		// Send SIGKILL to threadgroup.
 		if err := tg.SendSignal(&arch.SignalInfo{
 			Signo: int32(linux.SIGKILL),
@@ -265,7 +266,7 @@ func (cm *containerManager) Destroy(cid *string, _ *struct{}) error {
 		}
 
 		// Remove the container thread group from the map.
-		delete(cm.l.containerRootTGs, *cid)
+		delete(cm.l.processes, key)
 	}
 
 	// Clean up the filesystem by unmounting all mounts for this container
@@ -379,9 +380,9 @@ type RestoreOpts struct {
 }
 
 // Restore loads a container from a statefile.
-// The container's current kernel is destroyed, a restore environment is created,
-// and the kernel is recreated with the restore state file. The container then sends the
-// signal to start.
+// The container's current kernel is destroyed, a restore environment is
+// created, and the kernel is recreated with the restore state file. The
+// container then sends the signal to start.
 func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 	log.Debugf("containerManager.Restore")
 
@@ -455,12 +456,20 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 	cm.l.rootProcArgs = kernel.CreateProcessArgs{}
 	cm.l.restore = true
 
+	// Reinitialize the sandbox ID and processes map. Note that it doesn't
+	// restore the state of multiple containers, nor exec processes.
+	cm.l.sandboxID = o.SandboxID
+	cm.l.mu.Lock()
+	key := execID{cid: o.SandboxID}
+	cm.l.processes = map[execID]*kernel.ThreadGroup{key: cm.l.k.GlobalInit()}
+	cm.l.mu.Unlock()
+
 	// Tell the root container to start and wait for the result.
 	cm.startChan <- struct{}{}
 	if err := <-cm.startResultChan; err != nil {
 		return fmt.Errorf("failed to start sandbox: %v", err)
 	}
-	cm.l.setRootContainerID(o.SandboxID)
+
 	return nil
 }
 
@@ -511,6 +520,6 @@ type SignalArgs struct {
 // Signal sends a signal to the init process of the container.
 // TODO: Send signal to exec process.
 func (cm *containerManager) Signal(args *SignalArgs, _ *struct{}) error {
-	log.Debugf("containerManager.Signal")
+	log.Debugf("containerManager.Signal %q %d, all: %t", args.CID, args.Signo, args.All)
 	return cm.l.signal(args.CID, args.Signo, args.All)
 }
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index bd6e146fc..52c251812 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -104,26 +104,18 @@ type Loader struct {
 	// sandboxID is the ID for the whole sandbox.
 	sandboxID string
 
-	// mu guards containerRootTGs and execProcesses.
+	// mu guards processes.
 	mu sync.Mutex
 
-	// containerRootTGs maps container IDs to their root processes. It
-	// can be used to determine which process to manipulate when clients
-	// call methods on particular containers.
+	// processes maps containers root process and invocation of exec. Root
+	// processes are keyed with container ID and pid=0, while exec invocations
+	// have the corresponding pid set.
 	//
-	// containerRootTGs is guarded by mu.
+	// processes is guardded by mu.
 	//
 	// TODO: When containers are removed via `runsc delete`,
-	// containerRootTGs should be cleaned up.
-	containerRootTGs map[string]*kernel.ThreadGroup
-
-	// execProcesses maps each invocation of exec to the process it spawns.
-	//
-	// execProcesses is guardded by mu.
-	//
-	// TODO: When containers are removed via `runsc delete`,
-	// execProcesses should be cleaned up.
-	execProcesses map[execID]*kernel.ThreadGroup
+	// processes should be cleaned up.
+	processes map[execID]*kernel.ThreadGroup
 }
 
 // execID uniquely identifies a sentry process.
@@ -287,8 +279,7 @@ func New(id string, spec *specs.Spec, conf *Config, controllerFD, deviceFD int,
 		startSignalForwarding: startSignalForwarding,
 		rootProcArgs:          procArgs,
 		sandboxID:             id,
-		containerRootTGs:      make(map[string]*kernel.ThreadGroup),
-		execProcesses:         make(map[execID]*kernel.ThreadGroup),
+		processes:             make(map[execID]*kernel.ThreadGroup),
 	}
 	ctrl.manager.l = l
 	return l, nil
@@ -425,7 +416,8 @@ func (l *Loader) run() error {
 	}
 
 	l.mu.Lock()
-	l.containerRootTGs[l.sandboxID] = l.k.GlobalInit()
+	key := execID{cid: l.sandboxID}
+	l.processes[key] = l.k.GlobalInit()
 	l.mu.Unlock()
 
 	// Start signal forwarding only after an init process is created.
@@ -521,7 +513,8 @@ func (l *Loader) startContainer(k *kernel.Kernel, spec *specs.Spec, conf *Config
 
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	l.containerRootTGs[cid] = tg
+	key := execID{cid: cid}
+	l.processes[key] = tg
 
 	return nil
 }
@@ -530,7 +523,8 @@ func (l *Loader) executeAsync(args *control.ExecArgs) (kernel.ThreadID, error) {
 	// Get the container Root Dirent from the Task, since we must run this
 	// process with the same Root.
 	l.mu.Lock()
-	tg, ok := l.containerRootTGs[args.ContainerID]
+	rootKey := execID{cid: args.ContainerID}
+	tg, ok := l.processes[rootKey]
 	l.mu.Unlock()
 	if !ok {
 		return 0, fmt.Errorf("cannot exec in container %q: no such container", args.ContainerID)
@@ -549,13 +543,13 @@ func (l *Loader) executeAsync(args *control.ExecArgs) (kernel.ThreadID, error) {
 		return 0, fmt.Errorf("error executing: %+v: %v", args, err)
 	}
 
-	// Insert the process into execProcesses so that we can wait on it
+	// Insert the process into processes so that we can wait on it
 	// later.
 	l.mu.Lock()
 	defer l.mu.Unlock()
 	eid := execID{cid: args.ContainerID, pid: tgid}
-	l.execProcesses[eid] = tg
-	log.Debugf("updated execProcesses: %v", l.execProcesses)
+	l.processes[eid] = tg
+	log.Debugf("updated processes: %v", l.processes)
 
 	return tgid, nil
 }
@@ -567,12 +561,12 @@ 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()
-	tg, ok := l.containerRootTGs[cid]
+	key := execID{cid: cid}
+	tg, ok := l.processes[key]
+	l.mu.Unlock()
 	if !ok {
-		defer l.mu.Unlock()
-		return fmt.Errorf("can't find process for container %q in %v", cid, l.containerRootTGs)
+		return fmt.Errorf("can't find process for container %q in %v", cid, l.processes)
 	}
-	l.mu.Unlock()
 
 	// If the thread either has already exited or exits during waiting,
 	// consider the container exited.
@@ -590,10 +584,10 @@ func (l *Loader) waitPID(tgid kernel.ThreadID, cid string, clearStatus bool, wai
 	}*/
 
 	// If the process was started via runsc exec, it will have an
-	// entry in l.execProcesses.
+	// entry in l.processes.
 	l.mu.Lock()
 	eid := execID{cid: cid, pid: tgid}
-	tg, ok := l.execProcesses[eid]
+	tg, ok := l.processes[eid]
 	l.mu.Unlock()
 	if ok {
 		ws := l.wait(tg)
@@ -601,8 +595,8 @@ func (l *Loader) waitPID(tgid kernel.ThreadID, cid string, clearStatus bool, wai
 		if clearStatus {
 			// Remove tg from the cache.
 			l.mu.Lock()
-			delete(l.execProcesses, eid)
-			log.Debugf("updated execProcesses (removal): %v", l.execProcesses)
+			delete(l.processes, eid)
+			log.Debugf("updated processes (removal): %v", l.processes)
 			l.mu.Unlock()
 		}
 		return nil
@@ -626,13 +620,6 @@ func (l *Loader) wait(tg *kernel.ThreadGroup) uint32 {
 	return tg.ExitStatus().Status()
 }
 
-func (l *Loader) setRootContainerID(cid string) {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-	l.containerRootTGs = map[string]*kernel.ThreadGroup{cid: l.k.GlobalInit()}
-	l.sandboxID = cid
-}
-
 // WaitForStartSignal waits for a start signal from the control server.
 func (l *Loader) WaitForStartSignal() {
 	<-l.ctrl.manager.startChan
@@ -674,7 +661,8 @@ func newEmptyNetworkStack(conf *Config, clock tcpip.Clock) (inet.Stack, error) {
 
 func (l *Loader) signal(cid string, signo int32, all bool) error {
 	l.mu.Lock()
-	tg, ok := l.containerRootTGs[cid]
+	key := execID{cid: cid}
+	tg, ok := l.processes[key]
 	l.mu.Unlock()
 	if !ok {
 		return fmt.Errorf("failed to signal container %q: no such container", cid)