Return correct number of PIDs with multi-container

Updates #172

PiperOrigin-RevId: 552620987

pkg/sentry/kernel/threads.go

@@ -302,6 +302,20 @@ func (ns *PIDNamespace) NumTasks() int {
 	return len(ns.tids)
 }
 
+// NumTasksPerContainer returns the number of tasks in ns that belongs to given container.
+func (ns *PIDNamespace) NumTasksPerContainer(cid string) int {
+	ns.owner.mu.RLock()
+	defer ns.owner.mu.RUnlock()
+
+	tasks := 0
+	for t := range ns.tids {
+		if t.ContainerID() == cid {
+			tasks++
+		}
+	}
+	return tasks
+}
+
 // ThreadGroups returns a snapshot of the thread groups in ns.
 func (ns *PIDNamespace) ThreadGroups() []*ThreadGroup {
 	return ns.ThreadGroupsAppend(nil)

runsc/boot/events.go

@@ -91,6 +91,13 @@ func (cm *containerManager) Event(cid *string, out *EventOut) error {
 		},
 	}
 
+	// PIDs and check that container exists before going further.
+	pids, err := cm.l.pidsCount(*cid)
+	if err != nil {
+		return err
+	}
+	out.Event.Data.Pids.Current = uint64(pids)
+
 	// Memory usage.
 	mem := cm.l.k.MemoryFile()
 	_ = mem.UpdateUsage() // best effort to update.
@@ -118,10 +125,6 @@ func (cm *containerManager) Event(cid *string, out *EventOut) error {
 
 	out.Event.Data.Memory.Usage.Usage = totalUsage
 
-	// PIDs.
-	// TODO(gvisor.dev/issue/172): Per-container accounting.
-	out.Event.Data.Pids.Current = uint64(len(cm.l.k.TaskSet().Root.ThreadGroups()))
-
 	// CPU usage by container.
 	out.ContainerUsage = control.ContainerUsage(cm.l.k)
 

runsc/boot/loader.go

@@ -1623,3 +1623,14 @@ func (l *Loader) containerCount() int {
 	}
 	return containers
 }
+
+func (l *Loader) pidsCount(cid string) (int, error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if _, err := l.tryThreadGroupFromIDLocked(execID{cid: cid}); err != nil {
+		// Container doesn't exist.
+		return 0, err
+	}
+	return l.k.TaskSet().Root.NumTasksPerContainer(cid), nil
+}

runsc/container/multi_container_test.go

@@ -2073,7 +2073,7 @@ func TestMultiContainerEvent(t *testing.T) {
 	conf.RootDir = rootDir
 
 	// Setup the containers.
-	sleep := []string{"/bin/sleep", "100"}
+	sleep := []string{"/bin/sh", "-c", "/bin/sleep 100 | grep 123"}
 	busy := []string{"/bin/bash", "-c", "i=0 ; while true ; do (( i += 1 )) ; done"}
 	quick := []string{"/bin/true"}
 	podSpec, ids := createSpecs(sleep, busy, quick)
@@ -2087,23 +2087,16 @@ func TestMultiContainerEvent(t *testing.T) {
 	t.Logf("Running container busy %s", containers[1].ID)
 	t.Logf("Running container quick %s", containers[2].ID)
 
-	// Wait for last container to stabilize the process count that is
-	// checked further below.
-	if ws, err := containers[2].Wait(); err != nil || ws != 0 {
-		t.Fatalf("Container.Wait, status: %v, err: %v", ws, err)
-	}
-	expectedPL := []*control.Process{
-		newProcessBuilder().Cmd("sleep").Process(),
-	}
-	if err := waitForProcessList(containers[0], expectedPL); err != nil {
+	// Wait for containers to start (last container should complete).
+	if err := waitForProcessCount(containers[0], 3); err != nil {
 		t.Errorf("failed to wait for sleep to start: %v", err)
 	}
-	expectedPL = []*control.Process{
-		newProcessBuilder().Cmd("bash").Process(),
-	}
-	if err := waitForProcessList(containers[1], expectedPL); err != nil {
+	if err := waitForProcessCount(containers[1], 1); err != nil {
 		t.Errorf("failed to wait for bash to start: %v", err)
 	}
+	if ws, err := containers[2].Wait(); err != nil || ws != 0 {
+		t.Fatalf("Container.Wait, status: %v, err: %v", ws, err)
+	}
 
 	// Check events for running containers.
 	for i, cont := range containers[:2] {
@@ -2118,9 +2111,14 @@ func TestMultiContainerEvent(t *testing.T) {
 		if cont.ID != evt.ID {
 			t.Errorf("Wrong container ID, want: %s, got: %s", cont.ID, evt.ID)
 		}
-		// One process per remaining container.
-		if got, want := evt.Data.Pids.Current, uint64(2); got != want {
-			t.Errorf("Wrong number of PIDs, cid: %q, want: %d, got: %d", cont.ID, want, got)
+
+		// container[0] expects 3 processes, while container[1] expects 1.
+		wantPids := 3
+		if i == 1 {
+			wantPids = 1
+		}
+		if got := evt.Data.Pids.Current; got != uint64(wantPids) {
+			t.Errorf("Wrong number of PIDs, cid: %q, want: %d, got: %d", cont.ID, wantPids, got)
 		}
 
 		switch i {