Refactor task start and exit from a PID namespace into separate functions.

PiperOrigin-RevId: 426083905

pkg/sentry/kernel/BUILD

@@ -200,6 +200,7 @@ go_library(
         "task_work.go",
         "thread_group.go",
         "threads.go",
+        "threads_impl.go",
         "timekeeper.go",
         "timekeeper_state.go",
         "tty.go",

pkg/sentry/kernel/task_exit.go

@@ -651,12 +651,7 @@ func (t *Task) exitNotifyLocked(fromPtraceDetach bool) {
 	if t.exitTracerAcked && t.exitParentAcked {
 		t.advanceExitStateLocked(TaskExitZombie, TaskExitDead)
 		for ns := t.tg.pidns; ns != nil; ns = ns.parent {
-			tid := ns.tids[t]
-			delete(ns.tasks, tid)
-			delete(ns.tids, t)
-			if t == t.tg.leader {
-				delete(ns.tgids, t.tg)
-			}
+			ns.deleteTask(t)
 		}
 		t.userCounters.decRLimitNProc()
 		t.tg.exitedCPUStats.Accumulate(t.CPUStats())

pkg/sentry/kernel/task_start.go

@@ -248,28 +248,25 @@ func (ts *TaskSet) assignTIDsLocked(t *Task) error {
 		tid ThreadID
 	}
 	var allocatedTIDs []allocatedTID
+	var tid ThreadID
+	var err error
 	for ns := t.tg.pidns; ns != nil; ns = ns.parent {
-		tid, err := ns.allocateTID()
-		if err != nil {
-			// Failure. Remove the tids we already allocated in descendant
-			// namespaces.
-			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
+		if tid, err = ns.allocateTID(); err != nil {
+			break
 		}
-		ns.tasks[tid] = t
-		ns.tids[t] = tid
-		if t.tg.leader == nil {
-			// New thread group.
-			ns.tgids[t.tg] = tid
+		if err = ns.addTask(t, tid); err != nil {
+			break
 		}
 		allocatedTIDs = append(allocatedTIDs, allocatedTID{ns, tid})
 	}
+	if err != nil {
+		// Failure. Remove the tids we already allocated in descendant
+		// namespaces.
+		for _, a := range allocatedTIDs {
+			a.ns.deleteTask(t)
+		}
+		return err
+	}
 	return nil
 }
 

pkg/sentry/kernel/threads.go

@@ -187,6 +187,9 @@ type PIDNamespace struct {
 	// exiting indicates that the namespace's init process is exiting or has
 	// exited.
 	exiting bool
+
+	// pidNamespaceData contains additional per-PID-namespace data.
+	extra pidNamespaceData
 }
 
 func newPIDNamespace(ts *TaskSet, parent *PIDNamespace, userns *auth.UserNamespace) *PIDNamespace {
@@ -201,6 +204,7 @@ func newPIDNamespace(ts *TaskSet, parent *PIDNamespace, userns *auth.UserNamespa
 		sids:          make(map[*Session]SessionID),
 		processGroups: make(map[ProcessGroupID]*ProcessGroup),
 		pgids:         make(map[*ProcessGroup]ProcessGroupID),
+		extra:         newPIDNamespaceData(),
 	}
 }
 

pkg/sentry/kernel/threads_impl.go

@@ -0,0 +1,50 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.1
+// +build go1.1
+
+package kernel
+
+// pidNamespaceData may contain extra per-PID-namespace data.
+// +stateify savable
+type pidNamespaceData struct {
+}
+
+// newPIDNamespaceData returns a new `pidNamespaceData` struct.
+func newPIDNamespaceData() pidNamespaceData {
+	return pidNamespaceData{}
+}
+
+// addTask adds a Task into this PIDNamespace.
+// It is always performed under TaskSet lock.
+func (ns *PIDNamespace) addTask(t *Task, tid ThreadID) error {
+	ns.tasks[tid] = t
+	ns.tids[t] = tid
+	if t.tg.leader == nil {
+		// New thread group.
+		ns.tgids[t.tg] = tid
+	}
+	return nil
+}
+
+// deleteTask deletes a Task from this PIDNamespace.
+// It is always performed under TaskSet lock.
+func (ns *PIDNamespace) deleteTask(t *Task) {
+	delete(ns.tasks, ns.tids[t])
+	delete(ns.tids, t)
+	if t == t.tg.leader || t.tg.leader == nil {
+		delete(ns.tgids, t.tg)
+	}
+}

runsc/container/container_test.go

@@ -1016,6 +1016,7 @@ func TestKillPid(t *testing.T) {
 			if err != nil {
 				t.Fatalf("failed to get process list: %v", err)
 			}
+			t.Logf("current process list: %v", procs)
 			var pid int32
 			for _, p := range procs {
 				if pid < int32(p.PID) {
@@ -1028,7 +1029,8 @@ func TestKillPid(t *testing.T) {
 
 			// Verify that one process is gone.
 			if err := waitForProcessCount(cont, nProcs-1); err != nil {
-				t.Fatalf("error waiting for processes: %v", err)
+				procs, procsErr := cont.Processes()
+				t.Fatalf("error waiting for processes: %v; current processes: %v / %v", err, procs, procsErr)
 			}
 
 			procs, err = cont.Processes()