Pack Task.gostate{,Seq,Time} into a single 64-bit atomic int.

pkg/atomicbitops/aligned_32bit_unsafe.go

@@ -198,6 +198,16 @@ func (u *Uint64) RacyStore(v uint64) {
 	*u.ptr() = v
 }
 
+// StoreRelaxed is actually the same as `Store` on 32-bit architectures,
+// since 64-bit plain stores are not atomic there.
+// See comment on the analogous of this function in `aligned_64bit.go`.
+//
+//go:norace
+//go:nosplit
+func (u *Uint64) StoreRelaxed(v uint64) {
+	atomic.StoreUint64(u.ptr(), v)
+}
+
 // Add is analogous to atomic.AddUint64.
 //
 //go:nosplit

pkg/atomicbitops/aligned_64bit.go

@@ -174,6 +174,20 @@ func (u *Uint64) RacyStore(v uint64) {
 	u.value = v
 }
 
+// StoreRelaxed stores `v` with **no** ordering guarantee.
+// Useful only where readers tolerate a stale value.
+// Race detection is disabled, so this must be used sparingly.
+// On 64-bit architectures, readers are guaranteed to see either the old
+// or the new value, no "partial writes" cases.
+// The 32-bit-architecture variant of this function does a real atomic
+// write to guarantee the same no-partial-write property.
+//
+//go:norace
+//go:nosplit
+func (u *Uint64) StoreRelaxed(v uint64) {
+	u.value = v
+}
+
 // Add is analogous to atomic.AddUint64.
 //
 //go:nosplit

pkg/sentry/kernel/task.go

