Cache ThreadGroups in PIDNamespace

If there are thousands of threads, ThreadGroupsAppend becomes very
expensive as it must iterate over all Tasks to find the ThreadGroup
leaders.

Reduce the cost by maintaining a map of ThreadGroups which can be used
to grab them all directly.

The one somewhat visible change is to convert PID namespace init
children zapping to a group-directed SIGKILL, as Linux did in
82058d668465 "signal: Use group_send_sig_info to kill all processes in a
pid namespace".

In a benchmark that creates N threads which sleep for two minutes, we
see approximately this much CPU time in ThreadGroupsAppend:

Before:

1 thread: 0ms
1024 threads: 30ms - 9130ms
4096 threads: 50ms - 2000ms
8192 threads: 18160ms
16384 threads: 17210ms

After:

1 thread: 0ms
1024 threads: 0ms
4096 threads: 0ms
8192 threads: 0ms
16384 threads: 0ms

The profiling is actually extremely noisy (likely due to cache effects),
as some runs show almost no samples at 1024, 4096 threads, but obviously
this does not scale to lots of threads.

PiperOrigin-RevId: 241828039
Change-Id: I17827c90045df4b3c49b3174f3a05bca3026a72c

pkg/sentry/kernel/kernel.go

@@ -866,14 +866,14 @@ func (k *Kernel) SendContainerSignal(cid string, info *arch.SignalInfo) error {
 	defer k.tasks.mu.RUnlock()
 
 	var lastErr error
-	for t := range k.tasks.Root.tids {
-		if t == t.tg.leader && t.ContainerID() == cid {
-			t.tg.signalHandlers.mu.Lock()
-			defer t.tg.signalHandlers.mu.Unlock()
+	for tg := range k.tasks.Root.tgids {
+		if tg.leader.ContainerID() == cid {
+			tg.signalHandlers.mu.Lock()
 			infoCopy := *info
-			if err := t.sendSignalLocked(&infoCopy, true /*group*/); err != nil {
+			if err := tg.leader.sendSignalLocked(&infoCopy, true /*group*/); err != nil {
 				lastErr = err
 			}
+			tg.signalHandlers.mu.Unlock()
 		}
 	}
 	return lastErr

pkg/sentry/kernel/sessions.go

@@ -240,14 +240,14 @@ func (pg *ProcessGroup) SendSignal(info *arch.SignalInfo) error {
 	defer tasks.mu.RUnlock()
 
 	var lastErr error
-	for t := range tasks.Root.tids {
-		if t == t.tg.leader && t.tg.ProcessGroup() == pg {
-			t.tg.signalHandlers.mu.Lock()
-			defer t.tg.signalHandlers.mu.Unlock()
+	for tg := range tasks.Root.tgids {
+		if tg.ProcessGroup() == pg {
+			tg.signalHandlers.mu.Lock()
 			infoCopy := *info
-			if err := t.sendSignalLocked(&infoCopy, true /*group*/); err != nil {
+			if err := tg.leader.sendSignalLocked(&infoCopy, true /*group*/); err != nil {
 				lastErr = err
 			}
+			tg.signalHandlers.mu.Unlock()
 		}
 	}
 	return lastErr
@@ -268,7 +268,7 @@ func (tg *ThreadGroup) CreateSession() error {
 // Precondition: callers must hold TaskSet.mu for writing.
 func (tg *ThreadGroup) createSession() error {
 	// Get the ID for this thread in the current namespace.
-	id := tg.pidns.tids[tg.leader]
+	id := tg.pidns.tgids[tg]
 
 	// Check if this ThreadGroup already leads a Session, or
 	// if the proposed group is already taken.
@@ -337,7 +337,7 @@ func (tg *ThreadGroup) createSession() error {
 
 	// Ensure a translation is added to all namespaces.
 	for ns := tg.pidns; ns != nil; ns = ns.parent {
-		local := ns.tids[tg.leader]
+		local := ns.tgids[tg]
 		ns.sids[s] = SessionID(local)
 		ns.sessions[SessionID(local)] = s
 		ns.pgids[pg] = ProcessGroupID(local)
@@ -356,7 +356,7 @@ func (tg *ThreadGroup) CreateProcessGroup() error {
 	defer tg.pidns.owner.mu.Unlock()
 
 	// Get the ID for this thread in the current namespace.
-	id := tg.pidns.tids[tg.leader]
+	id := tg.pidns.tgids[tg]
 
 	// Per above, check for a Session leader or existing group.
 	for s := tg.pidns.owner.sessions.Front(); s != nil; s = s.Next() {
@@ -401,7 +401,7 @@ func (tg *ThreadGroup) CreateProcessGroup() error {
 
 	// Ensure this translation is added to all namespaces.
 	for ns := tg.pidns; ns != nil; ns = ns.parent {
-		local := ns.tids[tg.leader]
+		local := ns.tgids[tg]
 		ns.pgids[pg] = ProcessGroupID(local)
 		ns.processGroups[ProcessGroupID(local)] = pg
 	}

pkg/sentry/kernel/task_exec.go

@@ -234,6 +234,8 @@ func (t *Task) promoteLocked() {
 		ns.tids[t] = leaderTID
 		ns.tasks[oldTID] = oldLeader
 		ns.tasks[leaderTID] = t
+		// Neither the ThreadGroup nor TGID change, so no need to
+		// update ns.tgids.
 	}
 
 	// Inherit the old leader's start time.

pkg/sentry/kernel/task_exit.go

@@ -329,14 +329,15 @@ func (t *Task) exitChildren() {
 		// signal." - pid_namespaces(7)
 		t.Debugf("Init process terminating, killing namespace")
 		t.tg.pidns.exiting = true
-		for other := range t.tg.pidns.tids {
-			if other.tg != t.tg {
-				other.tg.signalHandlers.mu.Lock()
-				other.sendSignalLocked(&arch.SignalInfo{
-					Signo: int32(linux.SIGKILL),
-				}, false /* group */)
-				other.tg.signalHandlers.mu.Unlock()
+		for other := range t.tg.pidns.tgids {
+			if other == t.tg {
+				continue
 			}
+			other.signalHandlers.mu.Lock()
+			other.leader.sendSignalLocked(&arch.SignalInfo{
+				Signo: int32(linux.SIGKILL),
+			}, true /* group */)
+			other.signalHandlers.mu.Unlock()
 		}
 		// TODO: The init process waits for all processes in the
 		// namespace to exit before completing its own exit
@@ -652,6 +653,9 @@ func (t *Task) exitNotifyLocked(fromPtraceDetach bool) {
 			tid := ns.tids[t]
 			delete(ns.tasks, tid)
 			delete(ns.tids, t)
+			if t == t.tg.leader {
+				delete(ns.tgids, t.tg)
+			}
 		}
 		t.tg.exitedCPUStats.Accumulate(t.CPUStats())
 		t.tg.ioUsage.Accumulate(t.ioUsage)

pkg/sentry/kernel/task_start.go

@@ -212,11 +212,18 @@ func (ts *TaskSet) assignTIDsLocked(t *Task) error {
 			for _, a := range allocatedTIDs {
 				delete(a.ns.tasks, a.tid)
 				delete(a.ns.tids, t)
+				if t.tg.leader == nil {
+					delete(a.ns.tgids, t.tg)
+				}
 			}
 			return err
 		}
 		ns.tasks[tid] = t
 		ns.tids[t] = tid
+		if t.tg.leader == nil {
+			// New thread group.
+			ns.tgids[t.tg] = tid
+		}
 		allocatedTIDs = append(allocatedTIDs, allocatedTID{ns, tid})
 	}
 	return nil

pkg/sentry/kernel/threads.go

@@ -100,10 +100,8 @@ func newTaskSet() *TaskSet {
 //
 // Preconditions: ts.mu must be locked (for reading or writing).
 func (ts *TaskSet) forEachThreadGroupLocked(f func(tg *ThreadGroup)) {
-	for t := range ts.Root.tids {
-		if t == t.tg.leader {
-			f(t.tg)
-		}
+	for tg := range ts.Root.tgids {
+		f(tg)
 	}
 }
 
@@ -145,6 +143,13 @@ type PIDNamespace struct {
 	// identifiers in this namespace.
 	tids map[*Task]ThreadID
 
+	// tgids is a mapping from thread groups visible in this namespace to
+	// their identifiers in this namespace.
+	//
+	// The content of tgids is equivalent to tids[tg.leader]. This exists
+	// primarily as an optimization to quickly find all thread groups.
+	tgids map[*ThreadGroup]ThreadID
+
 	// sessions is a mapping from SessionIDs in this namespace to sessions
 	// visible in the namespace.
 	sessions map[SessionID]*Session
@@ -173,6 +178,7 @@ func newPIDNamespace(ts *TaskSet, parent *PIDNamespace, userns *auth.UserNamespa
 		userns:        userns,
 		tasks:         make(map[ThreadID]*Task),
 		tids:          make(map[*Task]ThreadID),
+		tgids:         make(map[*ThreadGroup]ThreadID),
 		sessions:      make(map[SessionID]*Session),
 		sids:          make(map[*Session]SessionID),
 		processGroups: make(map[ProcessGroupID]*ProcessGroup),
@@ -227,7 +233,7 @@ func (ns *PIDNamespace) IDOfTask(t *Task) ThreadID {
 func (ns *PIDNamespace) IDOfThreadGroup(tg *ThreadGroup) ThreadID {
 	ns.owner.mu.RLock()
 	defer ns.owner.mu.RUnlock()
-	return ns.tids[tg.leader]
+	return ns.tgids[tg]
 }
 
 // Tasks returns a snapshot of the tasks in ns.
@@ -250,10 +256,8 @@ func (ns *PIDNamespace) ThreadGroups() []*ThreadGroup {
 func (ns *PIDNamespace) ThreadGroupsAppend(tgs []*ThreadGroup) []*ThreadGroup {
 	ns.owner.mu.RLock()
 	defer ns.owner.mu.RUnlock()
-	for t := range ns.tids {
-		if t == t.tg.leader {
-			tgs = append(tgs, t.tg)
-		}
+	for tg := range ns.tgids {
+		tgs = append(tgs, tg)
 	}
 	return tgs
 }
@@ -387,7 +391,7 @@ func (tg *ThreadGroup) MemberIDs(pidns *PIDNamespace) []ThreadID {
 func (tg *ThreadGroup) ID() ThreadID {
 	tg.pidns.owner.mu.RLock()
 	defer tg.pidns.owner.mu.RUnlock()
-	return tg.pidns.tids[tg.leader]
+	return tg.pidns.tgids[tg]
 }
 
 // A taskNode defines the relationship between a task and the rest of the