kvm: RSEQ: Add RSEQ support

This change adds RSEQ support for platform/kvm.

To support this, the KVM platform must provide two key capabilities:
 (1) A stable, unique CPU ID.
 (2) A way to detect when a thread as been preempted.

This implementation provides the necessary support as follows:
CPU ID:
  - platform/kvm now advertises the KVM vCPU ID as the cpu_id

Preemption Detection:
  - Compares the last context the CPU ran against the current context being
 scheduled.
  - Compares the context's rseqCPU and CPU ID retrieved by the platform

To facilitate this, several new methods are introduced to the platform
interface and implemented by platform/kvm:
 - HasCpuNumbers()
 - NumCPUs()
 - DetectsCPUPreemption()
 - PreemptCpu()
 - PreemptAllCpus()

PiperOrigin-RevId: 823228013

Author: tranji-cloud

pkg/sentry/kernel/kernel.go

@@ -490,6 +490,11 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 	k.cpuClockTickerWakeCh = make(chan struct{}, 1)
 	k.cpuClockTickerStopCond.L = &k.runningTasksMu
 	k.applicationCores = args.ApplicationCores
+	if args.UseHostCores && k.HasCPUNumbers() {
+		args.UseHostCores = false
+		log.Infof("UseHostCores enabled but the platform implements HasCPUNumbers(): setting UseHostCores to false")
+	}
+
 	if args.UseHostCores {
 		k.useHostCores = true
 		maxCPU, err := hostcpu.MaxPossibleCPU()
@@ -502,6 +507,15 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 			k.applicationCores = minAppCores
 		}
 	}
+
+	if k.HasCPUNumbers() {
+		if k.applicationCores < uint(k.NumCPUs()) {
+			log.Infof("ApplicationCores is less than NumCPUs: %d < %d", k.applicationCores, k.NumCPUs())
+			log.Infof("Setting applicationCores to NumCPUs: %d", k.NumCPUs())
+			k.applicationCores = uint(k.NumCPUs())
+		}
+	}
+
 	k.extraAuxv = args.ExtraAuxv
 	k.vdso = args.Vdso
 	k.vdsoParams = args.VdsoParams

pkg/sentry/kernel/rseq.go

@@ -20,7 +20,6 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
-	"gvisor.dev/gvisor/pkg/sentry/hostcpu"
 	"gvisor.dev/gvisor/pkg/usermem"
 )
 
