systrap: kick another stub thread after the handshake timeout

The idea of the handshake timeout is to see if here is any free active thread.

In addition, here are a few fixes that prevent waking up too many threads.

PiperOrigin-RevId: 519871285

pkg/sentry/platform/systrap/subprocess.go

@@ -787,44 +787,38 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 	return false, false, nil
 }
 
-const (
-	// deepSleepTimeout is the timeout after which we stop polling and fall asleep.
-	// The value is 100µs for 2GHz CPU.
-	decoupledDeepSleepTimeout = uint64(200000)
-	// threadKickTimeout is the timeout after which we will either wake up a sleeping
-	// thread or create a new one.
-	threadKickTimeout = uint64(20000)
-)
-
 func (s *subprocess) waitOnState(ctx *sharedContext) {
 	kicked := false
 	slowPath := false
 	start := cputicks()
 	handshake := false
 	if atomic.LoadUint32(&s.contextQueue.numActiveThreads) == 0 {
-		kicked = true
-		s.kickSysmsgThread()
+		kicked = s.kickSysmsgThread()
 	}
 	for curState := ctx.state(); curState == sysmsg.ContextStateNone; curState = ctx.state() {
 		if !slowPath {
 			delta := uint64(cputicks() - start)
-			if delta > decoupledDeepSleepTimeout {
+			if delta > deepSleepTimeout {
 				ctx.disableSentryFastPath()
 				slowPath = true
 				continue
 			}
 
-			if !handshake && ctx.isAcked() {
-				handshake = true
-				continue
+			if !handshake {
+				if ctx.isAcked() {
+					handshake = true
+					continue
+				}
+				if !kicked && delta > handshakeTimeout {
+					kicked = s.kickSysmsgThread()
+				}
 			}
 			spinloop()
 		} else {
 			// If the context already received a handshake then it knows it's being
 			// worked on.
 			if !kicked && !handshake {
-				kicked = true
-				s.kickSysmsgThread()
+				kicked = s.kickSysmsgThread()
 			}
 
 			ctx.sleepOnState(curState)
@@ -835,29 +829,46 @@ func (s *subprocess) waitOnState(ctx *sharedContext) {
 	ctx.enableSentryFastPath()
 }
 
-func (s *subprocess) kickSysmsgThread() {
-	s.sysmsgThreadsMu.Lock()
-
+func (s *subprocess) kickSysmsgThread() bool {
+	// numActiveContexts and numActiveThreads can be changed from stub
+	// threads that work with the contextQueue without any locks. The idea
+	// here is that any stub thread that gets CPU time can make some
+	// progress. In stub threads, we can use only spinlock-like
+	// synchronizations, but they don't work well because a thread that
+	// holds a lock can be preempted by another threads that is waiting for
+	// the same lock.
 	nrActiveContexts := atomic.LoadUint32(&s.contextQueue.numActiveContexts)
 	nrActiveThreads := atomic.LoadUint32(&s.contextQueue.numActiveThreads)
+	if nrActiveContexts != 0 && nrActiveThreads >= nrActiveContexts {
+		// This can happen when one or more stub threads are
+		// waiting for cpu time. The host probably has more
+		// running tasks than a number of cpu-s.
+		return false
+	}
 
-	if nrActiveThreads >= nrActiveContexts {
+	s.sysmsgThreadsMu.Lock()
+	nrActiveContexts = atomic.LoadUint32(&s.contextQueue.numActiveContexts)
+	nrActiveThreads = atomic.LoadUint32(&s.contextQueue.numActiveThreads)
+	if nrActiveContexts != 0 && nrActiveThreads >= nrActiveContexts {
 		s.sysmsgThreadsMu.Unlock()
-		return
+		return false
 	}
 
 	if s.numSysmsgThreads > int(nrActiveThreads) {
 		for _, t := range s.sysmsgThreads {
-			if t.msg.State.Get() == sysmsg.ThreadStateAsleep {
-				t.msg.WakeSysmsgThread()
+			if kicked, _ := t.msg.WakeSysmsgThread(); kicked {
 				s.sysmsgThreadsMu.Unlock()
-				return
+				return true
 			}
 		}
+		s.sysmsgThreadsMu.Unlock()
+		// Threads are kicked only here under sysmsgThreadsMu. It means
+		// that this case is possible only if one thread decides to
+		// fall asleep but then change its mind. Look at
+		// sysmsg_lib.c:get_context for more details.
+		return false
 	}
-	// It's also possible that we got here after iterating through all other
-	// threads and not finding anything asleep because other goroutines already
-	// woke up every other thread up.
+
 	if s.numSysmsgThreads < maxSysmsgThreads {
 		s.numSysmsgThreads++
 		s.sysmsgThreadsMu.Unlock()
@@ -868,10 +879,13 @@ func (s *subprocess) kickSysmsgThread() {
 			s.sysmsgThreadsMu.Lock()
 			s.numSysmsgThreads--
 			s.sysmsgThreadsMu.Unlock()
+			return false
 		}
-	} else {
-		s.sysmsgThreadsMu.Unlock()
+		return true
 	}
+
+	s.sysmsgThreadsMu.Unlock()
+	return false
 }
 
 // syscall executes the given system call without handling interruptions.

pkg/sentry/platform/systrap/sysmsg/sysmsg.go

@@ -93,6 +93,11 @@ func (s *ThreadState) Set(state ThreadState) {
 	atomic.StoreUint32((*uint32)(s), uint32(state))
 }
 
+// CompareAndSwap atomicaly compares and swaps the state value.
+func (s *ThreadState) CompareAndSwap(old, state ThreadState) bool {
+	return atomic.CompareAndSwapUint32((*uint32)(s), uint32(old), uint32(state))
+}
+
 // Get returns the current state value.
 //
 //go:nosplit

pkg/sentry/platform/systrap/sysmsg/sysmsg_lib.c

@@ -224,6 +224,8 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
         spinloop();
       }
     }
+
+    __atomic_store_n(&sysmsg->state, THREAD_STATE_ASLEEP, __ATOMIC_RELEASE);
     uint32_t nr_active_threads =
         __atomic_sub_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
     uint32_t nr_active_contexts =
@@ -236,12 +238,12 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
     if (nr_active_threads == 0 || nr_active_threads < nr_active_contexts) {
       ctx = queue_get_context(sysmsg);
       if (ctx) {
+        __atomic_store_n(&sysmsg->state, THREAD_STATE_PREP, __ATOMIC_RELEASE);
         __atomic_add_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
         return ctx;
       }
     }
 
-    __atomic_store_n(&sysmsg->state, THREAD_STATE_ASLEEP, __ATOMIC_RELEASE);
     while (__atomic_load_n(&sysmsg->state, __ATOMIC_ACQUIRE) ==
            THREAD_STATE_ASLEEP) {
       sys_futex(&sysmsg->state, FUTEX_WAIT, THREAD_STATE_ASLEEP, NULL, NULL, 0);
@@ -257,6 +259,7 @@ struct thread_context *switch_context(struct sysmsg *sysmsg,
                                       enum context_state new_context_state) {
   struct context_queue *queue = __export_context_queue_addr;
 
+  __atomic_sub_fetch(&queue->num_active_contexts, 1, __ATOMIC_ACQ_REL);
   __atomic_store_n(&ctx->thread_id, INVALID_THREAD_ID, __ATOMIC_RELEASE);
   __atomic_store_n(&ctx->last_thread_id, sysmsg->thread_id, __ATOMIC_RELEASE);
   __atomic_store_n(&ctx->state, new_context_state, __ATOMIC_RELEASE);
@@ -266,7 +269,6 @@ struct thread_context *switch_context(struct sysmsg *sysmsg,
       panic(ret);
     }
   }
-  __atomic_sub_fetch(&queue->num_active_contexts, 1, __ATOMIC_ACQ_REL);
 
   struct thread_context *old_ctx = sysmsg->context;
 

pkg/sentry/platform/systrap/sysmsg/sysmsg_unsafe.go

@@ -81,8 +81,10 @@ func (c *ThreadContext) SleepOnState(curState ContextState, interruptor interrup
 }
 
 // WakeSysmsgThread calls futex wake on Sysmsg.State.
-func (m *Msg) WakeSysmsgThread() syscall.Errno {
-	m.State.Set(ThreadStatePrep)
+func (m *Msg) WakeSysmsgThread() (bool, syscall.Errno) {
+	if !m.State.CompareAndSwap(ThreadStateAsleep, ThreadStatePrep) {
+		return false, 0
+	}
 	_, _, e := unix.RawSyscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&m.State)), linux.FUTEX_WAKE, 1, 0, 0, 0)
-	return e
+	return true, e
 }