Use atomic.Uint64. (tetratelabs#1620)

Signed-off-by: Nuno Cruces <[email protected]>

experimental/wazerotest/wazerotest.go

@@ -24,7 +24,6 @@ const (
 // WebAssembly module.
 type Module struct {
 	internalapi.WazeroOnlyType
-	exitStatus uint64
 
 	// The module name that will be returned by calling the Name method.
 	ModuleName string
@@ -40,6 +39,8 @@ type Module struct {
 	// "memory".
 	ExportMemory *Memory
 
+	exitStatus atomic.Uint64
+
 	once                        sync.Once
 	exportedFunctions           map[string]api.Function
 	exportedFunctionDefinitions map[string]api.FunctionDefinition
@@ -111,7 +112,7 @@ func (m *Module) Close(ctx context.Context) error {
 
 // CloseWithExitCode implements the same method as documented on api.Closer.
 func (m *Module) CloseWithExitCode(ctx context.Context, exitCode uint32) error {
-	atomic.CompareAndSwapUint64(&m.exitStatus, 0, exitStatusMarker|uint64(exitCode))
+	m.exitStatus.CompareAndSwap(0, exitStatusMarker|uint64(exitCode))
 	return nil
 }
 
@@ -144,7 +145,7 @@ func (m *Module) Function(i int) api.Function {
 }
 
 func (m *Module) ExitStatus() (exitCode uint32, exited bool) {
-	exitStatus := atomic.LoadUint64(&m.exitStatus)
+	exitStatus := m.exitStatus.Load()
 	return uint32(exitStatus), exitStatus != 0
 }
 

internal/engine/compiler/engine.go

@@ -376,13 +376,13 @@ const (
 	functionSize                     = 40
 
 	// Offsets for wasm.ModuleInstance.
-	moduleInstanceGlobalsOffset          = 32
-	moduleInstanceMemoryOffset           = 56
-	moduleInstanceTablesOffset           = 64
-	moduleInstanceEngineOffset           = 88
-	moduleInstanceTypeIDsOffset          = 104
-	moduleInstanceDataInstancesOffset    = 128
-	moduleInstanceElementInstancesOffset = 152
+	moduleInstanceGlobalsOffset          = 24
+	moduleInstanceMemoryOffset           = 48
+	moduleInstanceTablesOffset           = 56
+	moduleInstanceEngineOffset           = 80
+	moduleInstanceTypeIDsOffset          = 96
+	moduleInstanceDataInstancesOffset    = 120
+	moduleInstanceElementInstancesOffset = 144
 
 	// Offsets for wasm.TableInstance.
 	tableInstanceTableOffset    = 0
@@ -966,13 +966,13 @@ func newEngine(enabledFeatures api.CoreFeatures, fileCache filecache.Cache) *eng
 //
 // By declaring these values as `var`, slices created via `make([]..., var)`
 // will never be allocated on stack [1]. This means accessing these slices via
-// raw pointers is safe: As of version 1.18, Go's garbage collector never relocates
+// raw pointers is safe: As of version 1.21, Go's garbage collector never relocates
 // heap-allocated objects (aka no compaction of memory [2]).
 //
 // On Go upgrades, re-validate heap-allocation via `go build -gcflags='-m' ./internal/engine/compiler/...`.
 //
-//	[1] https://github.com/golang/go/blob/68ecdc2c70544c303aa923139a5f16caf107d955/src/cmd/compile/internal/escape/utils.go#L206-L208
-//	[2] https://github.com/golang/go/blob/68ecdc2c70544c303aa923139a5f16caf107d955/src/runtime/mgc.go#L9
+//	[1] https://github.com/golang/go/blob/c19c4c566c63818dfd059b352e52c4710eecf14d/src/cmd/compile/internal/escape/utils.go#L213-L215
+//	[2] https://github.com/golang/go/blob/c19c4c566c63818dfd059b352e52c4710eecf14d/src/runtime/mgc.go#L9
 //	[3] https://mayurwadekar2.medium.com/escape-analysis-in-golang-ee40a1c064c1
 //	[4] https://medium.com/@yulang.chu/go-stack-or-heap-2-slices-which-keep-in-stack-have-limitation-of-size-b3f3adfd6190
 var initialStackSize uint64 = 512

internal/wasm/module_instance.go

@@ -4,15 +4,14 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync/atomic"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/sys"
 )
 
 // FailIfClosed returns a sys.ExitError if CloseWithExitCode was called.
 func (m *ModuleInstance) FailIfClosed() (err error) {
-	if closed := atomic.LoadUint64(&m.Closed); closed != 0 {
+	if closed := m.Closed.Load(); closed != 0 {
 		switch closed & exitCodeFlagMask {
 		case exitCodeFlagResourceClosed:
 		case exitCodeFlagResourceNotClosed:
@@ -108,7 +107,7 @@ func (m *ModuleInstance) CloseWithExitCode(ctx context.Context, exitCode uint32)
 
 // IsClosed implements the same method as documented on api.Module.
 func (m *ModuleInstance) IsClosed() bool {
-	return atomic.LoadUint64(&m.Closed) != 0
+	return m.Closed.Load() != 0
 }
 
 func (m *ModuleInstance) closeWithExitCodeWithoutClosingResource(exitCode uint32) (err error) {
@@ -140,15 +139,14 @@ const (
 
 func (m *ModuleInstance) setExitCode(exitCode uint32, flag exitCodeFlag) bool {
 	closed := flag | uint64(exitCode)<<32 // Store exitCode as high-order bits.
-	return atomic.CompareAndSwapUint64(&m.Closed, 0, closed)
+	return m.Closed.CompareAndSwap(0, closed)
 }
 
 // ensureResourcesClosed ensures that resources assigned to ModuleInstance is released.
 // Only one call will happen per module, due to external atomic guards on Closed.
 func (m *ModuleInstance) ensureResourcesClosed(ctx context.Context) (err error) {
 	if closeNotifier := m.CloseNotifier; closeNotifier != nil { // experimental
-		closed := atomic.LoadUint64(&m.Closed)
-		closeNotifier.CloseNotify(ctx, uint32(closed>>32))
+		closeNotifier.CloseNotify(ctx, uint32(m.Closed.Load()>>32))
 		m.CloseNotifier = nil
 	}
 

internal/wasm/module_instance_test.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"sync"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -97,7 +96,7 @@ func TestModuleInstance_Close(t *testing.T) {
 				// Closing should not err.
 				require.NoError(t, tc.closer(ctx, m))
 
-				require.Equal(t, tc.expectedClosed, m.Closed)
+				require.Equal(t, tc.expectedClosed, m.Closed.Load())
 
 				// Outside callers should be able to know it was closed.
 				require.True(t, m.IsClosed())
@@ -205,7 +204,7 @@ func TestModuleInstance_CallDynamic(t *testing.T) {
 				// Closing should not err.
 				require.NoError(t, tc.closer(ctx, m))
 
-				require.Equal(t, tc.expectedClosed, m.Closed)
+				require.Equal(t, tc.expectedClosed, m.Closed.Load())
 
 				// Verify our intended side-effect
 				require.Nil(t, s.Module(moduleName))
@@ -268,7 +267,7 @@ func TestModuleInstance_CallDynamic(t *testing.T) {
 func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	s := newStore()
 	t.Run("timeout", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
+		cc := &ModuleInstance{ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		const duration = time.Second
 		ctx, cancel := context.WithTimeout(context.Background(), duration)
 		defer cancel()
@@ -277,7 +276,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		defer done()
 
 		// Resource shouldn't be released at this point.
-		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), atomic.LoadUint64(&cc.Closed)&exitCodeFlagMask)
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed.Load()&exitCodeFlagMask)
 		require.NotNil(t, cc.Sys)
 
 		err := cc.FailIfClosed()
@@ -288,7 +287,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 
 	t.Run("cancel", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
+		cc := &ModuleInstance{ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		ctx, cancel := context.WithCancel(context.Background())
 		done := cc.CloseModuleOnCanceledOrTimeout(context.WithValue(ctx, struct{}{}, 1)) // Wrapping arbitrary context.
 		cancel()
@@ -299,7 +298,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		time.Sleep(time.Second)
 
 		// Resource shouldn't be released at this point.
-		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), atomic.LoadUint64(&cc.Closed)&exitCodeFlagMask)
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed.Load()&exitCodeFlagMask)
 		require.NotNil(t, cc.Sys)
 
 		err := cc.FailIfClosed()
@@ -310,7 +309,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 
 	t.Run("timeout over cancel", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
+		cc := &ModuleInstance{ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		const duration = time.Second
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
@@ -322,7 +321,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		defer done()
 
 		// Resource shouldn't be released at this point.
-		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), atomic.LoadUint64(&cc.Closed)&exitCodeFlagMask)
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed.Load()&exitCodeFlagMask)
 		require.NotNil(t, cc.Sys)
 
 		err := cc.FailIfClosed()
@@ -333,7 +332,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 
 	t.Run("cancel over timeout", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
+		cc := &ModuleInstance{ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		ctx, cancel := context.WithCancel(context.Background())
 		// Wrap the timeout context by cancel context.
 		var timeoutDone context.CancelFunc
@@ -347,7 +346,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		time.Sleep(time.Second)
 
 		// Resource shouldn't be released at this point.
-		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), atomic.LoadUint64(&cc.Closed)&exitCodeFlagMask)
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed.Load()&exitCodeFlagMask)
 		require.NotNil(t, cc.Sys)
 
 		err := cc.FailIfClosed()
@@ -358,7 +357,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 
 	t.Run("cancel works", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{ModuleName: "test", s: s}
 		cancelChan := make(chan struct{})
 		var wg sync.WaitGroup
 		wg.Add(1)
@@ -373,7 +372,7 @@ func TestModuleInstance_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 
 	t.Run("no close on all resources canceled", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{ModuleName: "test", s: s}
 		cancelChan := make(chan struct{})
 		close(cancelChan)
 		ctx, cancel := context.WithCancel(context.Background())
@@ -390,7 +389,7 @@ func TestModuleInstance_CloseWithCtxErr(t *testing.T) {
 	s := newStore()
 
 	t.Run("context canceled", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{ModuleName: "test", s: s}
 		ctx, cancel := context.WithCancel(context.Background())
 		cancel()
 
@@ -401,7 +400,7 @@ func TestModuleInstance_CloseWithCtxErr(t *testing.T) {
 	})
 
 	t.Run("context timeout", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{ModuleName: "test", s: s}
 		duration := time.Second
 		ctx, cancel := context.WithTimeout(context.Background(), duration)
 		defer cancel()
@@ -415,7 +414,7 @@ func TestModuleInstance_CloseWithCtxErr(t *testing.T) {
 	})
 
 	t.Run("no error", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{ModuleName: "test", s: s}
 
 		cc.CloseWithCtxErr(context.Background())
 

internal/wasm/store.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"sync"
+	"sync/atomic"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/close"
@@ -71,17 +72,6 @@ type (
 	ModuleInstance struct {
 		internalapi.WazeroOnlyType
 
-		// Closed is used both to guard moduleEngine.CloseWithExitCode and to store the exit code.
-		//
-		// The update value is closedType + exitCode << 32. This ensures an exit code of zero isn't mistaken for never closed.
-		//
-		// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
-		// See /RATIONALE.md
-		//
-		// TODO: Retype this to atomic.Unit64 when Go 1.18 is no longer supported. Until then, keep Closed at the top of
-		// this struct. See PR #1299 for an implementation and discussion.
-		Closed uint64
-
 		ModuleName     string
 		Exports        map[string]*Export
 		Globals        []*GlobalInstance
@@ -116,6 +106,14 @@ type (
 		//	  security implications.
 		Sys *internalsys.Context
 
+		// Closed is used both to guard moduleEngine.CloseWithExitCode and to store the exit code.
+		//
+		// The update value is closedType + exitCode << 32. This ensures an exit code of zero isn't mistaken for never closed.
+		//
+		// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
+		// See /RATIONALE.md
+		Closed atomic.Uint64
+
 		// CodeCloser is non-nil when the code should be closed after this module.
 		CodeCloser api.Closer
 

runtime.go

@@ -159,7 +159,6 @@ func NewRuntimeWithConfig(ctx context.Context, rConfig RuntimeConfig) Runtime {
 		engine = config.newEngine(ctx, config.enabledFeatures, nil)
 	}
 	store := wasm.NewStore(config.enabledFeatures, engine)
-	zero := uint64(0)
 	return &runtime{
 		cache:                 cacheImpl,
 		store:                 store,
@@ -168,7 +167,6 @@ func NewRuntimeWithConfig(ctx context.Context, rConfig RuntimeConfig) Runtime {
 		memoryCapacityFromMax: config.memoryCapacityFromMax,
 		dwarfDisabled:         config.dwarfDisabled,
 		storeCustomSections:   config.storeCustomSections,
-		closed:                &zero,
 		ensureTermination:     config.ensureTermination,
 	}
 }
@@ -189,7 +187,7 @@ type runtime struct {
 	//
 	// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
 	// See /RATIONALE.md
-	closed *uint64
+	closed atomic.Uint64
 
 	ensureTermination bool
 }
@@ -259,7 +257,7 @@ func buildFunctionListeners(ctx context.Context, internal *wasm.Module) ([]exper
 
 // failIfClosed returns an error if CloseWithExitCode was called implicitly (by Close) or explicitly.
 func (r *runtime) failIfClosed() error {
-	if closed := atomic.LoadUint64(r.closed); closed != 0 {
+	if closed := r.closed.Load(); closed != 0 {
 		return fmt.Errorf("runtime closed with exit_code(%d)", uint32(closed>>32))
 	}
 	return nil
@@ -362,7 +360,7 @@ func (r *runtime) Close(ctx context.Context) error {
 // Note: it also marks the internal `closed` field
 func (r *runtime) CloseWithExitCode(ctx context.Context, exitCode uint32) error {
 	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
-	if !atomic.CompareAndSwapUint64(r.closed, 0, closed) {
+	if !r.closed.CompareAndSwap(0, closed) {
 		return nil
 	}
 	err := r.store.CloseWithExitCode(ctx, exitCode)