@@ -50,7 +49,7 @@ type OldRSeqCriticalRegion struct {
 
 // RSeqAvailable returns true if t supports (old and new) restartable sequences.
 func (t *Task) RSeqAvailable() bool {
-	return t.k.useHostCores && t.k.Platform.DetectsCPUPreemption()
+	return (t.k.useHostCores || t.k.Platform.HasCPUNumbers()) && t.k.Platform.DetectsCPUPreemption()
 }
 
 // SetRSeq registers addr as this thread's rseq structure.
@@ -201,7 +200,7 @@ func (t *Task) rseqUpdateCPU() error {
 		return nil
 	}
 
-	t.rseqCPU = int32(hostcpu.GetCPU())
+	t.rseqCPU = t.CPU()
 
 	// Update both CPUs, even if one fails.
 	rerr := t.rseqCopyOutCPU()

pkg/sentry/kernel/task_run.go

@@ -24,7 +24,6 @@ import (
 	"gvisor.dev/gvisor/pkg/goid"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/refs"
-	"gvisor.dev/gvisor/pkg/sentry/hostcpu"
 	"gvisor.dev/gvisor/pkg/sentry/ktime"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
@@ -207,7 +206,7 @@ func (app *runApp) execute(t *Task) taskRunState {
 	if t.rseqPreempted {
 		t.rseqPreempted = false
 		if t.rseqAddr != 0 || t.oldRSeqCPUAddr != 0 {
-			t.rseqCPU = int32(hostcpu.GetCPU())
+			t.rseqCPU = t.CPU()
 			if err := t.rseqCopyOutCPU(); err != nil {
 				t.Debugf("Failed to copy CPU to %#x for rseq: %v", t.rseqAddr, err)
 				t.forceSignal(linux.SIGSEGV, false)

pkg/sentry/kernel/task_sched.go

@@ -365,7 +365,7 @@ func (t *Task) SetCPUMask(mask sched.CPUSet) error {
 		return linuxerr.EINVAL
 	}
 
-	if t.k.useHostCores {
+	if t.k.useHostCores || t.k.Platform.HasCPUNumbers() {
 		// No-op; pretend the mask was immediately changed back.
 		return nil
 	}
@@ -383,6 +383,10 @@ func (t *Task) SetCPUMask(mask sched.CPUSet) error {
 
 // CPU returns the cpu id for a given task.
 func (t *Task) CPU() int32 {
+	if t.k.Platform.HasCPUNumbers() {
+		return t.p.LastCPUNumber()
+	}
+
 	if t.k.useHostCores {
 		return int32(hostcpu.GetCPU())
 	}

pkg/sentry/platform/kvm/context.go

@@ -16,6 +16,7 @@ package kvm
 
 import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/atomicbitops"
 	pkgcontext "gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/ring0"
@@ -36,6 +37,10 @@ type platformContext struct {
 
 	// interrupt is the interrupt platformContext.
 	interrupt interrupt.Forwarder
+
+	// lastUsedCPU is the last CPU ID used by this platformContext.
+	// It must be accessed atomically.
+	lastUsedCPU atomicbitops.Int32
 }
 
 // tryCPUIDError indicates that CPUID emulation should occur.
@@ -45,7 +50,7 @@ type tryCPUIDError struct{}
 func (tryCPUIDError) Error() string { return "cpuid emulation failed" }
 
 // Switch runs the provided platformContext in the given address space.
-func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, rseqCPU int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	localAS := as.(*addressSpace)
 
@@ -58,6 +63,19 @@ restart:
 		c.machine.Put(cpu) // Already preempted.
 		return nil, hostarch.NoAccess, platform.ErrContextInterrupt
 	}
+	// If this CPU was last used to run a different context, then we've
+	// been preempted.
+	last := cpu.lastCtx.Swap(c)
+	c.lastUsedCPU.Store(int32(cpu.id))
+	preempted := rseqCPU >= 0 && (last != c || rseqCPU != int32(cpu.id))
+	if preempted {
+		// Release resources.
+		c.machine.Put(cpu)
+
+		// All done.
+		c.interrupt.Disable()
+		return nil, hostarch.NoAccess, platform.ErrContextCPUPreempted
+	}
 
 	// Set the active address space.
 	//
@@ -136,3 +154,8 @@ func (c *platformContext) PullFullState(as platform.AddressSpace, ac *arch.Conte
 
 // PrepareSleep implements platform.Context.platform.Context.
 func (*platformContext) PrepareSleep() {}
+
+// LastCPUNumber implements platform.Context.LastCPUNumber.
+func (c *platformContext) LastCPUNumber() int32 {
+	return c.lastUsedCPU.Load()
+}

pkg/sentry/platform/kvm/kvm.go

@@ -62,8 +62,6 @@ type runData struct {
 
 // KVM represents a lightweight VM context.
 type KVM struct {
-	platform.NoCPUPreemptionDetection
-
 	// KVM never changes mm_structs.
 	platform.UseHostProcessMemoryBarrier
 
@@ -180,11 +178,52 @@ func (k *KVM) ConcurrencyCount() int {
 	return k.machine.maxVCPUs
 }
 
+// HasCPUNumbers implements platform.Platform.HasCPUNumbers.
+func (*KVM) HasCPUNumbers() bool {
+	return true
+}
+
+// NumCPUs implements platform.Platform.NumCPUs.
+func (k *KVM) NumCPUs() int32 {
+	return int32(k.machine.maxVCPUs)
+}
+
+// DetectsCPUPreemption implements platform.Platform.DetectsCPUPreemption.
+func (k *KVM) DetectsCPUPreemption() bool {
+	return true
+}
+
+// PreemptAllCPUs implements platform.Platform.PreemptAllCPUs.
+func (k *KVM) PreemptAllCPUs() error {
+	k.machine.mu.RLock()
+	defer k.machine.mu.RUnlock()
+	for _, c := range k.machine.vCPUsByID {
+		c.lastCtx.Store(nil)
+		c.BounceToHost()
+	}
+	return nil
+}
+
+// PreemptCPU implements platform.Platform.PreemptCPU.
+func (k *KVM) PreemptCPU(cpu int32) error {
+	if cpu < 0 || cpu >= k.NumCPUs() {
+		return fmt.Errorf("invalid CPU number: %d", cpu)
+	}
+	k.machine.mu.RLock()
+	defer k.machine.mu.RUnlock()
+	c := k.machine.vCPUsByID[cpu]
+	c.lastCtx.Store(nil)
+	c.BounceToHost()
+	return nil
+}
+
 // NewContext returns an interruptible context.
 func (k *KVM) NewContext(pkgcontext.Context) platform.Context {
-	return &platformContext{
+	pc := &platformContext{
 		machine: k.machine,
 	}
+	pc.lastUsedCPU.Store(0)
+	return pc
 }
 
 type constructor struct{}

pkg/sentry/platform/kvm/machine.go

@@ -216,6 +216,9 @@ type vCPU struct {
 
 	// dieState holds state related to vCPU death.
 	dieState dieState
+
+	// lastCtx is the last context that was scheduled on this vCPU
+	lastCtx atomic.Pointer[platformContext]
 }
 
 type dieState struct {

pkg/sentry/platform/platform.go

@@ -96,6 +96,14 @@ type Platform interface {
 	// NewContext returns a new execution context.
 	NewContext(context.Context) Context
 
+	// PreemptCPU causes the first following to to Context.Switch() to return
+	// ErrContextCPUPreempted.
+	//
+	// PreemptCPU is only supported if DetectsCPUPremption() && HasCPUNumbers() == true.
+	// Platforms for which this does not hold may panic if PreemptCPU is
+	// called.
+	PreemptCPU(cpu int32) error
+
 	// PreemptAllCPUs causes all concurrent calls to Context.Switch(), as well
 	// as the first following call to Context.Switch() for each Context, to
 	// return ErrContextCPUPreempted.
@@ -121,6 +129,12 @@ type Platform interface {
 	// in parallel. Concurrent calls to Context.Switch() beyond
 	// ConcurrencyCount() may block until previous calls have returned.
 	ConcurrencyCount() int
+
+	// HasCPUNumbers returns true if platform assigns CPU numbers to contexts.
+	HasCPUNumbers() bool
+
+	// NumCPUs returns the number of CPUs on the platform.
+	NumCPUs() int32
 }
 
 // NoCPUPreemptionDetection implements Platform.DetectsCPUPreemption and
@@ -137,6 +151,25 @@ func (NoCPUPreemptionDetection) PreemptAllCPUs() error {
 	panic("This platform does not support CPU preemption detection")
 }
 
+// NoCPUNumbers implements Platform.HasCPUNumbers for platforms that do
+// not support it.
+type NoCPUNumbers struct{}
+
+// HasCPUNumbers implements Platform.HasCPUNumbers.
+func (NoCPUNumbers) HasCPUNumbers() bool {
+	return false
+}
+
+// NumCPUs implements Platform.NumCPUs.
+func (NoCPUNumbers) NumCPUs() int32 {
+	panic("platform does not support CPU numbers")
+}
+
+// PreemptCPU implements Platform.PreemptCPU.
+func (NoCPUNumbers) PreemptCPU(cpu int32) error {
+	panic("platform does not support preempting a specific CPU")
+}
+
 // UseHostGlobalMemoryBarrier implements Platform.HaveGlobalMemoryBarrier and
 // Platform.GlobalMemoryBarrier by invoking equivalent functionality on the
 // host.
@@ -264,6 +297,16 @@ type Context interface {
 	// PrepareSleep() is called when the thread switches to the
 	// interruptible sleep state.
 	PrepareSleep()
+
+	// LastCPUNumber returns the last CPU number that this context was running on.
+	// If the context never ran on a CPU, it may return any valid CPU number, as long as the first
+	// call to Switch will detect that the CPU number is incorrect and return ErrContextCPUPreempted.
+	LastCPUNumber() int32
+}
+
+// LastCPUNumber implements Context.LastCPUNumber.
+func (NoCPUNumbers) LastCPUNumber() int32 {
+	panic("context does not support last CPU number")
 }
 
 // ContextError is one of the possible errors returned by Context.Switch().

pkg/sentry/platform/ptrace/ptrace.go

@@ -74,6 +74,7 @@ var (
 
 type context struct {
 	archContext
+	platform.NoCPUNumbers
 
 	// signalInfo is the signal info, if and when a signal is received.
 	signalInfo linux.SignalInfo
@@ -214,6 +215,7 @@ type PTrace struct {
 	platform.MMapMinAddr
 	platform.NoCPUPreemptionDetection
 	platform.UseHostGlobalMemoryBarrier
+	platform.NoCPUNumbers
 }
 
 // New returns a new ptrace-based implementation of the platform interface.

pkg/sentry/platform/systrap/systrap.go

@@ -125,6 +125,8 @@ var (
 
 // platformContext is an implementation of the platform context.
 type platformContext struct {
+	platform.NoCPUNumbers
+
 	// signalInfo is the signal info, if and when a signal is received.
 	signalInfo linux.SignalInfo
 
@@ -239,6 +241,7 @@ func (c *platformContext) PrepareSleep() {
 type Systrap struct {
 	platform.NoCPUPreemptionDetection
 	platform.UseHostGlobalMemoryBarrier
+	platform.NoCPUNumbers
 
 	// memoryFile is used to create a stub sysmsg stack which is shared with
 	// the Sentry. Since memoryFile is platform-private, it is never restored,