@@ -108,22 +108,20 @@ type Task struct {
 	// interruptChan is always notified after restore (see Task.run).
 	interruptChan chan struct{} `state:"nosave"`
 
-	// gostateSeq allows Task.TaskGoroutineStateTime() to read gostate and
-	// gostateTime atomically.
-	//
-	// gostateSeq is owned by the task goroutine.
-	gostateSeq sync.SeqCount `state:"nosave"`
-
-	// gostate is the current scheduling state of the task goroutine.
+	// gostate combines two fields in one 64-bit integer:
+	//  - First `gostateBits` bits are the `TaskGoroutineState` enum value.
+	//  - Rest of the bits are the value of Kernel.cpuClock when the state was
+	//   last updated or refreshed.
+	// Packing them this way allows efficient atomic reads and writes, which
+	// is critical for performance on the syscall hot path.
+	//
+	// Despite the use of `atomicbitops.Uint64`, `gostate` is written to
+	// **with no barrier guarantee** from the task goroutine, for syscall hot
+	// path performance reasons. This means all readers (other than from the
+	// task goroutine) **must** tolerate stale reads.
 	//
 	// gostate is owned by the task goroutine.
-	gostate atomicbitops.Uint32
-
-	// gostateTime was the value of Kernel.cpuClock when gostate was last
-	// updated or refreshed.
-	//
-	// gostateTime is owned by the task goroutine.
-	gostateTime atomicbitops.Int64
+	gostate atomicbitops.Uint64
 
 	// appCPUClock approximates the amount of time the task goroutine has spent
 	// in TaskGoroutineRunningApp.
@@ -754,7 +752,7 @@ func (t *Task) afterLoad(gocontext.Context) {
 		ts.populateCache(t)
 	}
 	t.interruptChan = make(chan struct{}, 1)
-	t.gostate.Store(uint32(TaskGoroutineNonexistent))
+	t.setGostate(TaskGoroutineNonexistent)
 	if t.stop != nil {
 		t.stopCount = atomicbitops.FromInt32(1)
 	}

pkg/sentry/kernel/task_sched.go

@@ -64,34 +64,49 @@ const (
 	TaskGoroutineStopped
 )
 
+// gostateBits is the width of the `TaskGoroutineState` bits in `Task`'s
+// `gostate` field. The remainder of the bits holds `Kernel.cpuClock`
+// nanoseconds. With gostateBits = 3, this is enough for 73 years.
+const gostateBits = 3
+
 // TaskGoroutineState returns the current state of the task goroutine.
 func (t *Task) TaskGoroutineState() TaskGoroutineState {
-	return TaskGoroutineState(t.gostate.Load())
+	return TaskGoroutineState(t.gostate.Load() & (1<<gostateBits - 1))
+}
+
+// ownTaskGoroutineState returns the current state of the task goroutine,
+// without using an atomic load. This is appropriate to do when running on
+// the task goroutine, since the task goroutine is the only writer.
+//
+// Preconditions: The caller must be running on the task goroutine.
+func (t *Task) ownTaskGoroutineState() TaskGoroutineState {
+	return TaskGoroutineState(t.gostate.RacyLoad() & (1<<gostateBits - 1))
 }
 
 // TaskGoroutineStateTime returns the current state of the task goroutine, and
 // the value of Kernel.CPUClockNow() when that state was last updated or
 // refreshed.
 func (t *Task) TaskGoroutineStateTime() (state TaskGoroutineState, time ktime.Time) {
-	for {
-		epoch := t.gostateSeq.BeginRead()
-		state = t.TaskGoroutineState()
-		time = ktime.FromNanoseconds(t.gostateTime.Load())
-		if t.gostateSeq.ReadOk(epoch) {
-			return
-		}
-	}
+	v := t.gostate.Load()
+	return TaskGoroutineState(v & (1<<gostateBits - 1)), ktime.FromNanoseconds(int64(v >> gostateBits))
+}
+
+// setGostate sets the task goroutine state, timestamped with the current
+// `Kernel.cpuClock`.
+//
+// Preconditions: The caller must be running on the task goroutine.
+func (t *Task) setGostate(state TaskGoroutineState) {
+	// StoreRelaxed due to this being on the syscall hot path, and all readers
+	// are expected to tolerate stale reads.
+	t.gostate.StoreRelaxed(uint64(t.k.cpuClock.Load())<<gostateBits | uint64(state))
 }
 
 // Preconditions: The caller must be running on the task goroutine.
 func (t *Task) accountTaskGoroutineEnter(state TaskGoroutineState) {
-	if oldState := t.TaskGoroutineState(); oldState != TaskGoroutineRunningSys {
+	if oldState := t.ownTaskGoroutineState(); oldState != TaskGoroutineRunningSys {
 		panic(fmt.Sprintf("Task goroutine switching from state %v (expected %v) to %v", oldState, TaskGoroutineRunningSys, state))
 	}
-	t.gostateSeq.BeginWrite()
-	t.gostate.Store(uint32(state))
-	t.touchGostateTime()
-	t.gostateSeq.EndWrite()
+	t.setGostate(state)
 	if state != TaskGoroutineRunningApp {
 		// Task is blocking/stopping.
 		t.k.decRunningTasks()
@@ -115,18 +130,17 @@ func (t *Task) accountTaskGoroutineLeave(state TaskGoroutineState) {
 		// Task is leaving uninterruptible sleep.
 		t.k.blockedTasks.Add(-1)
 	}
-	if oldState := t.TaskGoroutineState(); oldState != state {
+	if oldState := t.ownTaskGoroutineState(); oldState != state {
 		panic(fmt.Sprintf("Task goroutine switching from state %v (expected %v) to %v", oldState, state, TaskGoroutineRunningSys))
 	}
-	t.gostateSeq.BeginWrite()
-	t.gostate.Store(uint32(TaskGoroutineRunningSys))
-	t.touchGostateTime()
-	t.gostateSeq.EndWrite()
+	t.setGostate(TaskGoroutineRunningSys)
 }
 
+// touchGostateTime refreshes the timestamp of the current task goroutine state.
+//
 // Preconditions: The caller must be running on the task goroutine.
 func (t *Task) touchGostateTime() {
-	t.gostateTime.Store(t.k.cpuClock.Load())
+	t.setGostate(t.ownTaskGoroutineState())
 }
 
 // CPUClockNow returns the current value of the kernel CPU clock, which