Implement systrap context queue.

This is the initial implementation of the systrap context queue via a ringbuffer
in shared memory between stub threads and the sentry.

In this new model there is no longer a bound sysmsg thread for every context;
instead each subprocess starts with one initial sysmsg thread, which starts
polling the context queue for new contexts arriving from the sentry. If the
sentry detects that contexts are spending too much time in the context queue
without being processed, it will create new sysmsg threads or wake sleeping
ones. Tangentially, sysmsg threads will go to sleep if they spend too much time
busy looping without new context arrivals.

This model does not yet take into account the full load of the host system or
even multiple subprocesses in the same sandbox. Multiple overloaded subprocesses
are liable to make each other run slower by kicking sysmsg threads more often
than they need to; this will be remedied in follow up CLs.

PiperOrigin-RevId: 516680504

pkg/sentry/platform/systrap/BUILD

@@ -22,6 +22,7 @@ go_library(
     srcs = [
         "context_decoupling_disable.go",
         "context_decoupling_enable.go",
+        "context_queue.go",
         "filters.go",
         "filters_amd64.go",
         "filters_arm64.go",

pkg/sentry/platform/systrap/context_queue.go

@@ -0,0 +1,86 @@
+// Copyright 2023 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 (
+	"sync/atomic"
+)
+
+// LINT.IfChange
+const (
+	// maxEntries is the size of the ringbuffer.
+	maxContextQueueEntries uint32 = uint32(maxGuestContexts) + 1
+)
+
+type queuedContext struct {
+	contextID uint32
+	threadID  uint32
+}
+
+// contextQueue is a structure shared with the each stub thread that is used to
+// signal to stub threads which contexts are ready to resume running.
+//
+// It is a lockless ringbuffer where threads try to police themselves on whether
+// they should continue waiting for a context or go to sleep if they are
+// unneeded.
+type contextQueue struct {
+	// start is an index used for taking contexts out of the ringbuffer.
+	start uint32
+	// end is an index used for putting new contexts into the ringbuffer.
+	end uint32
+	// stubPollingIndex is used by stubs to indicate polling order.
+	stubPollingIndex uint32
+	// 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
+	// ringbuffer is the mmapped region of memory that's shared with the stub
+	// threads.
+	ringbuffer [maxContextQueueEntries]uint32
+}
+
+// LINT.ThenChange(./sysmsg/sysmsg_lib.c)
+
+func (q *contextQueue) init() {
+	for i := uint32(0); i < maxContextQueueEntries; i++ {
+		q.ringbuffer[i] = invalidContextID
+	}
+	atomic.StoreUint32(&q.start, 0)
+	atomic.StoreUint32(&q.end, 0)
+	atomic.StoreUint32(&q.stubPollingIndex, 0)
+	atomic.StoreUint32(&q.stubPollingIndexBase, 0)
+	atomic.StoreUint32(&q.numSleepingThreads, 0)
+}
+
+func (q *contextQueue) isEmpty() bool {
+	return atomic.LoadUint32(&q.start) == atomic.LoadUint32(&q.end)
+}
+
+func (q *contextQueue) queuedContexts() uint32 {
+	return (atomic.LoadUint32(&q.end) + maxContextQueueEntries - atomic.LoadUint32(&q.start)) % maxContextQueueEntries
+}
+
+func (q *contextQueue) add(contextID uint32) uint32 {
+	next := atomic.AddUint32(&q.end, 1)
+	if (next % maxContextQueueEntries) ==
+		(atomic.LoadUint32(&q.start) % maxContextQueueEntries) {
+		// should be unreacheable
+		panic("contextQueue is full")
+	}
+	next = (next - 1) % maxContextQueueEntries
+	atomic.StoreUint32(&q.ringbuffer[next], contextID)
+	return next // remove me
+}

pkg/sentry/platform/systrap/stub_unsafe.go

@@ -109,6 +109,9 @@ func stubInit() {
 	// |--------stubSysmsgStack-------------|
 	// | Reserved space for per-thread      |
 	// | sysmsg stacks.                     |
+	// |----------stubContextQueue----------|
+	// | Shared ringbuffer queue for stubs  |
+	// | to select the next context.        |
 	// |--------stubThreadContextRegion-----|
 	// | Reserved space for thread contexts |
 	// *------------------------------------*
@@ -135,6 +138,14 @@ func stubInit() {
 	// has to be aligned to sysmsg.PerThreadMemSize.
 	// Look at sysmsg/sighandler.c:sysmsg_addr() for more details.
 	mapLen, _ = hostarch.PageRoundUp(mapLen + sysmsg.PerThreadMemSize*(maxSystemThreads+1))
+
+	// Allocate context queue region
+	if contextDecouplingExp {
+		stubContextQueueRegion = mapLen
+		stubContextQueueRegionLen, _ = hostarch.PageRoundUp(unsafe.Sizeof(contextQueue{}))
+		mapLen += stubContextQueueRegionLen
+	}
+
 	// Allocate thread context region
 	stubContextRegion = mapLen
 	stubContextRegionLen = sysmsg.AllocatedSizeofThreadContextStruct * (maxGuestContexts + 1)
@@ -177,6 +188,7 @@ func stubInit() {
 	// Randomize stubSysmsgStack address.
 	gap := uintptr(rand.Uint64()) * hostarch.PageSize % (maximumUserAddress - stubStart - mapLen)
 	stubSysmsgStack += uintptr(gap)
+	stubContextQueueRegion += uintptr(gap)
 	stubContextRegion += uintptr(gap)
 
 	// Copy the stub to the address.
@@ -187,6 +199,7 @@ func stubInit() {
 	stubSysmsgStart += stubStart
 	stubSysmsgStack += stubStart
 	stubROMapEnd += stubStart
+	stubContextQueueRegion += stubStart
 	stubContextRegion += stubStart
 
 	// Align stubSysmsgStack to the per-thread stack size.
@@ -209,6 +222,8 @@ func stubInit() {
 	exp := (*uint64)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_context_decoupling_exp)))
 	if contextDecouplingExp {
 		*exp = 1
+		contextQueue := (*uint64)(unsafe.Pointer(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_context_queue_addr)))
+		*contextQueue = uint64(stubContextQueueRegion)
 	}
 
 	prepareSeccompRules(stubSysmsgStart, stubSysmsgRules, stubSysmsgRulesLen)
@@ -224,7 +239,7 @@ func stubInit() {
 
 	// Set the end.
 	stubEnd = stubStart + mapLen + uintptr(gap)
-	log.Debugf("stubStart %x stubSysmsgStart %x stubSysmsgStack %x, stubThreadContextRegion %x, mapLen %x", stubStart, stubSysmsgStart, stubSysmsgStack, stubContextRegion, mapLen)
+	log.Debugf("stubStart %x stubSysmsgStart %x stubSysmsgStack %x, stubContextQueue %x, stubThreadContextRegion %x, mapLen %x", stubStart, stubSysmsgStart, stubSysmsgStack, stubContextQueueRegion, stubContextRegion, mapLen)
 	log.Debugf(archState.String())
 	log.Debugf("contextDecouplingExp=%t", contextDecouplingExp)
 }

pkg/sentry/platform/systrap/subprocess.go

@@ -19,6 +19,7 @@ import (
 	"os"
 	"runtime"
 	"sync"
+	"sync/atomic"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/abi/linux"
@@ -97,6 +98,11 @@ type requestStub struct {
 	done chan *thread
 }
 
+// maxSysmsgThreads specifies the maximum number of system threads that a
+// subprocess can create in context decoupled mode.
+// TODO(b/268366549): Replace maxSystemThreads below.
+var maxSysmsgThreads = runtime.GOMAXPROCS(0)
+
 const (
 	// maxSystemThreads specifies the maximum number of system threads that a
 	// subprocess may create in order to process the contexts.
@@ -119,6 +125,9 @@ type subprocess struct {
 	// requests is used to signal creation of new threads.
 	requests chan any
 
+	// sysmsgInitRegs is used to reset sysemu regs.
+	sysmsgInitRegs arch.Registers
+
 	// mu protects the following fields.
 	mu sync.Mutex
 
@@ -147,8 +156,19 @@ type subprocess struct {
 	syscallThreadMu sync.Mutex
 	syscallThread   *syscallThread
 
+	// sysmsgThreadsMu protects sysmsgThreads and numSysmsgThreads
 	sysmsgThreadsMu sync.Mutex
-	sysmsgThreads   map[uint32]*sysmsgThread
+	// sysmsgThreads is a collection of all active sysmsg threads in the
+	// subprocess.
+	sysmsgThreads map[uint32]*sysmsgThread
+	// numSysmsgThreads counts the number of active sysmsg threads; we use a
+	// counter instead of using len(sysmsgThreads) because we need to synchronize
+	// how many threads get created _before_ the creation happens.
+	numSysmsgThreads int
+
+	// contextQueue is a queue of all contexts that are ready to switch back to
+	// user mode.
+	contextQueue *contextQueue
 }
 
 func (s *subprocess) initSyscallThread(ptraceThread *thread) error {
@@ -266,11 +286,12 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 	runtime.LockOSThread()
 	defer runtime.UnlockOSThread()
 
-	// Initialize the first thread.
+	// Initialize the syscall thread.
 	ptraceThread, err := create()
 	if err != nil {
 		return nil, err
 	}
+	sp.sysmsgInitRegs = ptraceThread.initRegs
 
 	if err := sp.initSyscallThread(ptraceThread); err != nil {
 		return nil, err
@@ -291,6 +312,14 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 	sp.usertrap = usertrap.New()
 	sp.mapSharedRegions()
 
+	// Create the initial sysmsg thread.
+	if contextDecouplingExp {
+		if _, err := sp.createSysmsgThread(nil, nil, nil); err != nil {
+			return nil, err
+		}
+		sp.numSysmsgThreads++
+	}
+
 	return sp, nil
 }
 
@@ -301,7 +330,7 @@ func newSubprocess(create func() (*thread, error), memoryFile *pgalloc.MemoryFil
 // Should be called before any sysmsg threads are created.
 // Initializes s.contextQueue and s.threadContextRegion.
 func (s *subprocess) mapSharedRegions() {
-	if s.threadContextRegion != 0 {
+	if s.contextQueue != nil || s.threadContextRegion != 0 {
 		panic("contextQueue or threadContextRegion was already initialized")
 	}
 
@@ -310,6 +339,27 @@ func (s *subprocess) mapSharedRegions() {
 		Dir:  pgalloc.TopDown,
 	}
 
+	if contextDecouplingExp {
+		// Map shared regions into the sentry.
+		contextQueueFR, contextQueue := mmapContextQueueForSentry(s.memoryFile, opts)
+		contextQueue.init()
+
+		// Map thread context region into the syscall thread.
+		_, err := s.syscallThread.syscall(
+			unix.SYS_MMAP,
+			arch.SyscallArgument{Value: uintptr(stubContextQueueRegion)},
+			arch.SyscallArgument{Value: uintptr(contextQueueFR.Length())},
+			arch.SyscallArgument{Value: uintptr(unix.PROT_READ | unix.PROT_WRITE)},
+			arch.SyscallArgument{Value: uintptr(unix.MAP_SHARED | unix.MAP_FILE | unix.MAP_FIXED)},
+			arch.SyscallArgument{Value: uintptr(s.memoryFile.FD())},
+			arch.SyscallArgument{Value: uintptr(contextQueueFR.Start)})
+		if err != nil {
+			panic(fmt.Sprintf("failed to mmap context queue region into syscall thread: %v", err))
+		}
+
+		s.contextQueue = contextQueue
+	}
+
 	// Map thread context region into the sentry.
 	threadContextFR, err := s.memoryFile.Allocate(uint64(stubContextRegionLen), opts)
 	if err != nil {
@@ -327,8 +377,6 @@ func (s *subprocess) mapSharedRegions() {
 	}
 
 	// Map thread context region into the syscall thread.
-	// Map shared regions that will be the same and used for all sysmsg threads
-	// in this subprocess.
 	if _, err := s.syscallThread.syscall(
 		unix.SYS_MMAP,
 		arch.SyscallArgument{Value: uintptr(stubContextRegion)},
@@ -631,18 +679,16 @@ func (t *thread) NotifyInterrupt() {
 func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool, shouldPatchSyscall bool, err error) {
 	// Reset necessary registers.
 	regs := &ac.StateData().Regs
+	s.resetSysemuRegs(regs)
+	ctx := s.getThreadContextFromID(c.cid)
+	ctx.Regs = regs.PtraceRegs
+	restoreArchSpecificState(ctx, ac)
+
+	// Get sysmsg thread bound to the context; no-op if contextDecoupling is on.
 	sysThread, err := s.getSysmsgThread(regs, c, ac)
 	if err != nil {
 		return false, false, err
 	}
-	msg := sysThread.msg
-	ctx := s.getThreadContextFromID(c.cid)
-	t := sysThread.thread
-	t.resetSysemuRegs(regs)
-
-	s.restoreFPState(msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
-	ctx.Regs = regs.PtraceRegs
-	restoreArchSpecificState(regs, t, sysThread, msg, ac)
 
 	// Check for interrupts, and ensure that future interrupts signal the context.
 	if !c.interrupt.Enable(c) {
@@ -653,11 +699,46 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 	}
 	defer c.interrupt.Disable()
 
-	msg.EnableSentryFastPath()
-	sysThread.waitEvent(sysmsg.ThreadStateDone)
+	if contextDecouplingExp {
+		s.restoreFPState(nil, ctx, 0, c, ac)
 
-	if msg.Err != 0 {
-		panic(fmt.Sprintf("stub thread %d failed: err %d line %d: %s", t.tid, msg.Err, msg.Line, msg))
+		// Place the context onto the context queue.
+		ctx.State.Set(sysmsg.ContextStateNone)
+		s.contextQueue.add(uint32(c.cid))
+		s.waitOnState(ctx)
+
+		// Check if there's been an error.
+		tid := atomic.LoadUint32(&ctx.ThreadID)
+		if tid != invalidThreadID {
+			if sysThread, ok := s.sysmsgThreads[tid]; ok && sysThread.msg.Err != 0 {
+				msg := sysThread.msg
+				panic(fmt.Sprintf("stub thread %d failed: err 0x%x line %d: %s", sysThread.thread.tid, msg.Err, msg.Line, msg))
+			}
+			log.Warningf("systrap: found unexpected ThreadContext.ThreadID field, expected %d found %d", invalidThreadID, tid)
+		}
+	} else {
+		msg := sysThread.msg
+		t := sysThread.thread
+
+		s.restoreFPState(msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
+
+		msg.EnableSentryFastPath()
+		sysThread.waitEvent(sysmsg.ThreadStateDone)
+
+		// Check if there's been an error.
+		if msg.Err != 0 {
+			panic(fmt.Sprintf("stub thread %d failed: err %d line %d: %s", t.tid, msg.Err, msg.Line, msg))
+		}
+
+		if ctx.State != sysmsg.ContextStateSyscallTrap {
+			var err error
+			sysThread.fpuStateToMsgOffset, err = msg.FPUStateOffset()
+			if err != nil {
+				return false, false, err
+			}
+		}
+
+		retrieveArchSpecificState(ctx, ac)
 	}
 
 	regs.PtraceRegs = ctx.Regs
@@ -665,16 +746,6 @@ func (s *subprocess) switchToApp(c *context, ac *arch.Context64) (isSyscall bool
 	// either delivered from the kernel or from this process. We
 	// don't respect other signals.
 	c.signalInfo = ctx.SignalInfo
-	if !contextDecouplingExp && ctx.State != sysmsg.ContextStateSyscallTrap {
-		var err error
-		sysThread.fpuStateToMsgOffset, err = msg.FPUStateOffset()
-		if err != nil {
-			return false, false, err
-		}
-	}
-
-	retrieveArchSpecificState(regs, msg, t, ac)
-
 	if ctx.State == sysmsg.ContextStateSyscallCanBePatched {
 		ctx.State = sysmsg.ContextStateSyscall
 		shouldPatchSyscall = true
@@ -693,6 +764,80 @@ 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 *sysmsg.ThreadContext) {
+	// ackedEvents is always reset to 0 at the end of this function.
+	ackedEvents := uint32(0)
+	kicked := false
+	slowPath := false
+	start := cputicks()
+	handshake := false
+	for curState := ctx.State.Get(); curState == sysmsg.ContextStateNone; curState = ctx.State.Get() {
+		if !slowPath {
+			delta := uint64(cputicks() - start)
+			if delta > decoupledDeepSleepTimeout {
+				ctx.DisableSentryFastPath()
+				slowPath = true
+				continue
+			}
+
+			if !handshake && ackedEvents != atomic.LoadUint32(&ctx.Acked) {
+				handshake = true
+				continue
+			}
+			spinloop()
+		} else {
+			// If the context already received a handshake then it knows it's being
+			// worked on.
+			if !kicked && !handshake {
+				kicked = true
+				s.kickSysmsgThread()
+			}
+
+			ctx.SleepOnState(curState)
+		}
+	}
+
+	atomic.StoreUint32(&ctx.Acked, 0)
+	ctx.EnableSentryFastPath()
+}
+
+func (s *subprocess) kickSysmsgThread() {
+	s.sysmsgThreadsMu.Lock()
+
+	if atomic.LoadUint32(&s.contextQueue.numSleepingThreads) > 0 {
+		for _, t := range s.sysmsgThreads {
+			if t.msg.State.Get() == sysmsg.ThreadStateAsleep {
+				t.msg.WakeSysmsgThread()
+				s.sysmsgThreadsMu.Unlock()
+				return
+			}
+		}
+	}
+	// 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()
+		if _, err := s.createSysmsgThread(nil, nil, nil); err != nil {
+			s.sysmsgThreadsMu.Lock()
+			s.numSysmsgThreads--
+			s.sysmsgThreadsMu.Unlock()
+		}
+	} else {
+		s.sysmsgThreadsMu.Unlock()
+	}
+}
+
 // syscall executes the given system call without handling interruptions.
 func (s *subprocess) syscall(sysno uintptr, args ...arch.SyscallArgument) (uintptr, error) {
 	s.syscallThreadMu.Lock()
@@ -753,16 +898,24 @@ func (s *subprocess) PullFullState(c *context, ac *arch.Context64) error {
 		panic("Attempted to PullFullState for context that is not used in subprocess")
 	}
 	ctx := s.getThreadContextFromID(c.cid)
-	sysThread, err := s.getSysmsgThread(&ac.StateData().Regs, c, ac)
-	if err != nil {
-		return err
+	if contextDecouplingExp {
+		s.saveFPState(nil, ctx, 0, c, ac)
+	} else {
+		sysThread, err := s.getSysmsgThread(&ac.StateData().Regs, c, ac)
+		if err != nil {
+			return err
+		}
+		s.saveFPState(sysThread.msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
 	}
-	s.saveFPState(sysThread.msg, ctx, sysThread.fpuStateToMsgOffset, c, ac)
 	return nil
 }
 
 // getSysmsgThread returns a sysmsg thread for the specified context.
+// (Unused if contextDecouplingExp=true).
 func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch.Context64) (*sysmsgThread, error) {
+	if contextDecouplingExp {
+		return nil, nil
+	}
 	sysThread := c.sysmsgThread
 	if sysThread != nil && sysThread.subproc != s {
 		// This can happen if a new address space
@@ -773,6 +926,19 @@ func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch
 	if sysThread != nil {
 		return sysThread, nil
 	}
+	return s.createSysmsgThread(tregs, c, ac)
+}
+
+// createSysmsgThread creates a new sysmsg thread.
+// If contextDecouplingExp=false, the thread starts working on the given context.
+// Otherwise the given function parameters are not used, and the thread starts
+// processing any available context in the context queue.
+func (s *subprocess) createSysmsgThread(tregs *arch.Registers, c *context, ac *arch.Context64) (*sysmsgThread, error) {
+	if contextDecouplingExp {
+		// We will not bind any specific context to this thread. We will still use
+		// tregs to setup the thread though.
+		tregs = &arch.Registers{}
+	}
 
 	// Create a new seccomp process.
 	var r requestThread
@@ -803,12 +969,14 @@ func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch
 		// TODO(b/144063246): Need to fail the clone system call.
 		panic(fmt.Sprintf("failed to allocate a new stack: %v", err))
 	}
-	sysThread = &sysmsgThread{
+	sysThread := &sysmsgThread{
 		thread:     p,
 		subproc:    s,
 		stackRange: fr,
 	}
-	tid := uint32(p.tid)
+	// Use the sysmsgStackID as a handle on this thread instead of host tid in
+	// order to be able to reliably specify invalidThreadID.
+	threadID := uint32(p.sysmsgStackID)
 
 	// Map the stack into the sentry.
 	sentryStackAddr, _, errno := unix.RawSyscall6(
@@ -851,15 +1019,19 @@ func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch
 	}
 
 	sysThread.setMsg(sysmsg.StackAddrToMsg(sentryStackAddr))
-	sysThread.msg.Init(tid)
-	s.getThreadContextFromID(c.cid).ThreadID = tid
-	sysThread.msg.ContextID = c.cid
+	sysThread.msg.Init(threadID)
+	if contextDecouplingExp {
+		sysThread.msg.ContextID = uint64(invalidContextID)
+	} else {
+		s.getThreadContextFromID(c.cid).ThreadID = threadID
+		sysThread.msg.ContextID = c.cid
+	}
 	sysThread.msg.Self = uint64(sysmsgStackAddr + sysmsg.MsgOffsetFromSharedStack)
 	sysThread.msg.SyshandlerStack = uint64(sysmsg.StackAddrToSyshandlerStack(sysThread.sysmsgPerThreadMemAddr()))
 	sysThread.msg.ContextRegion = uint64(stubContextRegion)
 	sysThread.msg.Syshandler = uint64(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_syshandler))
 
-	sysThread.msg.State.Set(sysmsg.ThreadStateDone)
+	sysThread.msg.State.Set(sysmsg.ThreadStateInitializing)
 
 	// Install a pre-compiled seccomp rules for the BPF process.
 	_, err = p.syscallIgnoreInterrupt(&p.initRegs, unix.SYS_PRCTL,
@@ -879,15 +1051,12 @@ func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch
 	}
 
 	// Prepare to start the BPF process.
-	p.resetSysemuRegs(tregs)
-	archSpecificSysThreadInit(sysThread, tregs)
+	s.resetSysemuRegs(tregs)
+	setArchSpecificRegs(sysThread, tregs)
 	if err := p.setRegs(tregs); err != nil {
 		panic(fmt.Sprintf("ptrace set regs failed: %v", err))
 	}
-	// Send a fake event to stop the BPF process.
-	if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(p.tgid), uintptr(p.tid), uintptr(unix.SIGSEGV)); e != 0 {
-		panic(fmt.Sprintf("tkill failed: %v", e))
-	}
+	archSpecificSysmsgThreadInit(sysThread)
 	// Skip SIGSTOP.
 	if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(p.tgid), uintptr(p.tid), uintptr(unix.SIGCONT)); e != 0 {
 		panic(fmt.Sprintf("tkill failed: %v", e))
@@ -897,23 +1066,23 @@ func (s *subprocess) getSysmsgThread(tregs *arch.Registers, c *context, ac *arch
 		panic(fmt.Sprintf("can't detach new clone: %v", errno))
 	}
 
-	sysThread.waitEvent(sysmsg.ThreadStateNone)
-	if msg := sysThread.msg; msg.Err != 0 {
-		panic(fmt.Sprintf("stub thread failed: %v (line %v)", msg.Err, msg.Line))
-	}
-
 	if !contextDecouplingExp {
+		sysThread.waitEvent(sysmsg.ThreadStateNone)
+		if msg := sysThread.msg; msg.Err != 0 {
+			panic(fmt.Sprintf("stub thread failed: %v (line %v)", msg.Err, msg.Line))
+		}
+
 		sysThread.fpuStateToMsgOffset, err = sysThread.msg.FPUStateOffset()
 		if err != nil {
 			sysThread.destroy()
 			return nil, err
 		}
+
+		c.sysmsgThread = sysThread
 	}
 
-	c.sysmsgThread = sysThread
-
 	s.sysmsgThreadsMu.Lock()
-	s.sysmsgThreads[tid] = sysThread
+	s.sysmsgThreads[threadID] = sysThread
 	s.sysmsgThreadsMu.Unlock()
 
 	return sysThread, nil
@@ -960,6 +1129,7 @@ func (s *subprocess) registerContext(c *context) error {
 	s.IncRef()
 	c.cid = id
 	c.subprocess = s
+	c.FullStateChanged()
 	unlock()
 
 	threadContext := s.getThreadContextFromID(id)

pkg/sentry/platform/systrap/subprocess_amd64.go

@@ -19,7 +19,6 @@ package systrap
 
 import (
 	"fmt"
-	"runtime"
 	"strings"
 
 	"golang.org/x/sys/unix"
@@ -37,13 +36,13 @@ const (
 // resetSysemuRegs sets up emulation registers.
 //
 // This should be called prior to calling sysemu.
-func (t *thread) resetSysemuRegs(regs *arch.Registers) {
-	regs.Cs = t.initRegs.Cs
-	regs.Ss = t.initRegs.Ss
-	regs.Ds = t.initRegs.Ds
-	regs.Es = t.initRegs.Es
-	regs.Fs = t.initRegs.Fs
-	regs.Gs = t.initRegs.Gs
+func (s *subprocess) resetSysemuRegs(regs *arch.Registers) {
+	regs.Cs = s.sysmsgInitRegs.Cs
+	regs.Ss = s.sysmsgInitRegs.Ss
+	regs.Ds = s.sysmsgInitRegs.Ds
+	regs.Es = s.sysmsgInitRegs.Es
+	regs.Fs = s.sysmsgInitRegs.Fs
+	regs.Gs = s.sysmsgInitRegs.Gs
 }
 
 // createSyscallRegs sets up syscall registers.
@@ -211,38 +210,31 @@ func appendArchSeccompRules(rules []seccomp.RuleSet) []seccomp.RuleSet {
 	}...)
 }
 
-func restoreArchSpecificState(regs *arch.Registers, t *thread, sysThread *sysmsgThread, msg *sysmsg.Msg, _ *arch.Context64) {
-	regs.Gs_base = msg.Self
+func restoreArchSpecificState(ctx *sysmsg.ThreadContext, ac *arch.Context64) {
+}
 
-	// Switching gs_base is a rare operation, therefore checking that we need to do
-	// so is better done in the sentry, because doing so on a host that doesn't
-	// have FSGSBASE instructions enabled is quite expensive since it would require
-	// an ARCH_PRCTL syscall.
-	if regs.Gs_base != sysThread.gsBase {
-		runtime.LockOSThread()
-		defer runtime.UnlockOSThread()
+func setArchSpecificRegs(sysThread *sysmsgThread, regs *arch.Registers) {
+	if contextDecouplingExp {
+		// Set the start function and initial stack.
+		regs.PtraceRegs.Rip = uint64(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_start))
+		regs.PtraceRegs.Rsp = uint64(sysmsg.StackAddrToSyshandlerStack(sysThread.sysmsgPerThreadMemAddr()))
+	}
 
-		t.attach()
+	// Set gs_base; this is the only time we set it and we don't expect it to ever
+	// change for any thread.
+	regs.Gs_base = sysThread.msg.Self
+}
 
-		var r arch.Registers
-		if err := t.getRegs(&r); err != nil {
-			panic(fmt.Sprintf("ptrace get regs failed: %v", err))
+func retrieveArchSpecificState(ctx *sysmsg.ThreadContext, ac *arch.Context64) {
+}
+
+func archSpecificSysmsgThreadInit(sysThread *sysmsgThread) {
+	// Send a fake event to stop the BPF process so that it enters the sighandler.
+	// If there is no coupled context we don't want that to happen because the
+	// thread needs to find a context first.
+	if !contextDecouplingExp {
+		if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(sysThread.thread.tgid), uintptr(sysThread.thread.tid), uintptr(unix.SIGSEGV)); e != 0 {
+			panic(fmt.Sprintf("tkill failed: %v", e))
 		}
-		r.Gs_base = regs.Gs_base
-		if err := t.setRegs(&r); err != nil {
-			panic(fmt.Sprintf("ptrace set regs failed: %v", err))
-		}
-		if _, _, errno := unix.RawSyscall6(unix.SYS_PTRACE, unix.PTRACE_DETACH, uintptr(t.tid), 0, 0, 0, 0); errno != 0 {
-			panic(fmt.Sprintf("ptrace detach failed: %v", errno))
-		}
-		sysThread.gsBase = regs.Gs_base
 	}
 }
-
-func archSpecificSysThreadInit(sysThread *sysmsgThread, regs *arch.Registers) {
-	regs.Gs_base = sysThread.msg.Self
-	sysThread.gsBase = regs.Gs_base
-}
-
-func retrieveArchSpecificState(regs *arch.Registers, msg *sysmsg.Msg, _ *thread, ac *arch.Context64) {
-}

pkg/sentry/platform/systrap/subprocess_arm64.go

@@ -36,7 +36,7 @@ const (
 // resetSysemuRegs sets up emulation registers.
 //
 // This should be called prior to calling sysemu.
-func (t *thread) resetSysemuRegs(regs *arch.Registers) {
+func (s *subprocess) resetSysemuRegs(regs *arch.Registers) {
 }
 
 // createSyscallRegs sets up syscall registers.
@@ -187,15 +187,27 @@ func (s *subprocess) arm64SyscallWorkaround(t *thread, regs *arch.Registers) {
 	}
 }
 
-func restoreArchSpecificState(regs *arch.Registers, t *thread, _ *sysmsgThread, msg *sysmsg.Msg, ac *arch.Context64) {
-	msg.TLS = uint64(ac.TLS())
+func restoreArchSpecificState(ctx *sysmsg.ThreadContext, ac *arch.Context64) {
+	ctx.TLS = uint64(ac.TLS())
 }
 
-func archSpecificSysThreadInit(sysThread *sysmsgThread, regs *arch.Registers) {
+func setArchSpecificRegs(sysThread *sysmsgThread, regs *arch.Registers) {
+	if contextDecouplingExp {
+		// Set the start function and initial stack.
+		regs.PtraceRegs.Pc = uint64(stubSysmsgStart + uintptr(sysmsg.Sighandler_blob_offset____export_start))
+		regs.PtraceRegs.Sp = uint64(sysmsg.StackAddrToSyshandlerStack(sysThread.sysmsgPerThreadMemAddr()))
+	}
 }
 
-func retrieveArchSpecificState(regs *arch.Registers, msg *sysmsg.Msg, t *thread, ac *arch.Context64) {
-	if !ac.SetTLS(uintptr(msg.TLS)) {
-		panic(fmt.Sprintf("ac.SetTLS(%+v) failed", msg.TLS))
+func retrieveArchSpecificState(ctx *sysmsg.ThreadContext, ac *arch.Context64) {
+	if !ac.SetTLS(uintptr(ctx.TLS)) {
+		panic(fmt.Sprintf("ac.SetTLS(%+v) failed", ctx.TLS))
+	}
+}
+
+func archSpecificSysmsgThreadInit(sysThread *sysmsgThread) {
+	// Send a fake event to stop the BPF process so that it enters the sighandler.
+	if _, _, e := unix.RawSyscall(unix.SYS_TGKILL, uintptr(sysThread.thread.tgid), uintptr(sysThread.thread.tid), uintptr(unix.SIGSEGV)); e != 0 {
+		panic(fmt.Sprintf("tkill failed: %v", e))
 	}
 }

pkg/sentry/platform/systrap/subprocess_unsafe.go

@@ -22,8 +22,12 @@
 package systrap
 
 import (
+	"fmt"
 	"unsafe"
 
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
 )
 
@@ -45,3 +49,22 @@ func (s *subprocess) getThreadContextFromID(cid uint64) *sysmsg.ThreadContext {
 	tcSlot := s.threadContextRegion + uintptr(cid)*sysmsg.AllocatedSizeofThreadContextStruct
 	return (*sysmsg.ThreadContext)(unsafe.Pointer(tcSlot))
 }
+
+func mmapContextQueueForSentry(memoryFile *pgalloc.MemoryFile, opts pgalloc.AllocOpts) (memmap.FileRange, *contextQueue) {
+	fr, err := memoryFile.Allocate(uint64(stubContextQueueRegionLen), opts)
+	if err != nil {
+		panic(fmt.Sprintf("failed to allocate a new subprocess context memory region"))
+	}
+	addr, _, errno := unix.RawSyscall6(
+		unix.SYS_MMAP,
+		0,
+		uintptr(fr.Length()),
+		unix.PROT_WRITE|unix.PROT_READ,
+		unix.MAP_SHARED|unix.MAP_FILE,
+		uintptr(memoryFile.FD()), uintptr(fr.Start))
+	if errno != 0 {
+		panic(fmt.Sprintf("mmap failed for subprocess context memory region: %v", errno))
+	}
+
+	return fr, (*contextQueue)(unsafe.Pointer(addr))
+}

pkg/sentry/platform/systrap/sysmsg/BUILD

@@ -118,6 +118,7 @@ go_library(
         "sysmsg.go",
         "sysmsg_amd64.go",
         "sysmsg_arm64.go",
+        "sysmsg_unsafe.go",
         ":sighandler_go_arch",
     ],
     embedsrcs = [
@@ -130,5 +131,6 @@ go_library(
         "//pkg/cpuid",
         "//pkg/errors",
         "//pkg/hostarch",
+        "@org_golang_x_sys//unix:go_default_library",
     ],
 )

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

@@ -54,15 +54,6 @@ long sys_futex(uint32_t *addr, int op, int val, struct __kernel_timespec *tv,
                    (long)addr2, (long)val3);
 }
 
-void check_sysmsg_thread_context(struct sysmsg *sysmsg,
-                                 struct thread_context **ctx) {
-  struct thread_context *new_ctx = thread_context_addr(sysmsg);
-  if (*ctx != new_ctx) {
-    *ctx = new_ctx;
-    __atomic_store_n(&(*ctx)->fpstate_changed, 1, __ATOMIC_RELEASE);
-  }
-}
-
 union csgsfs {
   uint64_t csgsfs;  // REG_CSGSFS
   struct {
@@ -164,19 +155,61 @@ static void set_fsbase(struct user_regs_struct *ptregs) {
   }
 }
 
+// switch_context_amd64 is a wrapper of switch_context() which does checks
+// specific to amd64.
+struct thread_context *switch_context_amd64(
+    struct sysmsg *sysmsg, struct thread_context *ctx,
+    enum thread_state new_thread_state, enum context_state new_context_state) {
+  get_fsbase(&ctx->ptregs);
+  long fs_base = ctx->ptregs.fs_base;
+
+  for (;;) {
+    // TODO(b/271631387): Once stub code globals can be used between objects
+    // move this check into sysmsg_lib:switch_context().
+    if (__export_context_decoupling_exp) {
+      ctx = switch_context(sysmsg, ctx, new_context_state);
+    } else {
+      ctx->state = new_context_state;
+      wait_state(sysmsg, new_thread_state);
+    }
+
+    if (__atomic_load_n(&ctx->interrupt, __ATOMIC_ACQUIRE) != 0) {
+      // This context got interrupted while it was waiting in the queue.
+      // Setup all the necessary bits to let the sentry know this context has
+      // switched back because of it.
+      __atomic_store_n(&ctx->interrupt, 0, __ATOMIC_RELEASE);
+      new_context_state = CONTEXT_STATE_FAULT;
+      ctx->signo = SIGCHLD;
+      ctx->siginfo.si_signo = SIGCHLD;
+      ctx->ptregs.orig_rax = -1;
+    } else {
+      break;
+    }
+  }
+  if (fs_base != ctx->ptregs.fs_base) {
+    set_fsbase(&ctx->ptregs);
+  }
+  return ctx;
+}
+
 void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
   ucontext_t *ucontext = _ucontext;
   void *sp = sysmsg_sp();
   struct sysmsg *sysmsg = sysmsg_addr(sp);
 
   if (sysmsg != sysmsg->self) panic(0xdeaddead);
+  int32_t thread_state = __atomic_load_n(&sysmsg->state, __ATOMIC_ACQUIRE);
+  if (__export_context_decoupling_exp &&
+      thread_state == THREAD_STATE_INITIALIZING) {
+    // This thread was interrupted before it even had a context.
+    return;
+  }
+
   struct thread_context *ctx = thread_context_addr(sysmsg);
 
   if (signo == SIGCHLD) {
     // If the current thread is in syshandler, an interrupt has to be postponed,
     // because sysmsg can't be changed.
-    int32_t thread_state;
-    thread_state = __atomic_load_n(&sysmsg->state, __ATOMIC_ACQUIRE);
     if (thread_state != THREAD_STATE_NONE) {
       // There are two possibilities for when we received the interrupt:
       //   1. Before syshandler switched to the sentry.
@@ -221,6 +254,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     return;
   }
 
+  enum context_state ctx_state = CONTEXT_STATE_INVALID;
   ctx->signo = signo;
   ctx->siginfo = *siginfo;
   gregs_to_ptregs(ucontext, &ctx->ptregs);
@@ -237,7 +271,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     case SIGSYS: {
       int si_sysno = siginfo->si_syscall;
       int i;
-      ctx->state = CONTEXT_STATE_SYSCALL;
+      ctx_state = CONTEXT_STATE_SYSCALL;
 
       // Check whether this syscall can be replaced on a function call or not.
       // If a syscall instruction set is "mov sysno, %eax, syscall", it can be
@@ -292,7 +326,7 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
 
         if (need_trap) {
           // This syscall can be replaced on the function call.
-          ctx->state = CONTEXT_STATE_SYSCALL_NEED_TRAP;
+          ctx_state = CONTEXT_STATE_SYSCALL_NEED_TRAP;
         }
       }
       ctx->ptregs.orig_rax = ctx->ptregs.rax;
@@ -310,19 +344,14 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     case SIGTRAP:
     case SIGILL:
       ctx->ptregs.orig_rax = -1;
-      ctx->state = CONTEXT_STATE_FAULT;
+      ctx_state = CONTEXT_STATE_FAULT;
       break;
     default:
       return;
   }
-  get_fsbase(&ctx->ptregs);
-  long fs_base = ctx->ptregs.fs_base;
 
-  wait_state(sysmsg, THREAD_STATE_EVENT);
+  ctx = switch_context_amd64(sysmsg, ctx, THREAD_STATE_EVENT, ctx_state);
 
-  if (fs_base != __atomic_load_n(&ctx->ptregs.fs_base, __ATOMIC_ACQUIRE)) {
-    set_fsbase(&ctx->ptregs);
-  }
   if (__export_context_decoupling_exp &&
       __atomic_load_n(&ctx->fpstate_changed, __ATOMIC_ACQUIRE)) {
     memcpy((uint8_t *)ucontext->uc_mcontext.fpregs, ctx->fpstate,
@@ -342,22 +371,23 @@ void __syshandler() {
 
   struct thread_context *ctx = thread_context_addr(sysmsg);
 
-  ctx->state = CONTEXT_STATE_SYSCALL_TRAP;
+  enum context_state ctx_state = CONTEXT_STATE_SYSCALL_TRAP;
   ctx->signo = SIGSYS;
   ctx->siginfo.si_addr = 0;
   ctx->siginfo.si_syscall = ctx->ptregs.rax;
   ctx->ptregs.rax = (unsigned long)-ENOSYS;
   __atomic_store_n(&sysmsg->interrupt, 0, __ATOMIC_RELAXED);
 
-  get_fsbase(&ctx->ptregs);
-  long fs_base = ctx->ptregs.fs_base;
+  switch_context_amd64(sysmsg, ctx, THREAD_STATE_EVENT, ctx_state);
+}
 
-  state = wait_state(sysmsg, THREAD_STATE_EVENT);
+// asm_restore_state is implemented in syshandler_amd64.S
+void asm_restore_state();
 
-  // Restore state
-  if (fs_base != ctx->ptregs.fs_base) {
-    set_fsbase(&ctx->ptregs);
-  }
+// On x86 restore_state jumps straight to user code and does not return.
+void restore_state(struct sysmsg *sysmsg, struct thread_context *ctx, void *) {
+  set_fsbase(&ctx->ptregs);
+  asm_restore_state();
 }
 
 void verify_offsets_amd64() {

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

@@ -96,8 +96,17 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
   struct sysmsg *sysmsg = sysmsg_addr(sp);
 
   if (sysmsg != sysmsg->self) panic(0xdeaddead);
+  int32_t thread_state = __atomic_load_n(&sysmsg->state, __ATOMIC_ACQUIRE);
+  if (__export_context_decoupling_exp &&
+      thread_state == THREAD_STATE_INITIALIZING) {
+    // Find a new context and exit to restore it.
+    __export_start(sysmsg, _ucontext);
+    return;
+  }
+
   struct thread_context *ctx = thread_context_addr(sysmsg);
 
+  uint32_t ctx_state = CONTEXT_STATE_INVALID;
   ctx->signo = signo;
 
   gregs_to_ptregs(ucontext, &ctx->ptregs);
@@ -118,11 +127,11 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
   } else {
     sysmsg->fpstate = (uint64_t)(fpStatePointer) - (uint64_t)sysmsg;
   }
-  sysmsg->tls = get_tls();
+  ctx->tls = get_tls();
   ctx->siginfo = *siginfo;
   switch (signo) {
     case SIGSYS: {
-      ctx->state = CONTEXT_STATE_SYSCALL;
+      ctx_state = CONTEXT_STATE_SYSCALL;
       if (siginfo->si_arch != AUDIT_ARCH_AARCH64) {
         // gVisor doesn't support x32 system calls, so let's change the syscall
         // number so that it returns ENOSYS. The value added here is just a
@@ -138,19 +147,48 @@ void __export_sighandler(int signo, siginfo_t *siginfo, void *_ucontext) {
     case SIGFPE:
     case SIGTRAP:
     case SIGILL:
-      ctx->state = CONTEXT_STATE_FAULT;
+      ctx_state = CONTEXT_STATE_FAULT;
       break;
     default:
       return;
   }
 
-  wait_state(sysmsg, THREAD_STATE_EVENT);
+  for (;;) {
+    if (__export_context_decoupling_exp) {
+      ctx = switch_context(sysmsg, ctx, ctx_state);
+    } else {
+      ctx->state = ctx_state;
+      wait_state(sysmsg, THREAD_STATE_EVENT);
+    }
+
+    if (__atomic_load_n(&ctx->interrupt, __ATOMIC_ACQUIRE) != 0) {
+      // This context got interrupted while it was waiting in the queue.
+      // Setup all the necessary bits to let the sentry know this context has
+      // switched back because of it.
+      __atomic_store_n(&ctx->interrupt, 0, __ATOMIC_RELEASE);
+      ctx_state = CONTEXT_STATE_FAULT;
+      ctx->signo = SIGCHLD;
+      ctx->siginfo.si_signo = SIGCHLD;
+    } else {
+      break;
+    }
+  }
+  restore_state(sysmsg, ctx, _ucontext);
+}
+
+// On ARM restore_state sets up a correct restore from the sighandler by
+// populating _ucontext.
+void restore_state(struct sysmsg *sysmsg, struct thread_context *ctx,
+                   void *_ucontext) {
+  ucontext_t *ucontext = _ucontext;
+  struct fpsimd_context *fpctx = &ucontext->uc_mcontext.__reserved;
+  uint8_t *fpStatePointer = (uint8_t *)&fpctx->fpsr;
 
   if (__export_context_decoupling_exp &&
       __atomic_load_n(&ctx->fpstate_changed, __ATOMIC_ACQUIRE)) {
     memcpy(fpStatePointer, ctx->fpstate, __export_arch_state.fp_len);
   }
   ptregs_to_gregs(ucontext, &ctx->ptregs);
-  set_tls(sysmsg->tls);
+  set_tls(ctx->tls);
   __atomic_store_n(&sysmsg->state, THREAD_STATE_NONE, __ATOMIC_RELEASE);
 }

pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S

@@ -196,6 +196,9 @@ __export_syshandler:
 
   callq __syshandler
 
+.globl asm_restore_state;
+.type asm_restore_state, @function;
+asm_restore_state:
   // thread_context may have changed, therefore we reload it into %rcx anew.
   load_thread_context_addr
   restore_fpstate

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

@@ -112,6 +112,16 @@ const (
 	// that there is a postponed interrupt from the syshandler.
 	// The sentry should never see this event.
 	ThreadStateInterrupt
+	// ThreadStateContextRestore means that the thread is in the process of doing
+	// a context restore.
+	ThreadStateContextRestore
+	// ThreadStateAsleep means that this thread fell asleep because there was not
+	// enough contexts to process in the context queue.
+	ThreadStateAsleep
+	// ThreadStateInitializing is only set once at sysmsg thread creation time. It
+	// is used to tell the signal handler that the thread does not yet have a
+	// context.
+	ThreadStateInitializing
 )
 
 // Msg contains the current state of the sysmsg thread.
@@ -162,9 +172,6 @@ type Msg struct {
 	// fpState is an offset relative to the sighandler stack to the fpState, stored
 	// by the sighandler.
 	fpState uint64
-	// TLS is a pointer to a thread local storage.
-	// It is is only populated on ARM64.
-	TLS uint64
 	// The fast path is the mode when a thread is polling msg->state to
 	// wait for a required state instead of calling FUTEX_WAIT.
 	//
@@ -227,7 +234,7 @@ const (
 const (
 	// MaxFPStateLen is the largest possible FPState that we will save.
 	// Note: This value was chosen to be able to fit ThreadContext into one page.
-	MaxFPStateLen uint32 = 3648
+	MaxFPStateLen uint32 = 3584
 
 	// AllocatedSizeofThreadContextStruct defines how much memory to allocate for
 	// one instance of ThreadContext.
@@ -265,6 +272,21 @@ type ThreadContext struct {
 	// ThreadID is the ID of the sysmsg thread that's currently working on the
 	// context.
 	ThreadID uint32
+	// LastThreadID is the ID of the previous sysmsg thread that ran the context
+	// (not the one currently working on it). This field is used by sysmsg threads
+	// to detect whether fpstate may have changed since the last time they ran a
+	// context.
+	LastThreadID uint32
+	// SentryFastPath is used to indicate to the stub thread that the sentry
+	// goroutine used for this thread context is busy-polling for a response
+	// instead of using FUTEX_WAIT.
+	SentryFastPath uint32
+	// Acked is used by sysmsg threads to signal to the sentry that this context
+	// has been picked up from the context queue and is actively being worked on.
+	Acked uint32
+	// TLS is a pointer to a thread local storage.
+	// It is is only populated on ARM64.
+	TLS uint64
 	// Debug is a variable to use to get visibility into the stub from the sentry.
 	Debug uint64
 }
@@ -277,6 +299,7 @@ func (m *Msg) Init(threadID uint32) {
 	m.Line = -1
 	m.stubFastPath = 0
 	m.sentryFastPath = 1
+	m.ThreadID = threadID
 }
 
 // Init initializes the ThreadContext instance.
@@ -301,15 +324,35 @@ func (m *Msg) DisableStubFastPath() {
 
 // EnableSentryFastPath enables the polling mode for the Sentry. It has to be
 // called before switching controls to the stub process.
+// This function is used if contextDecouplingExp=false because the fastpath
+// is negotiated in Sysmsg.
 func (m *Msg) EnableSentryFastPath() {
 	m.sentryFastPath = 1
 }
 
 // DisableSentryFastPath disables the polling mode for the Sentry.
+// This function is used if contextDecouplingExp=false because the fastpath
+// is negotiated in Sysmsg.
 func (m *Msg) DisableSentryFastPath() {
 	atomic.StoreUint32(&m.sentryFastPath, 0)
 }
 
+// EnableSentryFastPath indicates that the polling mode is enabled for the
+// Sentry. It has to be called before putting the context into the context queue.
+// This function is used if contextDecouplingExp=true because the fastpath
+// is negotiated in ThreadContext
+func (c *ThreadContext) EnableSentryFastPath() {
+	c.SentryFastPath = 1
+}
+
+// DisableSentryFastPath indicates that the polling mode for the sentry is
+// disabled for the Sentry.
+// This function is used if contextDecouplingExp=true because the fastpath
+// is negotiated in ThreadContext.
+func (c *ThreadContext) DisableSentryFastPath() {
+	atomic.StoreUint32(&c.SentryFastPath, 0)
+}
+
 // FPUStateOffset returns the offset of a saved FPU state to the msg.
 func (m *Msg) FPUStateOffset() (uint64, error) {
 	offset := m.fpState

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

@@ -37,12 +37,15 @@ struct arch_state {
 #endif
 
 // LINT.IfChange
-enum {
+enum thread_state {
   THREAD_STATE_NONE,
   THREAD_STATE_DONE,
   THREAD_STATE_EVENT,
   THREAD_STATE_PREP,
   THREAD_STATE_INTERRUPT,
+  THREAD_STATE_CONTEXT_RESTORE,
+  THREAD_STATE_ASLEEP,
+  THREAD_STATE_INITIALIZING,
 };
 
 // sysmsg contains the current state of the sysmsg thread. See: sysmsg.go:Msg
@@ -65,8 +68,6 @@ struct sysmsg {
   int32_t err_line;
   uint64_t debug;
   uint64_t fpstate;
-  // tls is only populated on ARM64.
-  uint64_t tls;
   uint32_t stub_fast_path;
   uint32_t sentry_fast_path;
   uint32_t acked_events;
@@ -96,6 +97,10 @@ struct thread_context {
   uint32_t state;
   uint32_t interrupt;
   uint32_t thread_id;
+  uint32_t last_thread_id;
+  uint32_t sentry_fast_path;
+  uint32_t acked;
+  uint64_t tls;
   uint64_t debug;
 };
 
@@ -118,6 +123,7 @@ extern uint64_t __export_pr_sched_core;
 extern uint64_t __export_deep_sleep_timeout;
 extern struct arch_state __export_arch_state;
 extern uint64_t __export_context_decoupling_exp;
+extern uint64_t __export_context_queue_addr;
 
 // NOLINTBEGIN(runtime/int)
 static void *sysmsg_sp() {
@@ -151,10 +157,15 @@ long sys_futex(uint32_t *addr, int op, int val, struct __kernel_timespec *tv,
 static void __panic(int err, long line) {
   void *sp = sysmsg_sp();
   struct sysmsg *sysmsg = sysmsg_addr(sp);
+  struct thread_context *ctx = thread_context_addr(sysmsg);
   sysmsg->err = err;
   sysmsg->err_line = line;
+  // Normally sentry waits on sysmsg->state.
   __atomic_store_n(&sysmsg->state, THREAD_STATE_EVENT, __ATOMIC_RELEASE);
   sys_futex(&sysmsg->state, FUTEX_WAKE, 1, NULL, NULL, 666);
+  // Under context-decoupling the sentry waits on ctx->state.
+  __atomic_store_n(&ctx->state, CONTEXT_STATE_FAULT, __ATOMIC_RELEASE);
+  sys_futex(&ctx->state, FUTEX_WAKE, 1, NULL, NULL, 666);
   // crash the stub process.
   //
   // Normal user processes cannot map addresses lower than vm.mmap_min_addr
@@ -165,7 +176,16 @@ static void __panic(int err, long line) {
 
 void memcpy(uint8_t *dest, uint8_t *src, size_t n);
 
-int wait_state(struct sysmsg *sysmsg, uint32_t state);
+void __export_start(struct sysmsg *sysmsg, void *_ucontext);
+
+void restore_state(struct sysmsg *sysmsg, struct thread_context *ctx,
+                   void *_ucontext);
+
+struct thread_context *switch_context(struct sysmsg *sysmsg,
+                                      struct thread_context *ctx,
+                                      enum context_state new_context_state);
+
+int wait_state(struct sysmsg *sysmsg, enum thread_state new_thread_state);
 
 #define panic(err) __panic(err, __LINE__)
 // NOLINTEND(runtime/int)

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

@@ -27,25 +27,36 @@
 // polling and fall asleep.
 uint64_t __export_deep_sleep_timeout;
 uint64_t __export_handshake_timeout;
+uint64_t __export_context_queue_addr;
 
-// A per-thread memory region is always align to STACK_SIZE.
-// *------------*
-// | guard page |
-// |------------|
-// | syshandler |
-// |   stack    |
-// |            |
-// |------------|
-// | guard page |
-// |------------|
-// |            |
-// |     ^      |
-// |    / \     |
-// |     |      |
-// |  altstack  |
-// |------------|
-// |   sysmsg   |
-// *------------*
+// LINT.IfChange
+#define MAX_STUB_THREADS (4096)
+#define MAX_CONTEXT_QUEUE_ENTRIES (MAX_STUB_THREADS + 1)
+#define INVALID_CONTEXT_ID (MAX_STUB_THREADS + 1)
+#define INVALID_THREAD_ID (MAX_STUB_THREADS + 1)
+
+// See systrap/context_queue.go
+struct context_queue {
+  uint32_t start;
+  uint32_t end;
+  uint32_t polling_index;
+  uint32_t polling_index_base;
+  uint32_t num_sleeping_threads;
+  uint32_t ringbuffer[MAX_CONTEXT_QUEUE_ENTRIES];
+};
+// LINT.ThenChange(../context_queue.go)
+
+uint32_t is_empty(struct context_queue *queue) {
+  return __atomic_load_n(&queue->start, __ATOMIC_ACQUIRE) ==
+         __atomic_load_n(&queue->end, __ATOMIC_ACQUIRE);
+}
+
+int32_t queued_contexts(struct context_queue *queue) {
+  return (__atomic_load_n(&queue->end, __ATOMIC_ACQUIRE) +
+          MAX_CONTEXT_QUEUE_ENTRIES -
+          __atomic_load_n(&queue->start, __ATOMIC_ACQUIRE)) %
+         MAX_CONTEXT_QUEUE_ENTRIES;
+}
 
 #if defined(__x86_64__)
 static __inline__ unsigned long rdtsc(void) {
@@ -71,7 +82,92 @@ void memcpy(uint8_t *dest, uint8_t *src, size_t n) {
   }
 }
 
-int wait_state(struct sysmsg *sysmsg, uint32_t state) {
+// 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 =
+      (struct context_queue *)(__export_context_queue_addr);
+  for (;;) {
+    // Change sysmsg thread state just to indicate thread is not asleep.
+    __atomic_store_n(&sysmsg->state, THREAD_STATE_PREP, __ATOMIC_RELEASE);
+    unsigned long start = rdtsc();
+    for (;;) {
+      if (!is_empty(queue)) {
+        uint32_t next = __atomic_load_n(&queue->start, __ATOMIC_ACQUIRE) %
+                        MAX_CONTEXT_QUEUE_ENTRIES;
+        uint32_t context_id = __atomic_exchange_n(
+            &queue->ringbuffer[next], INVALID_CONTEXT_ID, __ATOMIC_ACQ_REL);
+        if (context_id != INVALID_CONTEXT_ID) {
+          __atomic_add_fetch(&queue->start, 1, __ATOMIC_ACQ_REL);
+          if (context_id > MAX_STUB_THREADS) {
+            panic(context_id);
+          }
+          sysmsg->context_id = context_id;
+          struct thread_context *ctx = thread_context_addr(sysmsg);
+          __atomic_store_n(&ctx->acked, 1, __ATOMIC_RELEASE);
+          __atomic_store_n(&ctx->thread_id, sysmsg->thread_id,
+                           __ATOMIC_RELEASE);
+          return ctx;
+        } else {
+          continue;
+        }
+      }
+      if ((rdtsc() - start) > __export_deep_sleep_timeout) {
+        break;
+      }
+
+      spinloop();
+    }
+    __atomic_store_n(&sysmsg->state, THREAD_STATE_ASLEEP, __ATOMIC_RELEASE);
+
+    __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);
+  }
+}
+
+// switch_context signals the sentry that the old context is ready to be worked
+// on and retrieves a new context to switch to.
+struct thread_context *switch_context(struct sysmsg *sysmsg,
+                                      struct thread_context *ctx,
+                                      enum context_state new_context_state) {
+  __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);
+  if (__atomic_load_n(&ctx->sentry_fast_path, __ATOMIC_ACQUIRE) == 0) {
+    int ret = sys_futex(&ctx->state, FUTEX_WAKE, 1, NULL, NULL, 0);
+    if (ret < 0) {
+      panic(ret);
+    }
+  }
+  uint32_t old_ctx_id = sysmsg->context_id;
+
+  ctx = get_context(sysmsg);
+
+  if (old_ctx_id != sysmsg->context_id ||
+      ctx->last_thread_id != sysmsg->thread_id) {
+    ctx->fpstate_changed = 1;
+  }
+
+  return ctx;
+}
+
+void __export_start(struct sysmsg *sysmsg, void *_ucontext) {
+#if defined(__x86_64__)
+  asm volatile("movq %%gs:0, %0\n" : "=r"(sysmsg) : :);
+  if (sysmsg->self != sysmsg) {
+    panic(0xdeaddead);
+  }
+#endif
+
+  struct thread_context *ctx = get_context(sysmsg);
+  __atomic_store_n(&ctx->fpstate_changed, 1, __ATOMIC_RELEASE);
+  __atomic_store_n(&ctx->thread_id, sysmsg->thread_id, __ATOMIC_RELEASE);
+
+  restore_state(sysmsg, ctx, _ucontext);
+}
+
+int wait_state(struct sysmsg *sysmsg, enum thread_state new_thread_state) {
   unsigned long handshake_timeout;
   uint64_t acked_events_prev;
   unsigned long start;
@@ -81,7 +177,7 @@ int wait_state(struct sysmsg *sysmsg, uint32_t state) {
   // stub_fast_path can be changed non-atomically before we change the state and
   // wake up the Sentry.
   sysmsg->stub_fast_path = 1;
-  __atomic_store_n(&sysmsg->state, state, __ATOMIC_SEQ_CST);
+  __atomic_store_n(&sysmsg->state, new_thread_state, __ATOMIC_SEQ_CST);
 
   fast_path = __atomic_load_n(&sysmsg->sentry_fast_path, __ATOMIC_SEQ_CST);
   if (!fast_path) {

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

@@ -21,7 +21,7 @@
 #define FAULT_OPCODE 0x06
 
 // LINT.IfChange
-#define MAX_FPSTATE_LEN 3648
+#define MAX_FPSTATE_LEN 3584
 // Note: To be explicit, 2^12 = 4096; if ALLOCATED_SIZEOF_THREAD_CONTEXT_STRUCT
 //       is changed, make sure to change the code that relies on the bitshift.
 #define ALLOCATED_SIZEOF_THREAD_CONTEXT_STRUCT 4096

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

@@ -0,0 +1,40 @@
+// Copyright 2023 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 sysmsg
+
+import (
+	"fmt"
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/abi/linux"
+)
+
+// SleepOnState makes the caller sleep on the ThreadContext.State futex.
+func (c *ThreadContext) SleepOnState(curState ContextState) {
+	_, _, errno := unix.Syscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&c.State)),
+		linux.FUTEX_WAIT, uintptr(curState), 0, 0, 0)
+	if errno != 0 && errno != unix.EAGAIN && errno != unix.EINTR {
+		panic(fmt.Sprintf("error waiting for state: %v", errno))
+	}
+}
+
+// WakeSysmsgThread calls futex wake on Sysmsg.State.
+func (m *Msg) WakeSysmsgThread() syscall.Errno {
+	m.State.Set(ThreadStatePrep)
+	_, _, e := unix.RawSyscall6(unix.SYS_FUTEX, uintptr(unsafe.Pointer(&m.State)), linux.FUTEX_WAKE, 1, 0, 0, 0)
+	return e
+}

pkg/sentry/platform/systrap/sysmsg_thread.go

@@ -47,10 +47,6 @@ type sysmsgThread struct {
 	// context is the last context that ran on this thread.
 	context *context
 
-	// gsBase contains previous values of gs_base register to follow
-	// changes, because it's not restored by the kernel from a signal frame.
-	gsBase uint64
-
 	// stackRange is a sysmsg stack in the memory file.
 	stackRange memmap.FileRange
 

pkg/sentry/platform/systrap/systrap.go

@@ -80,6 +80,9 @@ var (
 	stubSysmsgStack uintptr
 	stubSysmsgStart uintptr
 	stubSysmsgEnd   uintptr
+	// Memory region to store the contextQueue.
+	stubContextQueueRegion    uintptr
+	stubContextQueueRegionLen uintptr
 	// Memory region to store instances of sysmsg.ThreadContext.
 	stubContextRegion    uintptr
 	stubContextRegionLen uintptr
@@ -131,7 +134,7 @@ type context struct {
 	lastFaultIP hostarch.Addr
 
 	// sysmsgThread is a sysmsg thread descriptor which is used to execute
-	// application code.
+	// application code. (Note: Unused if contextDecouplingExp=true).
 	sysmsgThread *sysmsgThread
 
 	// fpLen is the size of the floating point context.
@@ -281,10 +284,10 @@ func (c *context) Interrupt() {
 // NotifyInterrupt implements interrupt.Receiver.NotifyInterrupt.
 //
 // Another reasonable existing object to implement NotifyInterrupt would be
-// sysmsg.ThreadContext, because it already has the correct tid written into it
-// to know which thread to send the signal to. However we cannot do that because
-// it is in shared memory, which means that one subprocess can overwrite it to
-// have the sentry send an interrupt to a completely different subprocess.
+// sysmsg.ThreadContext, because we can write the correct host TID into it
+// to know which thread to send the signal to. However, because it is in shared
+// memory, one subprocess can overwrite it to have the sentry send an interrupt
+// to a completely different subprocess.
 // For this reason we use systrap.context and check that the target thread
 // is actually valid within the subprocess.
 func (c *context) NotifyInterrupt() {