Simplify KVM invalidation logic.

PiperOrigin-RevId: 196780209
Change-Id: I89f39eec914ce54a7c6c4f28e1b6d5ff5a7dd38d

7 files changed, 106 insertions, 22 deletions

diff --git a/pkg/sentry/platform/kvm/BUILD b/pkg/sentry/platform/kvm/BUILD
index adc43c21b..004938080 100644
--- a/pkg/sentry/platform/kvm/BUILD
+++ b/pkg/sentry/platform/kvm/BUILD
@@ -27,6 +27,7 @@ go_library(
     name = "kvm",
     srcs = [
         "address_space.go",
+        "address_space_unsafe.go",
         "bluepill.go",
         "bluepill_amd64.go",
         "bluepill_amd64.s",
diff --git a/pkg/sentry/platform/kvm/address_space.go b/pkg/sentry/platform/kvm/address_space.go
index 3d57ae0cb..e81cc0caf 100644
--- a/pkg/sentry/platform/kvm/address_space.go
+++ b/pkg/sentry/platform/kvm/address_space.go
@@ -16,8 +16,6 @@ package kvm
 
 import (
 	"reflect"
-	"sync"
-	"sync/atomic"
 
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform/filemem"
@@ -40,10 +38,10 @@ type addressSpace struct {
 
 	// dirtySet is the set of dirty vCPUs.
 	//
-	// The key is the vCPU, the value is a shared uint32 pointer that
-	// indicates whether or not the context is clean. A zero here indicates
-	// that the context should be cleaned prior to re-entry.
-	dirtySet sync.Map
+	// These are actually vCPU pointers that are stored iff the vCPU is
+	// dirty. If the vCPU is not dirty and requires invalidation, then a
+	// nil value is stored here instead.
+	dirtySet dirtySet
 
 	// files contains files mapped in the host address space.
 	//
@@ -53,22 +51,22 @@ type addressSpace struct {
 
 // Invalidate interrupts all dirty contexts.
 func (as *addressSpace) Invalidate() {
-	as.dirtySet.Range(func(key, value interface{}) bool {
-		c := key.(*vCPU)
-		v := value.(*uint32)
-		atomic.StoreUint32(v, 0) // Invalidation required.
-		c.BounceToKernel()       // Force a kernel transition.
-		return true              // Keep iterating.
-	})
+	for i := 0; i < as.dirtySet.size(); i++ {
+		if c := as.dirtySet.swap(i, nil); c != nil && c.active.get() == as {
+			c.BounceToKernel() // Force a kernel transition.
+		}
+	}
 }
 
 // Touch adds the given vCPU to the dirty list.
-func (as *addressSpace) Touch(c *vCPU) *uint32 {
-	value, ok := as.dirtySet.Load(c)
-	if !ok {
-		value, _ = as.dirtySet.LoadOrStore(c, new(uint32))
+//
+// The return value indicates whether a flush is required.
+func (as *addressSpace) Touch(c *vCPU) bool {
+	if old := as.dirtySet.swap(c.id, c); old == nil {
+		return true // Flush is required.
 	}
-	return value.(*uint32)
+	// Already dirty: no flush required.
+	return false
 }
 
 func (as *addressSpace) mapHost(addr usermem.Addr, m hostMapEntry, at usermem.AccessType) (inv bool) {
diff --git a/pkg/sentry/platform/kvm/address_space_unsafe.go b/pkg/sentry/platform/kvm/address_space_unsafe.go
new file mode 100644
index 000000000..b6c31ce10
--- /dev/null
+++ b/pkg/sentry/platform/kvm/address_space_unsafe.go
@@ -0,0 +1,44 @@
+// Copyright 2018 Google Inc.
+//
+// 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 kvm
+
+import (
+	"sync/atomic"
+	"unsafe"
+)
+
+// dirtySet tracks vCPUs for invalidation.
+type dirtySet struct {
+	vCPUs []unsafe.Pointer
+}
+
+// makeDirtySet makes a new dirtySet.
+func makeDirtySet(size int) dirtySet {
+	return dirtySet{
+		vCPUs: make([]unsafe.Pointer, size),
+	}
+}
+
+// size is the size of the set.
+func (ds *dirtySet) size() int {
+	return len(ds.vCPUs)
+}
+
+// swap sets the given index and returns the previous value.
+//
+// The index is typically the id for a non-nil vCPU.
+func (ds *dirtySet) swap(index int, c *vCPU) *vCPU {
+	return (*vCPU)(atomic.SwapPointer(&ds.vCPUs[index], unsafe.Pointer(c)))
+}
diff --git a/pkg/sentry/platform/kvm/context.go b/pkg/sentry/platform/kvm/context.go
index fd04a2c47..c9bfbc136 100644
--- a/pkg/sentry/platform/kvm/context.go
+++ b/pkg/sentry/platform/kvm/context.go
@@ -15,8 +15,6 @@
 package kvm
 
 import (
-	"sync/atomic"
-
 	"gvisor.googlesource.com/gvisor/pkg/sentry/arch"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform/interrupt"
@@ -54,10 +52,18 @@ func (c *context) Switch(as platform.AddressSpace, ac arch.Context, _ int32) (*a
 		return nil, usermem.NoAccess, platform.ErrContextInterrupt
 	}
 
+	// Set the active address space.
+	//
+	// This must be done prior to the call to Touch below. If the address
+	// space is invalidated between this line and the call below, we will
+	// flag on entry anyways. When the active address space below is
+	// cleared, it indicates that we don't need an explicit interrupt and
+	// that the flush can occur naturally on the next user entry.
+	cpu.active.set(localAS)
+
 	// Mark the address space as dirty.
 	flags := ring0.Flags(0)
-	dirty := localAS.Touch(cpu)
-	if v := atomic.SwapUint32(dirty, 1); v == 0 {
+	if localAS.Touch(cpu) {
 		flags |= ring0.FlagFlush
 	}
 	if ac.FullRestore() {
@@ -67,6 +73,9 @@ func (c *context) Switch(as platform.AddressSpace, ac arch.Context, _ int32) (*a
 	// Take the blue pill.
 	si, at, err := cpu.SwitchToUser(regs, fp, localAS.pageTables, flags)
 
+	// Clear the address space.
+	cpu.active.set(nil)
+
 	// Release resources.
 	c.machine.Put(cpu)
 
diff --git a/pkg/sentry/platform/kvm/kvm.go b/pkg/sentry/platform/kvm/kvm.go
index 31928c9f0..15a241f01 100644
--- a/pkg/sentry/platform/kvm/kvm.go
+++ b/pkg/sentry/platform/kvm/kvm.go
@@ -133,6 +133,7 @@ func (k *KVM) NewAddressSpace(_ interface{}) (platform.AddressSpace, <-chan stru
 		filemem:    k.FileMem,
 		machine:    k.machine,
 		pageTables: pageTables,
+		dirtySet:   makeDirtySet(len(k.machine.vCPUs)),
 	}, nil, nil
 }
 
diff --git a/pkg/sentry/platform/kvm/machine.go b/pkg/sentry/platform/kvm/machine.go
index 7a962e316..3ee21fe21 100644
--- a/pkg/sentry/platform/kvm/machine.go
+++ b/pkg/sentry/platform/kvm/machine.go
@@ -80,6 +80,9 @@ type vCPU struct {
 	// by the bluepill code (see bluepill_amd64.s).
 	ring0.CPU
 
+	// id is the vCPU id.
+	id int
+
 	// fd is the vCPU fd.
 	fd int
 
@@ -102,6 +105,10 @@ type vCPU struct {
 
 	// machine associated with this vCPU.
 	machine *machine
+
+	// active is the current addressSpace: this is set and read atomically,
+	// it is used to elide unnecessary interrupts due to invalidations.
+	active atomicAddressSpace
 }
 
 // newMachine returns a new VM context.
@@ -140,6 +147,7 @@ func newMachine(vm int, vCPUs int) (*machine, error) {
 			return nil, fmt.Errorf("error creating VCPU: %v", errno)
 		}
 		c := &vCPU{
+			id:      id,
 			fd:      int(fd),
 			machine: m,
 		}
diff --git a/pkg/sentry/platform/kvm/machine_unsafe.go b/pkg/sentry/platform/kvm/machine_unsafe.go
index 9f7fcd135..516098a2b 100644
--- a/pkg/sentry/platform/kvm/machine_unsafe.go
+++ b/pkg/sentry/platform/kvm/machine_unsafe.go
@@ -16,6 +16,7 @@ package kvm
 
 import (
 	"fmt"
+	"sync/atomic"
 	"syscall"
 	"unsafe"
 
@@ -68,6 +69,28 @@ func unmapRunData(r *runData) error {
 	return nil
 }
 
+// atomicAddressSpace is an atomic address space pointer.
+type atomicAddressSpace struct {
+	pointer unsafe.Pointer
+}
+
+// set sets the address space value.
+//
+//go:nosplit
+func (a *atomicAddressSpace) set(as *addressSpace) {
+	atomic.StorePointer(&a.pointer, unsafe.Pointer(as))
+}
+
+// get gets the address space value.
+//
+// Note that this should be considered best-effort, and may have changed by the
+// time this function returns.
+//
+//go:nosplit
+func (a *atomicAddressSpace) get() *addressSpace {
+	return (*addressSpace)(atomic.LoadPointer(&a.pointer))
+}
+
 // notify notifies that the vCPU has transitioned modes.
 //
 // This may be called by a signal handler and therefore throws on error.