kvm: remove machine.mu

pkg/sentry/platform/kvm/BUILD

@@ -17,6 +17,18 @@ go_template_instance(
     },
 )
 
+go_template_instance(
+    name = "atomicptrmap_tid_vcpu",
+    out = "atomicptrmap_tid_vcpu_unsafe.go",
+    package = "kvm",
+    prefix = "vCPUsByTID",
+    template = "//pkg/sync/atomicptrmap:generic_atomicptrmap",
+    types = {
+        "Key": "uint64",
+        "Value": "vCPU",
+    },
+)
+
 config_setting(
     name = "debug_build",
     values = {
@@ -31,6 +43,7 @@ go_library(
         "address_space_amd64.go",
         "address_space_arm64.go",
         "atomicptr_machine_unsafe.go",
+        "atomicptrmap_tid_vcpu_unsafe.go",
         "bluepill.go",
         "bluepill_allocator.go",
         "bluepill_amd64.go",
@@ -76,6 +89,7 @@ go_library(
         "//pkg/context",
         "//pkg/cpuid",
         "//pkg/fd",
+        "//pkg/gohacks",
         "//pkg/hostarch",
         "//pkg/hostos",
         "//pkg/hostsyscall",
@@ -130,6 +144,7 @@ go_test(
         "//pkg/sentry/platform",
         "//pkg/sentry/platform/kvm/testutil",
         "//pkg/sentry/time",
+        "//pkg/sync",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
@@ -164,6 +179,7 @@ go_test(
         "//pkg/sentry/platform",
         "//pkg/sentry/platform/kvm/testutil",
         "//pkg/sentry/time",
+        "//pkg/sync",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )

pkg/sentry/platform/kvm/bluepill_unsafe.go

@@ -70,14 +70,18 @@ func bluepillGuestExit(c *vCPU, context unsafe.Pointer) {
 	// Copy out registers.
 	bluepillArchExit(c, bluepillArchContext(context))
 
-	// Return to the vCPUReady state; notify any waiters.
-	user := c.state.Load() & vCPUUser
-	switch c.state.Swap(user) {
-	case user | vCPUGuest: // Expected case.
-	case user | vCPUGuest | vCPUWaiter:
+	// Unset vCPUGuest and vCPUWaiter; notify any waiters.
+	oldState := c.state.Load()
+	if oldState&vCPUGuest == 0 {
+		throw("invalid state in kvm.bluepillGuestExit")
+	}
+	newState := oldState &^ (vCPUGuest | vCPUWaiter)
+	oldState = c.state.Swap(newState)
+	if newState&vCPUStateMask == vCPUReady {
+		c.machine.availableNotify()
+	}
+	if oldState&vCPUWaiter != 0 {
 		c.notify()
-	default:
-		throw("invalid state")
 	}
 }
 
@@ -116,13 +120,14 @@ func bluepillHandler(context unsafe.Pointer) {
 	// Sanitize the registers; interrupts must always be disabled.
 	c := bluepillArchEnter(bluepillArchContext(context))
 
-	// Mark this as guest mode.
-	switch c.state.Swap(vCPUGuest | vCPUUser) {
-	case vCPUUser: // Expected case.
-	case vCPUUser | vCPUWaiter:
+	// Set vCPUGuest; unset vCPUWaiter; notify waiters.
+	state := c.state.Load()
+	if state&(vCPUUser|vCPUGuest) != vCPUUser {
+		throw("invalid state in bluepillHandler")
+	}
+	state = c.state.Swap((state | vCPUGuest) &^ vCPUWaiter)
+	if state&vCPUWaiter != 0 {
 		c.notify()
-	default:
-		throw("invalid state")
 	}
 
 	for {

pkg/sentry/platform/kvm/config.go

@@ -18,6 +18,16 @@
 package kvm
 
 // Config sets configuration options for each platform instance.
-type Config struct{}
+type Config struct {
+	// MaxVCPUs is the maximum number of vCPUs the platform instance will
+	// create. If MaxVCPUs is 0, the platform will choose a reasonable default.
+	MaxVCPUs int
+}
 
-func (*machine) applyConfig(config *Config) error { return nil }
+func (m *machine) applyConfig(config *Config) error {
+	if config.MaxVCPUs < 0 {
+		return fmt.Errorf("invalid Config.MaxVCPUs: %d", config.MaxVCPUs)
+	}
+	m.maxVCPUs = config.MaxVCPUs
+	return nil
+}

pkg/sentry/platform/kvm/kvm.go

@@ -16,6 +16,7 @@
 package kvm
 
 import (
+	"encoding/binary"
 	"fmt"
 
 	"golang.org/x/sys/unix"
@@ -76,6 +77,11 @@ type KVM struct {
 var (
 	globalOnce sync.Once
 	globalErr  error
+
+	// vCPUStateOffset is the offset in bytes from the start of vCPU.state to
+	// the start of the lower 4 bytes of vCPU.state, which varies depending on
+	// endianness.
+	vCPUStateOffset uintptr
 )
 
 // OpenDevice opens the KVM device and returns the File.
@@ -97,6 +103,18 @@ func New(deviceFile *fd.FD, config Config) (*KVM, error) {
 
 	// Ensure global initialization is done.
 	globalOnce.Do(func() {
+		var buf [8]byte
+		binary.NativeEndian.PutUint64(buf[:], 0x1122334455667788)
+		switch binary.NativeEndian.Uint32(buf[0:4]) {
+		case 0x11223344: // big endian
+			vCPUStateOffset = 4
+		case 0x55667788: // little endian
+			vCPUStateOffset = 0
+		default:
+			globalErr = fmt.Errorf("failed to determine endianness for kvm.vCPUStateOffset")
+			return
+		}
+
 		globalErr = updateGlobalOnce(int(fd))
 	})
 	if globalErr != nil {

pkg/sentry/platform/kvm/kvm_test.go

@@ -18,6 +18,7 @@ import (
 	"math/rand"
 	"os"
 	"reflect"
+	"runtime"
 	"testing"
 	"time"
 
@@ -32,6 +33,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/pkg/sentry/platform/kvm/testutil"
 	ktime "gvisor.dev/gvisor/pkg/sentry/time"
+	"gvisor.dev/gvisor/pkg/sync"
 )
 
 // dummyFPState is initialized in TestMain.
@@ -90,7 +92,7 @@ func bluepillTest(t testHarness, fn func(*vCPU)) {
 func TestKernelSyscall(t *testing.T) {
 	bluepillTest(t, func(c *vCPU) {
 		redpill() // Leave guest mode.
-		if got := c.state.Load(); got != vCPUUser {
+		if got := c.state.Load(); got&vCPUStateMask != vCPUUser {
 			t.Errorf("vCPU not in ready state: got %v", got)
 		}
 	})
@@ -108,7 +110,7 @@ func TestKernelFault(t *testing.T) {
 	hostFault() // Ensure recovery works.
 	bluepillTest(t, func(c *vCPU) {
 		hostFault()
-		if got := c.state.Load(); got != vCPUUser {
+		if got := c.state.Load(); got&vCPUStateMask != vCPUUser {
 			t.Errorf("vCPU not in ready state: got %v", got)
 		}
 	})
@@ -310,6 +312,12 @@ func TestBounceStress(t *testing.T) {
 				time.Sleep(time.Duration(n) * time.Microsecond)
 			}
 		}
+		m := c.machine
+		// Lock to this thread so that m.Get() always gives us the same vCPU,
+		// since applicationTest() => kvmTest() assumes that we return holding
+		// c.
+		runtime.LockOSThread()
+		defer runtime.UnlockOSThread()
 		for i := 0; i < 1000; i++ {
 			// Start an asynchronously executing goroutine that
 			// calls Bounce at pseudo-random point in time.
@@ -328,9 +336,11 @@ func TestBounceStress(t *testing.T) {
 			}, &si); err != platform.ErrContextInterrupt {
 				t.Errorf("application partial restore: got %v, wanted %v", err, platform.ErrContextInterrupt)
 			}
-			c.unlock()
+			m.Put(c)
 			randomSleep()
-			c.lock()
+			if c2 := m.Get(); c != c2 {
+				t.Fatalf("machine.Get() changed vCPUs from %p to %p", c, c2)
+			}
 		}
 		return false
 	})
@@ -480,6 +490,63 @@ func TestKernelVDSO(t *testing.T) {
 	})
 }
 
+// Regression test for b/404271139.
+func TestSingleVCPU(t *testing.T) {
+	// Create the machine.
+	deviceFile, err := OpenDevice("")
+	if err != nil {
+		t.Fatalf("error opening device file: %v", err)
+	}
+	k, err := New(deviceFile, Config{
+		MaxVCPUs: 1,
+	})
+	if err != nil {
+		t.Fatalf("error creating KVM instance: %v", err)
+	}
+	defer k.machine.Destroy()
+
+	// Ping-pong the single vCPU between two goroutines. The test passes if
+	// this does not deadlock.
+	stopC := make(chan struct{})
+	var doneWG sync.WaitGroup
+	defer func() {
+		close(stopC)
+		doneWG.Wait()
+	}()
+	var wakeC [2]chan struct{}
+	for i := range wakeC {
+		wakeC[i] = make(chan struct{}, 1)
+	}
+	for i := range wakeC {
+		doneWG.Add(1)
+		go func(i int) {
+			defer doneWG.Done()
+			for {
+				// Multiple ready channels in a select statement are chosen
+				// from randomly, so have a separate non-blocking receive from
+				// stopC first to ensure that it's honored in deterministic
+				// time.
+				select {
+				case <-stopC:
+					return
+				default:
+				}
+				select {
+				case <-stopC:
+					return
+				case <-wakeC[i]:
+					c := k.machine.Get()
+					bluepill(c)
+					wakeC[1-i] <- struct{}{}
+					k.machine.Put(c)
+				}
+			}
+		}(i)
+	}
+	wakeC[0] <- struct{}{}
+	time.Sleep(time.Second)
+}
+
 func BenchmarkApplicationSyscall(b *testing.B) {
 	var (
 		i int // Iteration includes machine.Get() / machine.Put().