systrap: limit a number of stub thread by a number of awake contexts

A context is awake if its guest thread isn't in the interruption sleep state.

The idea here is that a task in the sleep state will not return within the fast
path timeout and so we don't need to hold a stub thread for it.

PiperOrigin-RevId: 529006216

pkg/sentry/platform/systrap/context_queue.go

@@ -43,7 +43,11 @@ type contextQueue struct {
 	// numActiveThreads indicates to the sentry how many stubs are running.
 	numActiveThreads uint32
 	// numActiveContext is a number of running and waiting contexts
-	numActiveContexts   uint32
+	numActiveContexts uint32
+	// numAwakeContexts is the number of awake contexts. It includes all
+	// active contexts and contexts that are running in the Sentry.
+	numAwakeContexts uint32
+
 	fastPathDisabledTS  uint64
 	fastPathFailedInRow uint32
 	ringbuffer          [maxContextQueueEntries]uint32
@@ -61,6 +65,7 @@ func (q *contextQueue) init() {
 	atomic.StoreUint32(&q.fastPathFailedInRow, 0)
 	atomic.StoreUint32(&q.numActiveThreads, 0)
 	atomic.StoreUint32(&q.numActiveContexts, 0)
+	atomic.StoreUint32(&q.numAwakeContexts, 0)
 }
 
 func (q *contextQueue) isEmpty() bool {
@@ -71,7 +76,12 @@ func (q *contextQueue) queuedContexts() uint32 {
 	return (atomic.LoadUint32(&q.end) + maxContextQueueEntries - atomic.LoadUint32(&q.start)) % maxContextQueueEntries
 }
 
-func (q *contextQueue) add(contextID uint32) uint32 {
+func (q *contextQueue) add(ctx *sharedContext) uint32 {
+	contextID := ctx.contextID
+	if ctx.sleeping {
+		ctx.sleeping = false
+		atomic.AddUint32(&q.numAwakeContexts, 1)
+	}
 	atomic.AddUint32(&q.numActiveContexts, 1)
 	next := atomic.AddUint32(&q.end, 1)
 	if (next % maxContextQueueEntries) ==

pkg/sentry/platform/systrap/shared_context.go

@@ -57,6 +57,8 @@ type sharedContext struct {
 	sync           syncevent.Waiter
 	startWaitingTS int64
 	kicked         bool
+	// The task associated with the context fell asleep.
+	sleeping bool
 }
 
 const (
@@ -87,6 +89,7 @@ func (s *subprocess) getSharedContext() (*sharedContext, error) {
 	}
 	sc.shared.Init(invalidThreadID)
 	sc.sync.Init()
+	sc.sleeping = true
 
 	return &sc, nil
 }
@@ -95,6 +98,10 @@ func (sc *sharedContext) release() {
 	if sc == nil {
 		return
 	}
+	if !sc.sleeping {
+		atomic.AddUint32(&sc.subprocess.contextQueue.numAwakeContexts, ^uint32(0))
+
+	}
 	sc.subprocess.threadContextPool.Put(uint64(sc.contextID))
 	sc.subprocess.DecRef(sc.subprocess.release)
 }

pkg/sentry/platform/systrap/subprocess.go

@@ -726,7 +726,7 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 
 		// Place the context onto the context queue.
 		ctx.setState(sysmsg.ContextStateNone)
-		s.contextQueue.add(uint32(ctx.contextID))
+		s.contextQueue.add(ctx)
 		s.waitOnState(ctx)
 
 		// Check if there's been an error.

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

@@ -41,6 +41,7 @@ struct context_queue {
   uint32_t end;
   uint32_t num_active_threads;
   uint32_t num_active_contexts;
+  uint32_t num_awake_contexts;
   uint64_t fast_path_disalbed_ts;
   uint32_t fast_path_failed_in_row;
   uint32_t ringbuffer[MAX_CONTEXT_QUEUE_ENTRIES];
@@ -201,10 +202,67 @@ struct thread_context *queue_get_context(struct sysmsg *sysmsg) {
 #define FAILED_FAST_PATH_LIMIT 5
 #define FAILED_FAST_PATH_TIMEOUT 20000000  // 10ms
 
+// get_context_fast sets nr_active_threads_p only if it deactivates the thread.
+static struct thread_context *get_context_fast(struct sysmsg *sysmsg,
+                                               struct context_queue *queue,
+                                               uint32_t *nr_active_threads_p) {
+  uint32_t nr_active_threads, nr_awake_contexts;
+
+  if (!spinning_queue_push()) return NULL;
+
+  while (1) {
+    struct thread_context *ctx;
+
+    ctx = queue_get_context(sysmsg);
+    if (ctx) {
+      __atomic_store_n(&queue->fast_path_failed_in_row, 0, __ATOMIC_RELEASE);
+      spinning_queue_pop();
+      return ctx;
+    }
+
+    nr_active_threads =
+        __atomic_load_n(&queue->num_active_threads, __ATOMIC_ACQUIRE);
+    nr_awake_contexts =
+        __atomic_load_n(&queue->num_awake_contexts, __ATOMIC_ACQUIRE);
+
+    if (nr_awake_contexts < nr_active_threads) {
+      if (__atomic_compare_exchange_n(&queue->num_active_threads,
+                                      &nr_active_threads, nr_active_threads - 1,
+                                      false, __ATOMIC_SEQ_CST,
+                                      __ATOMIC_SEQ_CST)) {
+        nr_active_threads -= 1;
+        if (spinning_queue_remove_first(0)) {
+          *nr_active_threads_p = nr_active_threads;
+          break;
+        }
+
+        // spinning_queue_remove_first can fail due to a race with another
+        // thread.
+        __atomic_add_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
+      }
+    }
+
+    if (spinning_queue_remove_first(__export_deep_sleep_timeout)) {
+      uint32_t nr = __atomic_add_fetch(&queue->fast_path_failed_in_row, 1,
+                                       __ATOMIC_ACQ_REL);
+      if (nr >= FAILED_FAST_PATH_LIMIT) {
+        __atomic_store_n(&queue->fast_path_disalbed_ts, rdtsc(),
+                         __ATOMIC_RELEASE);
+      }
+      break;
+    }
+    spinloop();
+  }
+  return NULL;
+}
+
+#define NR_IF_THREAD_IS_ACTIVE (~0)
+
 // get_context retrieves a context that is ready to be restored to the user.
 // This populates sysmsg->thread_context_id.
 struct thread_context *get_context(struct sysmsg *sysmsg) {
   struct context_queue *queue = __export_context_queue_addr;
+  uint32_t nr_active_threads;
 
   for (;;) {
     struct thread_context *ctx;
@@ -230,31 +288,17 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
       }
     }
 
-    if (fast_path_enabled && spinning_queue_push()) {
-      while (1) {
-        ctx = queue_get_context(sysmsg);
-        if (ctx) {
-          __atomic_store_n(&queue->fast_path_failed_in_row, 0,
-                           __ATOMIC_RELEASE);
-          spinning_queue_pop();
-          return ctx;
-        }
-        if (spinning_queue_remove_first(__export_deep_sleep_timeout)) {
-          uint32_t nr = __atomic_add_fetch(&queue->fast_path_failed_in_row, 1,
-                                           __ATOMIC_ACQ_REL);
-          if (nr >= FAILED_FAST_PATH_LIMIT) {
-            __atomic_store_n(&queue->fast_path_disalbed_ts, rdtsc(),
-                             __ATOMIC_RELEASE);
-          }
-          break;
-        }
-        spinloop();
-      }
+    nr_active_threads = NR_IF_THREAD_IS_ACTIVE;
+    if (fast_path_enabled) {
+      ctx = get_context_fast(sysmsg, queue, &nr_active_threads);
+      if (ctx) return ctx;
+    }
+    if (nr_active_threads == NR_IF_THREAD_IS_ACTIVE) {
+      nr_active_threads =
+          __atomic_sub_fetch(&queue->num_active_threads, 1, __ATOMIC_ACQ_REL);
     }
 
     __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

pkg/sentry/platform/systrap/systrap.go

@@ -52,6 +52,7 @@ import (
 	"fmt"
 	"os"
 	"sync"
+	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	pkgcontext "gvisor.dev/gvisor/pkg/context"
@@ -290,9 +291,17 @@ func (c *context) Release() {
 // PrepareSleep implements platform.Context.platform.PrepareSleep.
 func (c *context) PrepareSleep() {
 	if contextDecouplingExp {
-		return // When this is called context hasn't entered the context queue.
-	}
-	if c.sysmsgThread != nil {
+		ctx := c.sharedContext
+		if ctx == nil {
+			return
+		}
+		s := ctx.subprocess
+		if !ctx.sleeping {
+			ctx.sleeping = true
+			atomic.AddUint32(&s.contextQueue.numAwakeContexts, ^uint32(0))
+		}
+		return
+	} else if c.sysmsgThread != nil {
 		c.sysmsgThread.msg.DisableStubFastPath()
 	}
 }