systrap: poll all contexts from one goroutine

Before this change, each task is polling its contexts. It reduces the number
of context switches but it creates extra CPU load.

This idea has been borrowed from Jamie's cl/488537531.

PiperOrigin-RevId: 528984717

pkg/sentry/platform/systrap/BUILD

@@ -6,6 +6,18 @@ package(
     licenses = ["notice"],
 )
 
+go_template_instance(
+    name = "context_list",
+    out = "context_list.go",
+    package = "systrap",
+    prefix = "context",
+    template = "//pkg/ilist:generic_list",
+    types = {
+        "Element": "*sharedContext",
+        "Linker": "*sharedContext",
+    },
+)
+
 go_template_instance(
     name = "subprocess_refs",
     out = "subprocess_refs.go",
@@ -22,6 +34,7 @@ go_library(
     srcs = [
         "context_decoupling_disable.go",
         "context_decoupling_enable.go",
+        "context_list.go",
         "context_queue.go",
         "filters.go",
         "filters_amd64.go",
@@ -78,6 +91,7 @@ go_library(
         "//pkg/sentry/platform/systrap/usertrap",
         "//pkg/sentry/usage",
         "//pkg/sync",
+        "//pkg/syncevent",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )

pkg/sentry/platform/systrap/shared_context.go

@@ -16,11 +16,14 @@ package systrap
 
 import (
 	"fmt"
+	"sync"
 	"sync/atomic"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
+	gsync "gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/syncevent"
 )
 
 const (
@@ -36,21 +39,38 @@ const (
 //     (trusted) synchronization between the sentry and the stub processes.
 //   - Data read from shared memory may require validation before it can be used.
 type sharedContext struct {
+	contextEntry
+
 	// subprocess is the subprocess that this sharedContext instance belongs to.
 	subprocess *subprocess
 	// contextID is the ID corresponding to the sysmsg.ThreadContext memory slot
 	// that is used for this sharedContext.
 	contextID uint32
-	// shared is the handle to the shared memory that the sentry task goroutine
+	// shared is the handle to the shared memory that the sentry task go-routine
 	// reads from and writes to.
 	// NOTE: Using this handle directly without a getter from this function should
 	//       most likely be avoided due to concerns listed above.
 	shared *sysmsg.ThreadContext
 
-	fastPathFailedInRow uint32
-	fastPathDisabledTS  uint64
+	// sync is used by the context go-routine to wait for events from the
+	// dispatcher.
+	sync           syncevent.Waiter
+	startWaitingTS int64
+	kicked         bool
 }
 
+const (
+	// sharedContextReady indicates that a context has new events.
+	sharedContextReady = syncevent.Set(1 << iota)
+	// sharedContextKicked indicates that a new stub thread should be woken up.
+	sharedContextKicked
+	// sharedContextSlowPath indicates that a context has to be waited for in the
+	// slow path.
+	sharedContextSlowPath
+	// sharedContextDispatch indicates that a context go-routine has to start the wait loop.
+	sharedContextDispatch
+)
+
 func (s *subprocess) getSharedContext() (*sharedContext, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -66,6 +86,7 @@ func (s *subprocess) getSharedContext() (*sharedContext, error) {
 		shared:     s.getThreadContextFromID(id),
 	}
 	sc.shared.Init(invalidThreadID)
+	sc.sync.Init()
 
 	return &sc, nil
 }
@@ -170,3 +191,115 @@ func (sc *sharedContext) sleepOnState(state sysmsg.ContextState) {
 		panic(fmt.Sprintf("error waiting for state: %v", errno))
 	}
 }
+
+type fastPathContextQueue struct {
+
+	// list is used only from the loop method and so it isn't protected by
+	// any lock.
+	list contextList
+
+	mu sync.Mutex
+
+	// nr is the number of contexts in the queue.
+	// +checklocks:mu
+	nr int
+
+	// entrants contains new contexts that haven't been added to `list` yet.
+	// +checklocks:mu
+	entrants contextList
+}
+
+var dispatcher fastPathContextQueue
+
+// loop is processing contexts in the queue. Only one instance of it can be
+// running, because it has exclusive access to the list.
+//
+// target is the context associated with the current go-routine.
+func (q *fastPathContextQueue) loop(target *sharedContext) {
+	done := false
+	processed := 0
+	slowPath := false
+	start := cputicks()
+	for {
+		var ctx, next *sharedContext
+
+		q.mu.Lock()
+		if processed != 0 || !q.entrants.Empty() {
+			start = cputicks()
+			slowPath = false
+		}
+		q.nr -= processed
+		// Add new contexts to the list.
+		q.list.PushBackList(&q.entrants)
+		ctx = q.list.Front()
+		q.mu.Unlock()
+
+		if done {
+			if ctx != nil {
+				// Wake up the next go-routine to run the loop.
+				ctx.sync.Receiver().Notify(sharedContextDispatch)
+			}
+			break
+		}
+
+		processed = 0
+		now := cputicks()
+		for ctx = q.list.Front(); ctx != nil; ctx = next {
+			next = ctx.Next()
+
+			event := sharedContextReady
+			if ctx.state() == sysmsg.ContextStateNone {
+				if slowPath {
+					event = sharedContextSlowPath
+				} else if !ctx.kicked && uint64(now-ctx.startWaitingTS) > handshakeTimeout {
+					if ctx.isAcked() {
+						ctx.kicked = true
+						continue
+					}
+					event = sharedContextKicked
+				} else {
+					continue
+				}
+			}
+			processed++
+			q.list.Remove(ctx)
+			if ctx == target {
+				done = true
+			}
+			ctx.sync.Receiver().Notify(event)
+		}
+		if processed == 0 {
+			if uint64(cputicks()-start) > deepSleepTimeout {
+				slowPath = true
+				// Do one more run to notify all contexts.
+				// q.list has to be empty at the end.
+				continue
+			}
+			gsync.Goyield()
+		}
+	}
+}
+
+func (q *fastPathContextQueue) waitFor(ctx *sharedContext) syncevent.Set {
+	events := syncevent.Set(0)
+
+	q.mu.Lock()
+	q.entrants.PushBack(ctx)
+	q.nr++
+	if q.nr == 1 {
+		events = sharedContextDispatch
+	}
+	q.mu.Unlock()
+
+	for {
+		if events&sharedContextDispatch != 0 {
+			ctx.sync.Ack(sharedContextDispatch)
+			q.loop(ctx)
+		}
+		events = ctx.sync.WaitAndAckAll()
+		if events&sharedContextDispatch == 0 {
+			break
+		}
+	}
+	return events
+}

pkg/sentry/platform/systrap/subprocess.go

@@ -34,7 +34,6 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/sysmsg"
 	"gvisor.dev/gvisor/pkg/sentry/platform/systrap/usertrap"
 	"gvisor.dev/gvisor/pkg/sentry/usage"
-	gsync "gvisor.dev/gvisor/pkg/sync"
 )
 
 var (
@@ -795,63 +794,44 @@ const (
 )
 
 func (s *subprocess) waitOnState(ctx *sharedContext) {
-	kicked := false
+	ctx.kicked = false
 	slowPath := false
 	start := cputicks()
-	handshake := false
+	ctx.startWaitingTS = start
 	if atomic.LoadUint32(&s.contextQueue.numActiveThreads) == 0 {
-		kicked = s.kickSysmsgThread()
-	}
-	fastPathEnabled := true
-	if ctx.fastPathDisabledTS != 0 {
-		if uint64(cputicks())-ctx.fastPathDisabledTS < fastPathDisabledTimeout {
-			fastPathEnabled = false
-		} else {
-			ctx.fastPathFailedInRow = 0
-			ctx.fastPathDisabledTS = 0
-		}
+		ctx.kicked = s.kickSysmsgThread()
 	}
 	for curState := ctx.state(); curState == sysmsg.ContextStateNone; curState = ctx.state() {
 		if !slowPath {
-			delta := uint64(cputicks() - start)
-			if !handshake {
-				if ctx.isAcked() {
-					handshake = true
+			events := dispatcher.waitFor(ctx)
+			if events&sharedContextKicked != 0 {
+				if ctx.kicked {
 					continue
 				}
-				if !kicked && delta > handshakeTimeout {
-					kicked = s.kickSysmsgThread()
+				if ctx.isAcked() {
+					ctx.kicked = true
+					continue
 				}
+				s.kickSysmsgThread()
+				ctx.kicked = true
+				continue
 			}
-			if !fastPathEnabled || delta > deepSleepTimeout {
+			if events&sharedContextSlowPath != 0 {
 				ctx.disableSentryFastPath()
 				slowPath = true
 				continue
 			}
-
-			gsync.Goyield()
 		} else {
 			// If the context already received a handshake then it knows it's being
 			// worked on.
-			if !kicked && !handshake {
-				kicked = s.kickSysmsgThread()
+			if !ctx.kicked && !ctx.isAcked() {
+				ctx.kicked = s.kickSysmsgThread()
 			}
 
 			ctx.sleepOnState(curState)
 		}
 	}
 
-	if fastPathEnabled {
-		if slowPath {
-			ctx.fastPathFailedInRow++
-			if ctx.fastPathFailedInRow > fastPathFailedInRowLimit {
-				ctx.fastPathDisabledTS = uint64(cputicks())
-			}
-		} else {
-			ctx.fastPathFailedInRow = 0
-		}
-	}
-
 	ctx.resetAcked()
 	ctx.enableSentryFastPath()
 }