introduce atomicbitops 32-bit types

Part of a series of changes that will end with prohibiting use of sync/atomic
(u)int32 functions. See cl/440484071 for more details.

PiperOrigin-RevId: 442673296

pkg/atomicbitops/32b_32bit.go

@@ -0,0 +1,201 @@
+// Copyright 2022 The gVisor Authors.
+//
+// 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.
+
+//go:build arm || mips || mipsle || 386
+// +build arm mips mipsle 386
+
+package atomicbitops
+
+import (
+	"sync/atomic"
+
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+// Note that this file is *identical* to 32b_64bit.go, as go_stateify gets
+// confused about build tags if these are not separated.
+
+// LINT.IfChange
+
+// Int32 is an atomic int32.
+//
+// The default value is zero.
+//
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
+// +stateify savable
+type Int32 struct {
+	_     sync.NoCopy
+	value int32
+}
+
+// FromInt32 returns an Int32 initialized to value v.
+//go:nosplit
+func FromInt32(v int32) Int32 {
+	return Int32{value: v}
+}
+
+// Load is analogous to atomic.LoadInt32.
+//go:nosplit
+func (i *Int32) Load() int32 {
+	return atomic.LoadInt32(&i.value)
+}
+
+// RacyLoad is analogous to reading an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (i *Int32) RacyLoad() int32 {
+	return i.value
+}
+
+// Store is analogous to atomic.StoreInt32.
+//go:nosplit
+func (i *Int32) Store(v int32) {
+	atomic.StoreInt32(&i.value, v)
+}
+
+// RacyStore is analogous to setting an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
+//go:nosplit
+func (i *Int32) RacyStore(v int32) {
+	i.value = v
+}
+
+// Add is analogous to atomic.AddInt32.
+//go:nosplit
+func (i *Int32) Add(v int32) int32 {
+	return atomic.AddInt32(&i.value, v)
+}
+
+// RacyAdd is analogous to adding to an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (i *Int32) RacyAdd(v int32) int32 {
+	i.value += v
+	return i.value
+}
+
+// Swap is analogous to atomic.SwapInt32.
+//go:nosplit
+func (i *Int32) Swap(v int32) int32 {
+	return atomic.SwapInt32(&i.value, v)
+}
+
+// CompareAndSwap is analogous to atomic.CompareAndSwapInt32.
+//go:nosplit
+func (i *Int32) CompareAndSwap(oldVal, newVal int32) bool {
+	return atomic.CompareAndSwapInt32(&i.value, oldVal, newVal)
+}
+
+//go:nosplit
+func (i *Int32) ptr() *int32 {
+	return &i.value
+}
+
+// Uint32 is an atomic uint32.
+//
+// See aligned_unsafe.go in this directory for justification.
+//
+// +stateify savable
+type Uint32 struct {
+	_     sync.NoCopy
+	value uint32
+}
+
+// FromUint32 returns an Uint32 initialized to value v.
+//go:nosplit
+func FromUint32(v uint32) Uint32 {
+	return Uint32{value: v}
+}
+
+// Load is analogous to atomic.LoadUint32.
+//go:nosplit
+func (u *Uint32) Load() uint32 {
+	return atomic.LoadUint32(&u.value)
+}
+
+// RacyLoad is analogous to reading an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyLoad() uint32 {
+	return u.value
+}
+
+// Store is analogous to atomic.StoreUint32.
+//go:nosplit
+func (u *Uint32) Store(v uint32) {
+	atomic.StoreUint32(&u.value, v)
+}
+
+// RacyStore is analogous to setting an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyStore(v uint32) {
+	u.value = v
+}
+
+// Add is analogous to atomic.AddUint32.
+//go:nosplit
+func (u *Uint32) Add(v uint32) uint32 {
+	return atomic.AddUint32(&u.value, v)
+}
+
+// RacyAdd is analogous to adding to an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyAdd(v uint32) uint32 {
+	u.value += v
+	return u.value
+}
+
+// Swap is analogous to atomic.SwapUint32.
+//go:nosplit
+func (u *Uint32) Swap(v uint32) uint32 {
+	return atomic.SwapUint32(&u.value, v)
+}
+
+// CompareAndSwap is analogous to atomic.CompareAndSwapUint32.
+//go:nosplit
+func (u *Uint32) CompareAndSwap(oldVal, newVal uint32) bool {
+	return atomic.CompareAndSwapUint32(&u.value, oldVal, newVal)
+}
+
+//go:nosplit
+func (u *Uint32) ptr() *uint32 {
+	return &u.value
+}
+
+// LINT.ThenChange(32b_64bit.go)

