Repository: tetratelabs/wazero (go)

Issue:
Potential data race on sys.FS when context cancelation enabled. 
Hitting a race condition that I don't quite understand.  Made a reproducible test for it here:

https://github.com/abraithwaite/wazero/commit/ca10894d


**Environment (please complete the relevant information):**
- Go version: 1.20
- wazero Version: tested v1.0.1 and 8a9d3f28
- Host architecture: arm64 (darwin) and amd64 (github runners)
- Runtime mode: Compiler

Not sure if the way we're using wazero is crazy, so please feel free to comment on that if you think this is improper usage of the library. 🙇‍♂️


Pull request description:
Defer resource closure on context cancelation to m.FailIfClosed
Fixes #1314. Previously, when context cancelation is enabled, the resource is 
also closed in the `closeModuleOnCanceledOrTimeout` function which is 
asynchronously called concurrently to the Wasm function call. On the other hand,
host functions, especially wasi functions, do not check if the module is closed or not
therefore, if we reached the context cancelation while executing wasi functions, we
had race conditions.

This makes change to `closeModuleOnCanceledOrTimeout` so that it won't close resources,
and defer closing them to in `FailIfClosed` which is called synchronously during `api.Function.Call`.

Patch:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 8a90e083..a7057825 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -4958,7 +4958,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 		amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceOffset, amd64CallingConventionDestinationFunctionModuleInstanceAddressRegister)
 	jmpIfModuleNotChange := c.assembler.CompileJump(amd64.JEQ)
 
-	// If engine.CallContext.ModuleInstanceAddress is not equal the value on amd64CallingConventionDestinationFunctionModuleInstanceAddressRegister,
+	// If engine.ModuleContext.ModuleInstance is not equal the value on amd64CallingConventionDestinationFunctionModuleInstanceAddressRegister,
 	// we have to put the new value there.
 	c.assembler.CompileRegisterToMemory(amd64.MOVQ, amd64CallingConventionDestinationFunctionModuleInstanceAddressRegister,
 		amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceOffset)
