From 7480450936dc2da5485d2599deea2ef83ae9d391 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 28 Feb 2024 17:03:28 -0800
Subject: [PATCH] Replace `Task.ptraceTracer` with `atomic.Pointer`.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This removes the compare-and-branch involved in casting
`ptraceTracer.Load()` to `*Task`.

This is on the hot syscall path, and in the overwhelmingly-likely case that
the task is not being traced, we should be as fast as possible.

Shaves off about 20 nanoseconds from the syscall hot path:

```
                                 │   before    │               after                │
                                 │   sec/op    │   sec/op     vs base               │
Syscallbench/syscall.getpidopt-8   1.355µ ± 2%   1.334µ ± 2%  -1.59% (p=0.030 n=64)
```

PiperOrigin-RevId: 611277514
---
 pkg/sentry/kernel/ptrace.go     | 4 ++--
 pkg/sentry/kernel/task.go       | 8 +++-----
 pkg/sentry/kernel/task_exit.go  | 2 +-
 pkg/sentry/kernel/task_start.go | 1 -
 4 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/pkg/sentry/kernel/ptrace.go b/pkg/sentry/kernel/ptrace.go
index 17ac1b56a..a8c591146 100644
--- a/pkg/sentry/kernel/ptrace.go
+++ b/pkg/sentry/kernel/ptrace.go
@@ -317,7 +317,7 @@ func (t *Task) SetYAMAException(tracer *Task) {
 
 // Tracer returns t's ptrace Tracer.
 func (t *Task) Tracer() *Task {
-	return t.ptraceTracer.Load().(*Task)
+	return t.ptraceTracer.Load()
 }
 
 // hasTracer returns true if t has a ptrace tracer attached.
@@ -605,7 +605,7 @@ func (t *Task) forgetTracerLocked() {
 	t.ptraceOpts = ptraceOptions{}
 	t.ptraceSyscallMode = ptraceSyscallNone
 	t.ptraceSinglestep = false
-	t.ptraceTracer.Store((*Task)(nil))
+	t.ptraceTracer.Store(nil)
 	if t.exitTracerNotified && !t.exitTracerAcked {
 		t.exitTracerAcked = true
 		t.exitNotifyLocked(true)
diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go
index c0af03092..4b0a131a4 100644
--- a/pkg/sentry/kernel/task.go
+++ b/pkg/sentry/kernel/task.go
@@ -319,15 +319,13 @@ type Task struct {
 	goroutineStopped sync.WaitGroup `state:"nosave"`
 
 	// ptraceTracer is the task that is ptrace-attached to this one. If
-	// ptraceTracer is nil, this task is not being traced. Note that due to
-	// atomic.Value limitations (atomic.Value.Store(nil) panics), a nil
-	// ptraceTracer is always represented as a typed nil (i.e. (*Task)(nil)).
+	// ptraceTracer is nil, this task is not being traced.
 	//
 	// ptraceTracer is protected by the TaskSet mutex, and accessed with atomic
 	// operations. This allows paths that wouldn't otherwise lock the TaskSet
 	// mutex, notably the syscall path, to check if ptraceTracer is nil without
 	// additional synchronization.
-	ptraceTracer atomic.Value `state:".(*Task)"`
+	ptraceTracer atomic.Pointer[Task] `state:".(*Task)"`
 
 	// ptraceTracees is the set of tasks that this task is ptrace-attached to.
 	//
@@ -616,7 +614,7 @@ var (
 )
 
 func (t *Task) savePtraceTracer() *Task {
-	return t.ptraceTracer.Load().(*Task)
+	return t.ptraceTracer.Load()
 }
 
 func (t *Task) loadPtraceTracer(tracer *Task) {
diff --git a/pkg/sentry/kernel/task_exit.go b/pkg/sentry/kernel/task_exit.go
index 0bca6a294..c4d9c8c6e 100644
--- a/pkg/sentry/kernel/task_exit.go
+++ b/pkg/sentry/kernel/task_exit.go
@@ -1111,7 +1111,7 @@ func (t *Task) waitCollectZombieLocked(target *Task, opts *WaitOptions, asPtrace
 	// will be reaped here.
 	if tracer := target.Tracer(); tracer != nil && tracer.tg == t.tg && target.exitTracerNotified {
 		target.exitTracerAcked = true
-		target.ptraceTracer.Store((*Task)(nil))
+		target.ptraceTracer.Store(nil)
 		delete(t.ptraceTracees, target)
 	}
 	if target.parent != nil && target.parent.tg == t.tg && target.exitParentNotified {
diff --git a/pkg/sentry/kernel/task_start.go b/pkg/sentry/kernel/task_start.go
index bb6a0606c..c5cb5411a 100644
--- a/pkg/sentry/kernel/task_start.go
+++ b/pkg/sentry/kernel/task_start.go
@@ -176,7 +176,6 @@ func (ts *TaskSet) newTask(ctx context.Context, cfg *TaskConfig) (*Task, error)
 	t.netns = cfg.NetworkNamespace
 	t.creds.Store(cfg.Credentials)
 	t.endStopCond.L = &t.tg.signalHandlers.mu
-	t.ptraceTracer.Store((*Task)(nil))
 	t.seccomp.Store((*taskSeccomp)(nil))
 	// We don't construct t.blockingTimer until Task.run(); see that function
 	// for justification.