pkg/atomicbitops/32b_64bit.go

@@ -0,0 +1,201 @@
+// Copyright 2022 The gVisor Authors.
+//
+// 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.
+
+//go:build !arm && !mips && !mipsle && !386
+// +build !arm,!mips,!mipsle,!386
+
+package atomicbitops
+
+import (
+	"sync/atomic"
+
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
+// Note that this file is *identical* to 32b_32bit.go, as go_stateify gets
+// confused about build tags if these are not separated.
+
+// LINT.IfChange
+
+// Int32 is an atomic int32.
+//
+// The default value is zero.
+//
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
+// +stateify savable
+type Int32 struct {
+	_     sync.NoCopy
+	value int32
+}
+
+// FromInt32 returns an Int32 initialized to value v.
+//go:nosplit
+func FromInt32(v int32) Int32 {
+	return Int32{value: v}
+}
+
+// Load is analogous to atomic.LoadInt32.
+//go:nosplit
+func (i *Int32) Load() int32 {
+	return atomic.LoadInt32(&i.value)
+}
+
+// RacyLoad is analogous to reading an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (i *Int32) RacyLoad() int32 {
+	return i.value
+}
+
+// Store is analogous to atomic.StoreInt32.
+//go:nosplit
+func (i *Int32) Store(v int32) {
+	atomic.StoreInt32(&i.value, v)
+}
+
+// RacyStore is analogous to setting an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
+//go:nosplit
+func (i *Int32) RacyStore(v int32) {
+	i.value = v
+}
+
+// Add is analogous to atomic.AddInt32.
+//go:nosplit
+func (i *Int32) Add(v int32) int32 {
+	return atomic.AddInt32(&i.value, v)
+}
+
+// RacyAdd is analogous to adding to an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (i *Int32) RacyAdd(v int32) int32 {
+	i.value += v
+	return i.value
+}
+
+// Swap is analogous to atomic.SwapInt32.
+//go:nosplit
+func (i *Int32) Swap(v int32) int32 {
+	return atomic.SwapInt32(&i.value, v)
+}
+
+// CompareAndSwap is analogous to atomic.CompareAndSwapInt32.
+//go:nosplit
+func (i *Int32) CompareAndSwap(oldVal, newVal int32) bool {
+	return atomic.CompareAndSwapInt32(&i.value, oldVal, newVal)
+}
+
+//go:nosplit
+func (i *Int32) ptr() *int32 {
+	return &i.value
+}
+
+// Uint32 is an atomic uint32.
+//
+// See aligned_unsafe.go in this directory for justification.
+//
+// +stateify savable
+type Uint32 struct {
+	_     sync.NoCopy
+	value uint32
+}
+
+// FromUint32 returns an Uint32 initialized to value v.
+//go:nosplit
+func FromUint32(v uint32) Uint32 {
+	return Uint32{value: v}
+}
+
+// Load is analogous to atomic.LoadUint32.
+//go:nosplit
+func (u *Uint32) Load() uint32 {
+	return atomic.LoadUint32(&u.value)
+}
+
+// RacyLoad is analogous to reading an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyLoad() uint32 {
+	return u.value
+}
+
+// Store is analogous to atomic.StoreUint32.
+//go:nosplit
+func (u *Uint32) Store(v uint32) {
+	atomic.StoreUint32(&u.value, v)
+}
+
+// RacyStore is analogous to setting an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyStore(v uint32) {
+	u.value = v
+}
+
+// Add is analogous to atomic.AddUint32.
+//go:nosplit
+func (u *Uint32) Add(v uint32) uint32 {
+	return atomic.AddUint32(&u.value, v)
+}
+
+// RacyAdd is analogous to adding to an atomic value without using
+// synchronization.
+//
+// It may be helpful to document why a racy operation is permitted.
+//
+//go:nosplit
+func (u *Uint32) RacyAdd(v uint32) uint32 {
+	u.value += v
+	return u.value
+}
+
+// Swap is analogous to atomic.SwapUint32.
+//go:nosplit
+func (u *Uint32) Swap(v uint32) uint32 {
+	return atomic.SwapUint32(&u.value, v)
+}
+
+// CompareAndSwap is analogous to atomic.CompareAndSwapUint32.
+//go:nosplit
+func (u *Uint32) CompareAndSwap(oldVal, newVal uint32) bool {
+	return atomic.CompareAndSwapUint32(&u.value, oldVal, newVal)
+}
+
+//go:nosplit
+func (u *Uint32) ptr() *uint32 {
+	return &u.value
+}
+
+// LINT.ThenChange(32b_32bit.go)

