kernel: periodically preempt tasks when platform max concurrency is exceeded

Add a new method `ConcurrencyCount()` to the `platform.Platform` interface. This
method returns the maximum number of tasks that can run in parallel on the
platform. For most platforms, this is effectively infinite (math.MaxInt). For
the KVM platform, this is the number of available vCPUs.

Introduce a new method `Preempt()` to the `platform.Context` interface. The CPU
clock ticker now checks if the number of running tasks is greater than
Platform.ConcurrencyCount(). If it is, it calls `t.p.Preempt()` on each running
task to trigger a preemption.

This mechanism ensures that on platforms with a fixed number of execution units,
like KVM with a limited number of vCPUs, tasks are periodically preempted to
allow other runnable tasks to be scheduled. This prevents any single task from
monopolizing a vCPU and improves overall fairness and responsiveness.

The `Preempt()` implementation in most platforms is a no-op. In the KVM
platform, `Preempt()` is implemented by leveraging the existing interrupt
mechanism, effectively signaling the vCPU to exit.

PiperOrigin-RevId: 805926296

Author: avagin

pkg/sentry/kernel/task_sched.go

@@ -197,6 +197,7 @@ func (k *Kernel) runCPUClockTicker() {
 		allTasks []*Task
 		incTasks = make([]*Task, k.applicationCores)
 	)
+	concurrencyCount := k.ConcurrencyCount()
 
 	for {
 		// Stop CPU clocks while nothing is running.
@@ -244,9 +245,12 @@ func (k *Kernel) runCPUClockTicker() {
 		// applicationCores running tasks (and their thread groups).
 		allTasks = k.tasks.Root.TasksAppend(allTasks)
 		runningTasks := 0
+		runningAppTasks := 0
 		for _, t := range allTasks {
 			state := t.TaskGoroutineState()
-			if state != TaskGoroutineRunningApp && state != TaskGoroutineRunningSys {
+			if state == TaskGoroutineRunningApp {
+				runningAppTasks++
+			} else if state != TaskGoroutineRunningSys {
 				continue
 			}
 			if runningTasks < len(incTasks) {
@@ -259,6 +263,7 @@ func (k *Kernel) runCPUClockTicker() {
 				incTasks[i] = t
 			}
 		}
+		preempt := runningAppTasks > concurrencyCount
 		numIncTasks := min(runningTasks, len(incTasks))
 		// Shuffle incTasks to ensure that if multiple tasks are in the same
 		// thread group, then all are equally likely to be
@@ -272,6 +277,9 @@ func (k *Kernel) runCPUClockTicker() {
 				t.appCPUClock.Add(linux.ClockTick)
 				t.tg.appCPUClockLast.Store(t)
 				t.tg.appCPUClock.Add(linux.ClockTick)
+				if preempt {
+					t.p.Preempt()
+				}
 				fallthrough
 			case TaskGoroutineRunningSys:
 				t.appSysCPUClock.Add(linux.ClockTick)

pkg/sentry/platform/interrupt/interrupt.go

@@ -97,3 +97,14 @@ func (f *Forwarder) NotifyInterrupt() {
 	}
 	f.mu.Unlock()
 }
+
+// Preempt preempts the running context. Preempt is a weaker version of
+// NotifyInterrupt, it doesn't set the pending flag which is set when a context
+// isn't actually running at this moment.
+func (f *Forwarder) Preempt() {
+	f.mu.Lock()
+	if f.dst != nil {
+		f.dst.NotifyInterrupt()
+	}
+	f.mu.Unlock()
+}

pkg/sentry/platform/kvm/context.go

@@ -118,6 +118,11 @@ func (c *platformContext) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
+// Preempt implements platform.Context.Preempt.
+func (c *platformContext) Preempt() {
+	c.interrupt.Preempt()
+}
+
 // Release implements platform.Context.Release().
 func (c *platformContext) Release() {}
 

pkg/sentry/platform/kvm/kvm.go

@@ -174,6 +174,12 @@ func (k *KVM) NewAddressSpace(any) (platform.AddressSpace, <-chan struct{}, erro
 	}, nil, nil
 }
 
+// ConcurrencyCount implements platform.Platform.ConcurrencyCount.
+// KVM can't run more than maxVCPUs contexts concurrently.
+func (k *KVM) ConcurrencyCount() int {
+	return k.machine.maxVCPUs
+}
+
 // NewContext returns an interruptible context.
 func (k *KVM) NewContext(pkgcontext.Context) platform.Context {
 	return &platformContext{

pkg/sentry/platform/platform.go

@@ -116,6 +116,11 @@ type Platform interface {
 
 	// SeccompInfo returns seccomp-related information about this platform.
 	SeccompInfo() SeccompInfo
+
+	// ConcurrencyCount returns the maximum number of contexts that can run
+	// in parallel. Concurrent calls to Context.Switch() beyond
+	// ConcurrencyCount() may block until previous calls have returned.
+	ConcurrencyCount() int
 }
 
 // NoCPUPreemptionDetection implements Platform.DetectsCPUPreemption and
@@ -246,6 +251,13 @@ type Context interface {
 	// ErrContextInterrupt.
 	Interrupt()
 
+	// Preempt interrupts a concurrent call to Switch().
+	// If Platform.ConcurrencyCount() == math.MaxInt, or if the context is not
+	// running application code (e.g. it is blocked waiting for the number of
+	// running contexts to drop below Platform.ConcurrencyCount(), Preempt may
+	// have no effect.
+	Preempt()
+
 	// Release() releases any resources associated with this context.
 	Release()
 

pkg/sentry/platform/ptrace/ptrace.go

@@ -45,6 +45,8 @@
 package ptrace
 
 import (
+	"math"
+
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	pkgcontext "gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/fd"
@@ -189,6 +191,9 @@ func (c *context) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
+// Preempt implements platform.Context.Preempt.
+func (c *context) Preempt() {}
+
 // Release implements platform.Context.Release().
 func (c *context) Release() {}
 
@@ -258,6 +263,11 @@ func (p *PTrace) NewAddressSpace(any) (platform.AddressSpace, <-chan struct{}, e
 	return as, nil, err
 }
 
+// ConcurrencyCount implements platform.Platform.ConcurrencyCount.
+func (*PTrace) ConcurrencyCount() int {
+	return math.MaxInt
+}
+
 type constructor struct{}
 
 func (*constructor) New(platform.Options) (platform.Platform, error) {

pkg/sentry/platform/systrap/systrap.go

@@ -212,6 +212,9 @@ func (c *platformContext) Interrupt() {
 	c.interrupt.NotifyInterrupt()
 }
 
+// Preempt implements platform.Context.Preempt.
+func (c *platformContext) Preempt() {}
+
 // Release releases all platform resources used by the platformContext.
 func (c *platformContext) Release() {
 	if c.sharedContext != nil {
@@ -341,6 +344,11 @@ func (*Systrap) NewContext(ctx pkgcontext.Context) platform.Context {
 	}
 }
 
+// ConcurrencyCount implements platform.Platform.ConcurrencyCount.
+func (*Systrap) ConcurrencyCount() int {
+	return maxSysmsgThreads
+}
+
 type constructor struct{}
 
 func (*constructor) New(opts platform.Options) (platform.Platform, error) {