systrap: check that minimum one stub thread is active after queueing a context

and don't activate more threads than contexts.

PiperOrigin-RevId: 519168272

pkg/sentry/platform/systrap/context_queue.go

@@ -45,8 +45,10 @@ type contextQueue struct {
 	// stubPollingIndexBase is used by stubs to indicate to each other how many
 	// threads went to sleep.
 	stubPollingIndexBase uint32
-	// numSleepingThreads indicates to the sentry how many stubs are asleep.
-	numSleepingThreads uint32
+	// numActiveThreads indicates to the sentry how many stubs are running.
+	numActiveThreads uint32
+	// numActiveContext is a number of running and waiting contexts
+	numActiveContexts uint32
 	// ringbuffer is the mmapped region of memory that's shared with the stub
 	// threads.
 	ringbuffer [maxContextQueueEntries]uint32
@@ -62,7 +64,8 @@ func (q *contextQueue) init() {
 	atomic.StoreUint32(&q.end, 0)
 	atomic.StoreUint32(&q.stubPollingIndex, 0)
 	atomic.StoreUint32(&q.stubPollingIndexBase, 0)
-	atomic.StoreUint32(&q.numSleepingThreads, 0)
+	atomic.StoreUint32(&q.numActiveThreads, 0)
+	atomic.StoreUint32(&q.numActiveContexts, 0)
 }
 
 func (q *contextQueue) isEmpty() bool {
@@ -74,6 +77,7 @@ func (q *contextQueue) queuedContexts() uint32 {
 }
 
 func (q *contextQueue) add(contextID uint32) uint32 {
+	atomic.AddUint32(&q.numActiveContexts, 1)
 	next := atomic.AddUint32(&q.end, 1)
 	if (next % maxContextQueueEntries) ==
 		(atomic.LoadUint32(&q.start) % maxContextQueueEntries) {

pkg/sentry/platform/systrap/subprocess.go

@@ -315,7 +315,9 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 
 	// Create the initial sysmsg thread.
 	if contextDecouplingExp {
+		atomic.AddUint32(&sp.contextQueue.numActiveThreads, 1)
 		if _, err := sp.createSysmsgThread(nil, nil, nil); err != nil {
+			atomic.AddUint32(&sp.contextQueue.numActiveThreads, ^uint32(0))
 			return nil, err
 		}
 		sp.numSysmsgThreads++
@@ -799,6 +801,10 @@ func (s *subprocess) waitOnState(ctx *sharedContext) {
 	slowPath := false
 	start := cputicks()
 	handshake := false
+	if atomic.LoadUint32(&s.contextQueue.numActiveThreads) == 0 {
+		kicked = true
+		s.kickSysmsgThread()
+	}
 	for curState := ctx.state(); curState == sysmsg.ContextStateNone; curState = ctx.state() {
 		if !slowPath {
 			delta := uint64(cputicks() - start)
@@ -832,7 +838,15 @@ func (s *subprocess) waitOnState(ctx *sharedContext) {
 func (s *subprocess) kickSysmsgThread() {
 	s.sysmsgThreadsMu.Lock()
 
-	if atomic.LoadUint32(&s.contextQueue.numSleepingThreads) > 0 {
+	nrActiveContexts := atomic.LoadUint32(&s.contextQueue.numActiveContexts)
+	nrActiveThreads := atomic.LoadUint32(&s.contextQueue.numActiveThreads)
+
+	if nrActiveThreads >= nrActiveContexts {
+		s.sysmsgThreadsMu.Unlock()
+		return
+	}
+
+	if s.numSysmsgThreads > int(nrActiveThreads) {
 		for _, t := range s.sysmsgThreads {
 			if t.msg.State.Get() == sysmsg.ThreadStateAsleep {
 				t.msg.WakeSysmsgThread()
@@ -847,7 +861,10 @@ func (s *subprocess) kickSysmsgThread() {
 	if s.numSysmsgThreads < maxSysmsgThreads {
 		s.numSysmsgThreads++
 		s.sysmsgThreadsMu.Unlock()
+		atomic.AddUint32(&s.contextQueue.numActiveThreads, 1)
 		if _, err := s.createSysmsgThread(nil, nil, nil); err != nil {
+			log.Warningf("Unable to create a new stub thread: %s", err)
+			atomic.AddUint32(&s.contextQueue.numActiveThreads, ^uint32(0))
 			s.sysmsgThreadsMu.Lock()
 			s.numSysmsgThreads--
 			s.sysmsgThreadsMu.Unlock()

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

@@ -41,7 +41,8 @@ struct context_queue {
   uint32_t end;
   uint32_t polling_index;
   uint32_t polling_index_base;
-  uint32_t num_sleeping_threads;
+  uint32_t num_active_threads;
+  uint32_t num_active_contexts;
   uint32_t ringbuffer[MAX_CONTEXT_QUEUE_ENTRIES];
 };
 
@@ -209,22 +210,43 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
     __atomic_store_n(&sysmsg->state, THREAD_STATE_PREP, __ATOMIC_RELEASE);
     ctx = queue_get_context(sysmsg);
     if (ctx) return ctx;
+
     if (spinning_queue_push()) {
-      while (!spinning_queue_remove_first(__export_deep_sleep_timeout)) {
+      while (1) {
         ctx = queue_get_context(sysmsg);
         if (ctx) {
           spinning_queue_pop();
           return ctx;
         }
-
+        if (spinning_queue_remove_first(__export_deep_sleep_timeout)) {
+          break;
+        }
         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 =
+        __atomic_load_n(&queue->num_active_contexts, __ATOMIC_ACQUIRE);
+    // We have to make another attempt to get a context here to prevent TOCTTOU
+    // races with waitOnState and kickSysmsgThread. There are two assumptions:
+    // * If the queue isn't empty, one or more threads have to be active.
+    // * A new thread isn't kicked, if the number of active threads are not less
+    //   than a number of active contexts.
+    if (nr_active_threads == 0 || nr_active_threads < nr_active_contexts) {
+      ctx = queue_get_context(sysmsg);
+      if (ctx) {
+        __atomic_add_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
+        return ctx;
+      }
+    }
 
-    __atomic_add_fetch(&queue->num_sleeping_threads, 1, __ATOMIC_ACQ_REL);
-    sys_futex(&sysmsg->state, FUTEX_WAIT, THREAD_STATE_ASLEEP, NULL, NULL, 0);
-    __atomic_sub_fetch(&queue->num_sleeping_threads, 1, __ATOMIC_ACQ_REL);
+    __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);
+    }
+    __atomic_add_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
   }
 }
 
@@ -233,6 +255,8 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
 struct thread_context *switch_context(struct sysmsg *sysmsg,
                                       struct thread_context *ctx,
                                       enum context_state new_context_state) {
+  struct context_queue *queue = __export_context_queue_addr;
+
   __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);
@@ -242,6 +266,7 @@ 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;