pkg/atomicbitops/BUILD

@@ -5,6 +5,8 @@ package(licenses = ["notice"])
 go_library(
     name = "atomicbitops",
     srcs = [
+        "32b_32bit.go",
+        "32b_64bit.go",
         "aligned_32bit_unsafe.go",
         "aligned_64bit.go",
         "atomicbitops.go",

pkg/atomicbitops/aligned_32bit_unsafe.go

@@ -27,6 +27,9 @@ import (
 // Int64 is an atomic int64 that is guaranteed to be 64-bit
 // aligned, even on 32-bit systems.
 //
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
 // Per https://golang.org/pkg/sync/atomic/#pkg-note-BUG:
 //
 // "On ARM, 386, and 32-bit MIPS, it is the caller's responsibility to arrange
@@ -121,6 +124,9 @@ func (i *Int64) CompareAndSwap(oldVal, newVal int64) bool {
 // Uint64 is an atomic uint64 that is guaranteed to be 64-bit
 // aligned, even on 32-bit systems.
 //
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
 // Per https://golang.org/pkg/sync/atomic/#pkg-note-BUG:
 //
 // "On ARM, 386, and 32-bit MIPS, it is the caller's responsibility to arrange

pkg/atomicbitops/aligned_64bit.go

@@ -29,6 +29,9 @@ import (
 //
 // The default value is zero.
 //
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
 // See aligned_32bit_unsafe.go in this directory for justification.
 //
 // +stateify savable
@@ -113,6 +116,9 @@ func (i *Int64) ptr() *int64 {
 // aligned, even on 32-bit systems. On most architectures, it's just a regular
 // uint64.
 //
+// Don't add fields to this struct. It is important that it remain the same
+// size as its builtin analogue.
+//
 // See aligned_unsafe.go in this directory for justification.
 //
 // +stateify savable

pkg/atomicbitops/aligned_test.go

@@ -16,6 +16,7 @@ package atomicbitops
 
 import (
 	"testing"
+	"unsafe"
 )
 
 func TestAtomiciInt64(t *testing.T) {
@@ -33,3 +34,18 @@ func TestAtomicUint64(t *testing.T) {
 	}{}
 	v.v64.Add(1)
 }
+
+func TestSize(t *testing.T) {
+	if size := unsafe.Sizeof(Int32{}); size != 4 {
+		t.Errorf("Int32 should be 4 bytes in size, but is %d bytes", size)
+	}
+	if size := unsafe.Sizeof(Uint32{}); size != 4 {
+		t.Errorf("Uint32 should be 4 bytes in size, but is %d bytes", size)
+	}
+	if size := unsafe.Sizeof(Int64{}); size != 8 {
+		t.Errorf("Int32 should be 8 bytes in size, but is %d bytes", size)
+	}
+	if size := unsafe.Sizeof(Uint64{}); size != 8 {
+		t.Errorf("Int32 should be 8 bytes in size, but is %d bytes", size)
+	}
+}

pkg/atomicbitops/atomicbitops.go

@@ -22,17 +22,33 @@
 package atomicbitops
 
 // AndUint32 atomically applies bitwise AND operation to *addr with val.
-func AndUint32(addr *uint32, val uint32)
+func AndUint32(addr *Uint32, val uint32) {
+	andUint32(&addr.value, val)
+}
+
+func andUint32(addr *uint32, val uint32)
 
 // OrUint32 atomically applies bitwise OR operation to *addr with val.
-func OrUint32(addr *uint32, val uint32)
+func OrUint32(addr *Uint32, val uint32) {
+	orUint32(&addr.value, val)
+}
+
+func orUint32(addr *uint32, val uint32)
 
 // XorUint32 atomically applies bitwise XOR operation to *addr with val.
-func XorUint32(addr *uint32, val uint32)
+func XorUint32(addr *Uint32, val uint32) {
+	xorUint32(&addr.value, val)
+}
+
+func xorUint32(addr *uint32, val uint32)
 
 // CompareAndSwapUint32 is like sync/atomic.CompareAndSwapUint32, but returns
 // the value previously stored at addr.
-func CompareAndSwapUint32(addr *uint32, old, new uint32) uint32
+func CompareAndSwapUint32(addr *Uint32, old, new uint32) uint32 {
+	return compareAndSwapUint32(&addr.value, old, new)
+}
+
+func compareAndSwapUint32(addr *uint32, old, new uint32) uint32
 
 // AndUint64 atomically applies bitwise AND operation to *addr with val.
 func AndUint64(addr *Uint64, val uint64) {

pkg/atomicbitops/atomicbitops_amd64.s

@@ -16,28 +16,28 @@
 
 #include "textflag.h"
 
-TEXT ·AndUint32(SB),NOSPLIT,$0-12
+TEXT ·andUint32(SB),NOSPLIT,$0-12
   MOVQ  addr+0(FP), BX
   MOVL  val+8(FP), AX
   LOCK
   ANDL   AX, 0(BX)
   RET
 
-TEXT ·OrUint32(SB),NOSPLIT,$0-12
+TEXT ·orUint32(SB),NOSPLIT,$0-12
   MOVQ  addr+0(FP), BX
   MOVL  val+8(FP), AX
   LOCK
   ORL   AX, 0(BX)
   RET
 
-TEXT ·XorUint32(SB),NOSPLIT,$0-12
+TEXT ·xorUint32(SB),NOSPLIT,$0-12
   MOVQ  addr+0(FP), BX
   MOVL  val+8(FP), AX
   LOCK
   XORL   AX, 0(BX)
   RET
 
-TEXT ·CompareAndSwapUint32(SB),NOSPLIT,$0-20
+TEXT ·compareAndSwapUint32(SB),NOSPLIT,$0-20
   MOVQ  addr+0(FP), DI
   MOVL  old+8(FP), AX
   MOVL  new+12(FP), DX

pkg/atomicbitops/atomicbitops_arm64.s

@@ -16,7 +16,7 @@
 
 #include "textflag.h"
 
-TEXT ·AndUint32(SB),NOSPLIT,$0-12
+TEXT ·andUint32(SB),NOSPLIT,$0-12
   MOVD    ptr+0(FP), R0
   MOVW    val+8(FP), R1
 again:
@@ -26,7 +26,7 @@ again:
   CBNZ    R3, again
   RET
 
-TEXT ·OrUint32(SB),NOSPLIT,$0-12
+TEXT ·orUint32(SB),NOSPLIT,$0-12
   MOVD    ptr+0(FP), R0
   MOVW    val+8(FP), R1
 again:
@@ -36,7 +36,7 @@ again:
   CBNZ    R3, again
   RET
 
-TEXT ·XorUint32(SB),NOSPLIT,$0-12
+TEXT ·xorUint32(SB),NOSPLIT,$0-12
   MOVD    ptr+0(FP), R0
   MOVW    val+8(FP), R1
 again:
@@ -46,7 +46,7 @@ again:
   CBNZ    R3, again
   RET
 
-TEXT ·CompareAndSwapUint32(SB),NOSPLIT,$0-20
+TEXT ·compareAndSwapUint32(SB),NOSPLIT,$0-20
   MOVD addr+0(FP), R0
   MOVW old+8(FP), R1
   MOVW new+12(FP), R2

pkg/atomicbitops/atomicbitops_noasm.go

@@ -17,51 +17,49 @@
 
 package atomicbitops
 
-import (
-	"sync/atomic"
-)
+import "sync/atomic"
 
 //go:nosplit
-func AndUint32(addr *uint32, val uint32) {
+func AndUint32(addr *Uint32, val uint32) {
 	for {
-		o := atomic.LoadUint32(addr)
+		o := addr.Load()
 		n := o & val
-		if atomic.CompareAndSwapUint32(addr, o, n) {
+		if atomic.CompareAndSwapUint32(&addr.value, o, n) {
 			break
 		}
 	}
 }
 
 //go:nosplit
-func OrUint32(addr *uint32, val uint32) {
+func OrUint32(addr *Uint32, val uint32) {
 	for {
-		o := atomic.LoadUint32(addr)
+		o := addr.Load()
 		n := o | val
-		if atomic.CompareAndSwapUint32(addr, o, n) {
+		if atomic.CompareAndSwapUint32(&addr.value, o, n) {
 			break
 		}
 	}
 }
 
 //go:nosplit
-func XorUint32(addr *uint32, val uint32) {
+func XorUint32(addr *Uint32, val uint32) {
 	for {
-		o := atomic.LoadUint32(addr)
+		o := addr.Load()
 		n := o ^ val
-		if atomic.CompareAndSwapUint32(addr, o, n) {
+		if atomic.CompareAndSwapUint32(&addr.value, o, n) {
 			break
 		}
 	}
 }
 
 //go:nosplit
-func CompareAndSwapUint32(addr *uint32, old, new uint32) (prev uint32) {
+func CompareAndSwapUint32(addr *Uint32, old, new uint32) (prev uint32) {
 	for {
-		prev = atomic.LoadUint32(addr)
+		prev = addr.Load()
 		if prev != old {
 			return
 		}
-		if atomic.CompareAndSwapUint32(addr, old, new) {
+		if atomic.CompareAndSwapUint32(&addr.value, old, new) {
 			return
 		}
 	}

pkg/atomicbitops/atomicbitops_test.go

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// +checkalignedignore
 package atomicbitops
 
 import (
@@ -23,20 +24,20 @@ import (
 
 const iterations = 100
 
-func detectRaces32(val, target uint32, fn func(*uint32, uint32)) bool {
+func detectRaces32(val, target uint32, fn func(*Uint32, uint32)) bool {
 	runtime.GOMAXPROCS(100)
 	for n := 0; n < iterations; n++ {
-		x := val
+		x := FromUint32(val)
 		var wg sync.WaitGroup
 		for i := uint32(0); i < 32; i++ {
 			wg.Add(1)
-			go func(a *uint32, i uint32) {
+			go func(a *Uint32, i uint32) {
 				defer wg.Done()
 				fn(a, uint32(1<<i))
 			}(&x, i)
 		}
 		wg.Wait()
-		if x != target {
+		if x != FromUint32(target) {
 			return true
 		}
 	}
@@ -137,12 +138,12 @@ func TestCompareAndSwapUint32(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		val := test.prev
+		val := FromUint32(test.prev)
 		prev := CompareAndSwapUint32(&val, test.old, test.new)
 		if got, want := prev, test.prev; got != want {
 			t.Errorf("%s: incorrect returned previous value: got %d, expected %d", test.name, got, want)
 		}
-		if got, want := val, test.next; got != want {
+		if got, want := val.Load(), test.next; got != want {
 			t.Errorf("%s: incorrect value stored in val: got %d, expected %d", test.name, got, want)
 		}
 	}

pkg/sentry/platform/kvm/bluepill_fault.go

@@ -87,10 +87,10 @@ func handleBluepillFault(m *machine, physical uintptr, phyRegions []physicalRegi
 	//
 	// First, we need to acquire the exclusive right to set a slot.  See
 	// machine.nextSlot for information about the protocol.
-	slot := atomic.SwapUint32(&m.nextSlot, ^uint32(0))
+	slot := m.nextSlot.Swap(^uint32(0))
 	for slot == ^uint32(0) {
 		yield() // Race with another call.
-		slot = atomic.SwapUint32(&m.nextSlot, ^uint32(0))
+		slot = m.nextSlot.Swap(^uint32(0))
 	}
 	flags := _KVM_MEM_FLAGS_NONE
 	if pr.readOnly {
@@ -104,12 +104,12 @@ func handleBluepillFault(m *machine, physical uintptr, phyRegions []physicalRegi
 		atomic.StoreUintptr(&m.usedSlots[slot], physicalStart)
 		// Successfully added region; we can increment nextSlot and
 		// allow another set to proceed here.
-		atomic.StoreUint32(&m.nextSlot, slot+1)
+		m.nextSlot.Store(slot + 1)
 		return virtualStart + (physical - physicalStart), true
 	}
 
 	// Release our slot (still available).
-	atomic.StoreUint32(&m.nextSlot, slot)
+	m.nextSlot.Store(slot)
 
 	switch errno {
 	case unix.EEXIST:

pkg/sentry/platform/kvm/bluepill_unsafe.go

@@ -22,7 +22,6 @@
 package kvm
 
 import (
-	"sync/atomic"
 	"unsafe"
 
 	"golang.org/x/sys/unix"
@@ -71,8 +70,8 @@ func bluepillGuestExit(c *vCPU, context unsafe.Pointer) {
 	bluepillArchExit(c, bluepillArchContext(context))
 
 	// Return to the vCPUReady state; notify any waiters.
-	user := atomic.LoadUint32(&c.state) & vCPUUser
-	switch atomic.SwapUint32(&c.state, user) {
+	user := c.state.Load() & vCPUUser
+	switch c.state.Swap(user) {
 	case user | vCPUGuest: // Expected case.
 	case user | vCPUGuest | vCPUWaiter:
 		c.notify()
@@ -102,7 +101,7 @@ func bluepillHandler(context unsafe.Pointer) {
 	c := bluepillArchEnter(bluepillArchContext(context))
 
 	// Mark this as guest mode.
-	switch atomic.SwapUint32(&c.state, vCPUGuest|vCPUUser) {
+	switch c.state.Swap(vCPUGuest | vCPUUser) {
 	case vCPUUser: // Expected case.
 	case vCPUUser | vCPUWaiter:
 		c.notify()

pkg/sentry/platform/kvm/kvm_test.go

@@ -17,7 +17,6 @@ package kvm
 import (
 	"math/rand"
 	"reflect"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -88,7 +87,7 @@ func bluepillTest(t testHarness, fn func(*vCPU)) {
 func TestKernelSyscall(t *testing.T) {
 	bluepillTest(t, func(c *vCPU) {
 		redpill() // Leave guest mode.
-		if got := atomic.LoadUint32(&c.state); got != vCPUUser {
+		if got := c.state.Load(); got != vCPUUser {
 			t.Errorf("vCPU not in ready state: got %v", got)
 		}
 	})
@@ -106,7 +105,7 @@ func TestKernelFault(t *testing.T) {
 	hostFault() // Ensure recovery works.
 	bluepillTest(t, func(c *vCPU) {
 		hostFault()
-		if got := atomic.LoadUint32(&c.state); got != vCPUUser {
+		if got := c.state.Load(); got != vCPUUser {
 			t.Errorf("vCPU not in ready state: got %v", got)
 		}
 	})

pkg/sentry/platform/kvm/machine.go

@@ -44,9 +44,9 @@ type machine struct {
 
 	// nextSlot is the next slot for setMemoryRegion.
 	//
-	// This must be accessed atomically. If nextSlot is ^uint32(0), then
-	// slots are currently being updated, and the caller should retry.
-	nextSlot uint32
+	// If nextSlot is ^uint32(0), then slots are currently being updated, and the
+	// caller should retry.
+	nextSlot atomicbitops.Uint32
 
 	// upperSharedPageTables tracks the read-only shared upper of all the pagetables.
 	upperSharedPageTables *pagetables.PageTables
@@ -130,7 +130,7 @@ type vCPU struct {
 	// state is the vCPU state.
 	//
 	// This is a bitmask of the three fields (vCPU*) described above.
-	state uint32
+	state atomicbitops.Uint32
 
 	// runData for this vCPU.
 	runData *runData
@@ -314,7 +314,7 @@ func newMachine(vm int) (*machine, error) {
 //
 //go:nosplit
 func (m *machine) hasSlot(physical uintptr) bool {
-	slotLen := int(atomic.LoadUint32(&m.nextSlot))
+	slotLen := int(m.nextSlot.Load())
 	// When slots are being updated, nextSlot is ^uint32(0). As this situation
 	// is less likely happen, we just set the slotLen to m.maxSlots, and scan
 	// the whole usedSlots array.
@@ -438,7 +438,7 @@ func (m *machine) Get() *vCPU {
 	for {
 		// Scan for an available vCPU.
 		for origTID, c := range m.vCPUsByTID {
-			if atomic.CompareAndSwapUint32(&c.state, vCPUReady, vCPUUser) {
+			if c.state.CompareAndSwap(vCPUReady, vCPUUser) {
 				delete(m.vCPUsByTID, origTID)
 				m.vCPUsByTID[tid] = c
 				m.mu.Unlock()
@@ -460,7 +460,7 @@ func (m *machine) Get() *vCPU {
 
 		// Scan for something not in user mode.
 		for origTID, c := range m.vCPUsByTID {
-			if !atomic.CompareAndSwapUint32(&c.state, vCPUGuest, vCPUGuest|vCPUWaiter) {
+			if !c.state.CompareAndSwap(vCPUGuest, vCPUGuest|vCPUWaiter) {
 				continue
 			}
 
@@ -471,7 +471,7 @@ func (m *machine) Get() *vCPU {
 			// just the vCPUReady state.
 			for {
 				c.waitUntilNot(vCPUGuest | vCPUWaiter)
-				if atomic.CompareAndSwapUint32(&c.state, vCPUReady, vCPUUser) {
+				if c.state.CompareAndSwap(vCPUReady, vCPUUser) {
 					break
 				}
 			}
@@ -591,7 +591,7 @@ func (c *vCPU) bounce(forceGuestExit bool) {
 	origGuestExits := c.guestExits.Load()
 	origUserExits := c.userExits.Load()
 	for {
-		switch state := atomic.LoadUint32(&c.state); state {
+		switch state := c.state.Load(); state {
 		case vCPUReady, vCPUWaiter:
 			// There is nothing to be done, we're already in the
 			// kernel pre-acquisition. The Bounce criteria have
@@ -602,7 +602,7 @@ func (c *vCPU) bounce(forceGuestExit bool) {
 			// transition. When the transition takes place, then we
 			// can inject an interrupt to ensure a return to host
 			// mode.
-			atomic.CompareAndSwapUint32(&c.state, state, state|vCPUWaiter)
+			c.state.CompareAndSwap(state, state|vCPUWaiter)
 		case vCPUUser | vCPUWaiter:
 			// Wait for the transition to guest mode. This should
 			// come from the bluepill handler.
@@ -615,7 +615,7 @@ func (c *vCPU) bounce(forceGuestExit bool) {
 			}
 			// The vCPU is in user or kernel mode. Attempt to
 			// register a notification on change.
-			if !atomic.CompareAndSwapUint32(&c.state, state, state|vCPUWaiter) {
+			if !c.state.CompareAndSwap(state, state|vCPUWaiter) {
 				break // Retry.
 			}
 			for {
@@ -707,7 +707,7 @@ const machinePoolSize = 16
 // machinePool is enumerated from the seccompMmapHandler signal handler
 var (
 	machinePool          [machinePoolSize]machineAtomicPtr
-	machinePoolLen       uint32
+	machinePoolLen       atomicbitops.Uint32
 	machinePoolMu        sync.Mutex
 	seccompMmapRulesOnce gosync.Once
 )
@@ -752,7 +752,7 @@ func seccompMmapRules(m *machine) {
 	})
 
 	machinePoolMu.Lock()
-	n := atomic.LoadUint32(&machinePoolLen)
+	n := machinePoolLen.Load()
 	i := uint32(0)
 	for ; i < n; i++ {
 		if machinePool[i].Load() == nil {
@@ -764,7 +764,7 @@ func seccompMmapRules(m *machine) {
 			machinePoolMu.Unlock()
 			panic("machinePool is full")
 		}
-		atomic.AddUint32(&machinePoolLen, 1)
+		machinePoolLen.Add(1)
 	}
 	machinePool[i].Store(m)
 	m.machinePoolIndex = i

pkg/sentry/platform/kvm/machine_unsafe.go

@@ -208,7 +208,7 @@ func seccompMmapHandler(context unsafe.Pointer) {
 	}
 
 	seccompMmapHandlerCnt.Add(1)
-	for i := uint32(0); i < atomic.LoadUint32(&machinePoolLen); i++ {
+	for i := uint32(0); i < machinePoolLen.Load(); i++ {
 		m := machinePool[i].Load()
 		if m == nil {
 			continue

pkg/usermem/bytes_io_unsafe.go

@@ -15,7 +15,6 @@
 package usermem
 
 import (
-	"sync/atomic"
 	"unsafe"
 
 	"gvisor.dev/gvisor/pkg/atomicbitops"
@@ -28,7 +27,8 @@ func (b *BytesIO) SwapUint32(ctx context.Context, addr hostarch.Addr, new uint32
 	if _, rngErr := b.rangeCheck(addr, 4); rngErr != nil {
 		return 0, rngErr
 	}
-	return atomic.SwapUint32((*uint32)(unsafe.Pointer(&b.Bytes[int(addr)])), new), nil
+	return (*atomicbitops.Uint32)(unsafe.Pointer(&b.Bytes[int(addr)])).Swap(new), nil
+
 }
 
 // CompareAndSwapUint32 implements IO.CompareAndSwapUint32.
@@ -36,7 +36,7 @@ func (b *BytesIO) CompareAndSwapUint32(ctx context.Context, addr hostarch.Addr,
 	if _, rngErr := b.rangeCheck(addr, 4); rngErr != nil {
 		return 0, rngErr
 	}
-	return atomicbitops.CompareAndSwapUint32((*uint32)(unsafe.Pointer(&b.Bytes[int(addr)])), old, new), nil
+	return atomicbitops.CompareAndSwapUint32((*atomicbitops.Uint32)(unsafe.Pointer(&b.Bytes[int(addr)])), old, new), nil
 }
 
 // LoadUint32 implements IO.LoadUint32.
@@ -44,5 +44,5 @@ func (b *BytesIO) LoadUint32(ctx context.Context, addr hostarch.Addr, opts IOOpt
 	if _, err := b.rangeCheck(addr, 4); err != nil {
 		return 0, err
 	}
-	return atomic.LoadUint32((*uint32)(unsafe.Pointer(&b.Bytes[int(addr)]))), nil
+	return (*atomicbitops.Uint32)(unsafe.Pointer(&b.Bytes[int(addr)])).Load(), nil
 }