systrap: don't wake up each thread separately

Now all threads are waiting on queue->num_thread_to_wakeup,
it is a single point for all threads.

This change allows us to avoid cases when num_active_threads
are inconsistent with threads states, because they can't be
changed atomically.

PiperOrigin-RevId: 531020012

pkg/sentry/platform/systrap/BUILD

@@ -46,6 +46,7 @@ go_library(
     srcs = [
         "context_list.go",
         "context_queue.go",
+        "context_queue_unsafe.go",
         "filters.go",
         "filters_amd64.go",
         "filters_arm64.go",

pkg/sentry/platform/systrap/context_queue.go

@@ -40,8 +40,13 @@ type contextQueue struct {
 	start uint32
 	// end is an index used for putting new contexts into the ringbuffer.
 	end uint32
+
 	// numActiveThreads indicates to the sentry how many stubs are running.
+	// It is changed only by stub threads.
 	numActiveThreads uint32
+	// numThreadsToWakeup is the number of threads requested by Sentry to wake up.
+	// The Sentry increments it and stub threads decrements.
+	numThreadsToWakeup uint32
 	// numActiveContext is a number of running and waiting contexts
 	numActiveContexts uint32
 	// numAwakeContexts is the number of awake contexts. It includes all
@@ -73,6 +78,7 @@ func (q *contextQueue) init() {
 	atomic.StoreUint64(&q.fastPathDisabledTS, 0)
 	atomic.StoreUint32(&q.fastPathFailedInRow, 0)
 	atomic.StoreUint32(&q.numActiveThreads, 0)
+	atomic.StoreUint32(&q.numThreadsToWakeup, 0)
 	atomic.StoreUint32(&q.numActiveContexts, 0)
 	atomic.StoreUint32(&q.numAwakeContexts, 0)
 	atomic.StoreUint32(&q.fastPathDisabled, 0)

pkg/sentry/platform/systrap/context_queue_unsafe.go

@@ -0,0 +1,28 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package systrap
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
+
+func (q *contextQueue) wakeupSysmsgThread() {
+	unix.RawSyscall6(unix.SYS_FUTEX,
+		uintptr(unsafe.Pointer(&q.numThreadsToWakeup)),
+		linux.FUTEX_WAKE, 1, 0, 0, 0)
+}

pkg/sentry/platform/systrap/subprocess.go

@@ -315,9 +315,8 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 	sp.mapPrivateRegions()
 
 	// Create the initial sysmsg thread.
-	atomic.AddUint32(&sp.contextQueue.numActiveThreads, 1)
+	atomic.AddUint32(&sp.contextQueue.numThreadsToWakeup, 1)
 	if err := sp.createSysmsgThread(); err != nil {
-		atomic.AddUint32(&sp.contextQueue.numActiveThreads, ^uint32(0))
 		return nil, err
 	}
 	sp.numSysmsgThreads++
@@ -829,62 +828,58 @@ func (s *subprocess) waitOnState(ctx *sharedContext, stubFastPathEnabled bool) {
 	ctx.enableSentryFastPath()
 }
 
-func (s *subprocess) kickSysmsgThread() bool {
+// canKickSysmsgThread returns true if a new thread can be kicked.
+// The second return value is the expected number of threads after kicking a
+// new one.
+func (s *subprocess) canKickSysmsgThread() (bool, uint32) {
 	// numActiveContexts and numActiveThreads can be changed from stub
-	// threads that work with the contextQueue without any locks. The idea
+	// threads that handles 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
+	// holds a lock can be preempted by another thread that is waiting for
 	// the same lock.
-	nrActiveContexts := atomic.LoadUint32(&s.contextQueue.numActiveContexts)
 	nrActiveThreads := atomic.LoadUint32(&s.contextQueue.numActiveThreads)
-	if nrActiveContexts != 0 && nrActiveThreads >= nrActiveContexts {
+	nrThreadsToWakeup := atomic.LoadUint32(&s.contextQueue.numThreadsToWakeup)
+	nrActiveContexts := atomic.LoadUint32(&s.contextQueue.numActiveContexts)
+
+	nrActiveThreads += nrThreadsToWakeup + 1
+	if 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, nrActiveThreads
+	}
+	return true, nrActiveThreads
+}
+
+func (s *subprocess) kickSysmsgThread() bool {
+	kick, _ := s.canKickSysmsgThread()
+	if !kick {
 		return false
 	}
 
 	s.sysmsgThreadsMu.Lock()
-	nrActiveContexts = atomic.LoadUint32(&s.contextQueue.numActiveContexts)
-	nrActiveThreads = atomic.LoadUint32(&s.contextQueue.numActiveThreads)
-	if nrActiveContexts != 0 && nrActiveThreads >= nrActiveContexts {
+	kick, nrThreads := s.canKickSysmsgThread()
+	if !kick {
 		s.sysmsgThreadsMu.Unlock()
 		return false
 	}
-
-	if s.numSysmsgThreads > int(nrActiveThreads) {
-		for _, t := range s.sysmsgThreads {
-			if kicked, _ := t.msg.WakeSysmsgThread(); kicked {
-				s.sysmsgThreadsMu.Unlock()
-				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
-	}
-
-	if s.numSysmsgThreads < maxSysmsgThreads {
+	atomic.AddUint32(&s.contextQueue.numThreadsToWakeup, 1)
+	if s.numSysmsgThreads < maxSysmsgThreads && s.numSysmsgThreads < int(nrThreads) {
 		s.numSysmsgThreads++
 		s.sysmsgThreadsMu.Unlock()
-		atomic.AddUint32(&s.contextQueue.numActiveThreads, 1)
 		if err := s.createSysmsgThread(); 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()
-			return false
 		}
-		return true
+	} else {
+		s.sysmsgThreadsMu.Unlock()
 	}
+	s.contextQueue.wakeupSysmsgThread()
 
-	s.sysmsgThreadsMu.Unlock()
 	return false
 }
 

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

@@ -364,12 +364,12 @@ void __syshandler() {
 }
 
 void __export_start(struct sysmsg *sysmsg, void *_ucontext) {
-#if defined(__x86_64__)
+  init_new_thread();
+
   asm volatile("movq %%gs:0, %0\n" : "=r"(sysmsg) : :);
   if (sysmsg->self != sysmsg) {
     panic(0xdeaddead);
   }
-#endif
 
   struct thread_context *ctx =
       switch_context_amd64(sysmsg, NULL, CONTEXT_STATE_INVALID);

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

@@ -107,6 +107,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
   struct thread_context *ctx = NULL, *old_ctx = NULL;
   if (thread_state == THREAD_STATE_INITIALIZING) {
     // Find a new context and exit to restore it.
+    init_new_thread();
     goto init;
   }
 

pkg/sentry/platform/systrap/sysmsg/sysmsg.h

@@ -171,6 +171,7 @@ struct thread_context *switch_context(struct sysmsg *sysmsg,
                                       enum context_state new_context_state);
 
 int wait_state(struct sysmsg *sysmsg, enum thread_state new_thread_state);
+void init_new_thread(void);
 
 #define panic(err) __panic(err, __LINE__)
 // NOLINTEND(runtime/int)

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

@@ -45,6 +45,7 @@ struct context_queue {
   uint32_t start;
   uint32_t end;
   uint32_t num_active_threads;
+  uint32_t num_threads_to_wakeup;
   uint32_t num_active_contexts;
   uint32_t num_awake_contexts;
   uint64_t fast_path_disalbed_ts;
@@ -152,6 +153,8 @@ static void spinning_queue_pop() {
 // spinning_queue_remove_first removes one thread from a queue that has been
 // spinning longer than others and longer than a specified timeout.
 //
+// If `timeout` is zero, it always removes one element and never returns false.
+//
 // Returns true if one thread has been removed from the queue.
 static bool spinning_queue_remove_first(uint64_t timeout)
     __attribute__((warn_unused_result));
@@ -160,16 +163,19 @@ static bool spinning_queue_remove_first(uint64_t timeout) {
   uint64_t ts;
   uint32_t idx;
 
-  idx = atomic_load(&queue->start);
-  ts = atomic_load(&queue->start_times[idx % SPINNING_QUEUE_SIZE]);
-  if (ts == 0 || rdtsc() - ts < timeout) return false;
+  while (1) {
+    idx = atomic_load(&queue->start);
+    ts = atomic_load(&queue->start_times[idx % SPINNING_QUEUE_SIZE]);
+    if (ts == 0 && timeout == 0) continue;
+    if (ts == 0 || rdtsc() - ts < timeout) return false;
 
-  // The current thread is still in a queue and the length of the queue is twice
-  // of the maximum number of threads, so we can zero the element and be sure
-  // that nobody is trying to set it in a  non-zero value.
-  atomic_store(&queue->start_times[idx % SPINNING_QUEUE_SIZE], 0);
-  if (!atomic_compare_exchange(&queue->start, &idx, idx + 1)) {
-    return false;
+    // The current thread is still in a queue and the length of the queue is
+    // twice of the maximum number of threads, so we can zero the element and be
+    // sure that nobody is trying to set it in a  non-zero value.
+    atomic_store(&queue->start_times[idx % SPINNING_QUEUE_SIZE], 0);
+    if (atomic_compare_exchange(&queue->start, &idx, idx + 1)) {
+      break;
+    }
   }
 
   return true;
@@ -232,8 +238,8 @@ static struct thread_context *get_context_fast(struct sysmsg *sysmsg,
       return ctx;
     }
 
-    if (atomic_load(&queue->fast_path_disabled) != 0 &&
-        spinning_queue_remove_first(0)) {
+    if (atomic_load(&queue->fast_path_disabled) != 0) {
+      if (!spinning_queue_remove_first(0)) panic(0);
       break;
     }
 
@@ -244,14 +250,9 @@ static struct thread_context *get_context_fast(struct sysmsg *sysmsg,
       if (atomic_compare_exchange(&queue->num_active_threads,
                                   &nr_active_threads, nr_active_threads - 1)) {
         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(&queue->num_active_threads, 1);
+        if (!spinning_queue_remove_first(0)) panic(0);
+        *nr_active_threads_p = nr_active_threads;
+        break;
       }
     }
 
@@ -269,6 +270,25 @@ static struct thread_context *get_context_fast(struct sysmsg *sysmsg,
 
 #define NR_IF_THREAD_IS_ACTIVE (~0)
 
+static bool try_to_dec_threads_to_wakeup(struct context_queue *queue) {
+  while (1) {
+    uint32_t nr = atomic_load(&queue->num_threads_to_wakeup);
+    if (nr == 0) {
+      return false;
+    }
+    if (atomic_compare_exchange(&queue->num_threads_to_wakeup, &nr, nr - 1)) {
+      return true;
+    };
+  }
+}
+
+void init_new_thread() {
+  struct context_queue *queue = __export_context_queue_addr;
+
+  atomic_add(&queue->num_active_threads, 1);
+  try_to_dec_threads_to_wakeup(queue);
+}
+
 // 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) {
@@ -317,7 +337,7 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
     // * 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) {
+    if (nr_active_threads < nr_active_contexts) {
       ctx = queue_get_context(sysmsg);
       if (ctx) {
         atomic_store(&sysmsg->state, THREAD_STATE_PREP);
@@ -326,10 +346,19 @@ struct thread_context *get_context(struct sysmsg *sysmsg) {
       }
     }
 
-    while (atomic_load(&sysmsg->state) == THREAD_STATE_ASLEEP) {
-      sys_futex(&sysmsg->state, FUTEX_WAIT, THREAD_STATE_ASLEEP, NULL, NULL, 0);
+    while (1) {
+      if (!try_to_dec_threads_to_wakeup(queue)) {
+        sys_futex(&queue->num_threads_to_wakeup, FUTEX_WAIT, 0, NULL, NULL, 0);
+        continue;
+      }
+      // Mark this thread as being active only if it can get a context.
+      ctx = queue_get_context(sysmsg);
+      if (ctx) {
+        atomic_store(&sysmsg->state, THREAD_STATE_PREP);
+        atomic_add(&queue->num_active_threads, 1);
+        return ctx;
+      }
     }
-    atomic_add(&queue->num_active_threads, 1);
   }
 }