diff --git a/internal/wasm/module_instance.go b/internal/wasm/module_instance.go
index d7d7fb68..a12f5b6d 100644
--- a/internal/wasm/module_instance.go
+++ b/internal/wasm/module_instance.go
@@ -13,6 +13,13 @@ import (
 // FailIfClosed returns a sys.ExitError if CloseWithExitCode was called.
 func (m *ModuleInstance) FailIfClosed() (err error) {
 	if closed := atomic.LoadUint64(&m.Closed); closed != 0 {
+		switch closed & exitCodeFlagMask {
+		case exitCodeFlagResourceClosed:
+		case exitCodeFlagResourceNotClosed:
+			// This happens when this module is closed asynchronously in CloseModuleOnCanceledOrTimeout,
+			// and the closure of resources have been deferred here.
+			_ = m.ensureResourcesClosed(context.Background())
+		}
 		return sys.NewExitError(uint32(closed >> 32)) // Unpack the high order bits as the exit code.
 	}
 	return nil
@@ -45,7 +52,16 @@ func (m *ModuleInstance) closeModuleOnCanceledOrTimeout(ctx context.Context, can
 			// case go will randomize which branch of the outer select to enter
 			// and we don't want to close the module.
 		default:
-			m.CloseWithCtxErr(ctx)
+			// This is the same logic as CloseWithCtxErr except this calls closeWithExitCodeWithoutClosingResource
+			// so that we can defer the resource closure in FailIfClosed.
+			switch {
+			case errors.Is(ctx.Err(), context.Canceled):
+				// TODO: figure out how to report error here.
+				_ = m.closeWithExitCodeWithoutClosingResource(sys.ExitCodeContextCanceled)
+			case errors.Is(ctx.Err(), context.DeadlineExceeded):
+				// TODO: figure out how to report error here.
+				_ = m.closeWithExitCodeWithoutClosingResource(sys.ExitCodeDeadlineExceeded)
+			}
 		}
 	case <-cancelChan:
 	}
@@ -83,27 +99,46 @@ func (m *ModuleInstance) Close(ctx context.Context) (err error) {
 
 // CloseWithExitCode implements the same method as documented on api.Module.
 func (m *ModuleInstance) CloseWithExitCode(ctx context.Context, exitCode uint32) (err error) {
-	if !m.setExitCode(exitCode) {
+	if !m.setExitCode(exitCode, exitCodeFlagResourceClosed) {
 		return nil // not an error to have already closed
 	}
 	_ = m.s.deleteModule(m.moduleListNode)
 	return m.ensureResourcesClosed(ctx)
 }
 
+func (m *ModuleInstance) closeWithExitCodeWithoutClosingResource(exitCode uint32) (err error) {
+	if !m.setExitCode(exitCode, exitCodeFlagResourceNotClosed) {
+		return nil // not an error to have already closed
+	}
+	_ = m.s.deleteModule(m.moduleListNode)
+	return nil
+}
+
 // closeWithExitCode is the same as CloseWithExitCode besides this doesn't delete it from Store.moduleList.
 func (m *ModuleInstance) closeWithExitCode(ctx context.Context, exitCode uint32) (err error) {
-	if !m.setExitCode(exitCode) {
+	if !m.setExitCode(exitCode, exitCodeFlagResourceClosed) {
 		return nil // not an error to have already closed
 	}
 	return m.ensureResourcesClosed(ctx)
 }
 
-func (m *ModuleInstance) setExitCode(exitCode uint32) bool {
-	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
+type exitCodeFlag = uint64
+
+const exitCodeFlagMask = 0xff
+
+const (
+	// exitCodeFlagResourceClosed indicates that the module was closed and resources were already closed.
+	exitCodeFlagResourceClosed = 1 << iota
+	// exitCodeFlagResourceNotClosed indicates that the module was closed while resources are not closed yet.
+	exitCodeFlagResourceNotClosed
+)
+
+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)
 }
 
-// ensureResourcesClosed ensures that resources assigned to CallContext is released.
+// ensureResourcesClosed ensures that resources assigned to ModuleInstance is released.
 // Multiple calls to this function is safe.
 func (m *ModuleInstance) ensureResourcesClosed(ctx context.Context) (err error) {
 	if sysCtx := m.Sys; sysCtx != nil { // nil if from HostModuleBuilder
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 0b1cfb00..c56c134e 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -99,7 +99,7 @@ type (
 		//
 		// # Notes
 		//
-		//   - This is a part of CallContext so that scope and Close is coherent.
+		//   - This is a part of ModuleInstance so that scope and Close is coherent.
 		//   - This is not exposed outside this repository (as a host function
 		//	  parameter) because we haven't thought through capabilities based
 		//	  security implications.
@@ -271,7 +271,7 @@ func NewStore(enabledFeatures api.CoreFeatures, engine Engine) *Store {
 //
 // * ctx: the default context used for function calls.
 // * name: the name of the module.
-// * sys: the system context, which will be closed (SysContext.Close) on CallContext.Close.
+// * sys: the system context, which will be closed (SysContext.Close) on ModuleInstance.Close.
 //
 // Note: Module.Validate must be called prior to instantiation.
 func (s *Store) Instantiate(
diff --git a/site/content/community/users.md b/site/content/community/users.md
index b917f692..007d23de 100644
--- a/site/content/community/users.md
+++ b/site/content/community/users.md
@@ -16,6 +16,7 @@ considering their efforts before starting your own!
 | [go-re2][7]      | high performance regular expressions                                                                 |
 | [go-sqlite3][11] | [SQLite][12] bindings, `database/sql` driver                                                         |
 | [mjml-go][19]    | Compile [MJML][20] to HTML directly in Go                                                            |
+| [wazergo][29]    | Generics library for type-safe and high performance wazero host modules                              |
 | [Wetware][28]    | Simple, secure & scalable clusters                                                                   |
 
 ### General purpose plugins
@@ -117,3 +118,5 @@ experience.
 [27]: https://github.com/yomorun/yomo/tree/master/example/7-wasm
 
 [28]: https://github.com/wetware/ww
+
+[29]: https://github.com/stealthrocket/wazergo


Test patch:
diff --git a/internal/wasm/module_instance_test.go b/internal/wasm/module_instance_test.go
index 9e848364..022edc3f 100644
--- a/internal/wasm/module_instance_test.go
+++ b/internal/wasm/module_instance_test.go
@@ -16,7 +16,7 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-func TestCallContext_String(t *testing.T) {
+func TestModuleInstance_String(t *testing.T) {
 	s := newStore()
 
 	tests := []struct {
@@ -57,7 +57,7 @@ func TestCallContext_String(t *testing.T) {
 	}
 }
 
-func TestCallContext_Close(t *testing.T) {
+func TestModuleInstance_Close(t *testing.T) {
 	s := newStore()
 
 	tests := []struct {
@@ -67,15 +67,15 @@ func TestCallContext_Close(t *testing.T) {
 	}{
 		{
 			name: "Close()",
-			closer: func(ctx context.Context, callContext *ModuleInstance) error {
-				return callContext.Close(ctx)
+			closer: func(ctx context.Context, m *ModuleInstance) error {
+				return m.Close(ctx)
 			},
 			expectedClosed: uint64(1),
 		},
 		{
 			name: "CloseWithExitCode(255)",
-			closer: func(ctx context.Context, callContext *ModuleInstance) error {
-				return callContext.CloseWithExitCode(ctx, 255)
+			closer: func(ctx context.Context, m *ModuleInstance) error {
+				return m.CloseWithExitCode(ctx, 255)
 			},
 			expectedClosed: uint64(255)<<32 + 1,
 		},
@@ -161,7 +161,7 @@ func TestCallContext_Close(t *testing.T) {
 	})
 }
 
-func TestCallContext_CallDynamic(t *testing.T) {
+func TestModuleInstance_CallDynamic(t *testing.T) {
 	s := newStore()
 
 	tests := []struct {
@@ -171,15 +171,15 @@ func TestCallContext_CallDynamic(t *testing.T) {
 	}{
 		{
 			name: "Close()",
-			closer: func(ctx context.Context, callContext *ModuleInstance) error {
-				return callContext.Close(ctx)
+			closer: func(ctx context.Context, m *ModuleInstance) error {
+				return m.Close(ctx)
 			},
 			expectedClosed: uint64(1),
 		},
 		{
 			name: "CloseWithExitCode(255)",
-			closer: func(ctx context.Context, callContext *ModuleInstance) error {
-				return callContext.CloseWithExitCode(ctx, 255)
+			closer: func(ctx context.Context, m *ModuleInstance) error {
+				return m.CloseWithExitCode(ctx, 255)
 			},
 			expectedClosed: uint64(255)<<32 + 1,
 		},
@@ -259,10 +259,10 @@ func TestCallContext_CallDynamic(t *testing.T) {
 	})
 }
 
-func TestCallContext_CloseModuleOnCanceledOrTimeout(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}
+		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		const duration = time.Second
 		ctx, cancel := context.WithTimeout(context.Background(), duration)
 		defer cancel()
@@ -270,12 +270,19 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		time.Sleep(duration * 2)
 		defer done()
 
+		// Resource shouldn't be released at this point.
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed&exitCodeFlagMask)
+		require.NotNil(t, cc.Sys)
+
 		err := cc.FailIfClosed()
 		require.EqualError(t, err, "module closed with context deadline exceeded")
+
+		// The resource must be closed in FailIfClosed.
+		require.Nil(t, cc.Sys)
 	})
 
 	t.Run("cancel", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{Closed: 0, 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()
@@ -283,14 +290,21 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		cancel()
 		cancel()
 		defer done()
-
 		time.Sleep(time.Second)
+
+		// Resource shouldn't be released at this point.
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed&exitCodeFlagMask)
+		require.NotNil(t, cc.Sys)
+
 		err := cc.FailIfClosed()
 		require.EqualError(t, err, "module closed with context canceled")
+
+		// The resource must be closed in FailIfClosed.
+		require.Nil(t, cc.Sys)
 	})
 
 	t.Run("timeout over cancel", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s, Sys: internalsys.DefaultContext(nil)}
 		const duration = time.Second
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
@@ -300,10 +314,20 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		done := cc.CloseModuleOnCanceledOrTimeout(context.WithValue(ctx, struct{}{}, 1)) // Wrapping arbitrary context.
 		time.Sleep(duration * 2)
 		defer done()
+
+		// Resource shouldn't be released at this point.
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed&exitCodeFlagMask)
+		require.NotNil(t, cc.Sys)
+
+		err := cc.FailIfClosed()
+		require.EqualError(t, err, "module closed with context deadline exceeded")
+
+		// The resource must be closed in FailIfClosed.
+		require.Nil(t, cc.Sys)
 	})
 
 	t.Run("cancel over timeout", func(t *testing.T) {
-		cc := &ModuleInstance{Closed: 0, ModuleName: "test", s: s}
+		cc := &ModuleInstance{Closed: 0, 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
@@ -315,8 +339,16 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 		defer done()
 
 		time.Sleep(time.Second)
+
+		// Resource shouldn't be released at this point.
+		require.Equal(t, exitCodeFlag(exitCodeFlagResourceNotClosed), cc.Closed&exitCodeFlagMask)
+		require.NotNil(t, cc.Sys)
+
 		err := cc.FailIfClosed()
 		require.EqualError(t, err, "module closed with context canceled")
+
+		// The resource must be closed in FailIfClosed.
+		require.Nil(t, cc.Sys)
 	})
 
 	t.Run("cancel works", func(t *testing.T) {
@@ -348,7 +380,7 @@ func TestCallContext_CloseModuleOnCanceledOrTimeout(t *testing.T) {
 	})
 }
 
-func TestCallContext_CloseWithCtxErr(t *testing.T) {
+func TestModuleInstance_CloseWithCtxErr(t *testing.T) {
 	s := newStore()
 
 	t.Run("context canceled", func(t *testing.T) {
@@ -393,7 +425,7 @@ func (m *mockCloser) Close(context.Context) error {
 	return nil
 }
 
-func TestCallContext_ensureResourcesClosed(t *testing.T) {
+func TestModuleInstance_ensureResourcesClosed(t *testing.T) {
 	closer := &mockCloser{}
 
 	for _, tc := range []struct {
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 3bbeb94e..eb8646a8 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -113,7 +113,7 @@ func TestStore_Instantiate(t *testing.T) {
 	require.NoError(t, err)
 	defer mod.Close(testCtx)
 
-	t.Run("CallContext defaults", func(t *testing.T) {
+	t.Run("ModuleInstance defaults", func(t *testing.T) {
 		require.Equal(t, s.nameToNode["bar"].module, mod)
 		require.Equal(t, s.nameToNode["bar"].module.MemoryInstance, mod.MemoryInstance)
 		require.Equal(t, s, mod.s)