Repository: tetratelabs/wazero
Pull request #988 (2022-12-31 03:52:26)

compiler: force moduleContext initialization after Go function calls
fixes https://github.com/tetratelabs/wazero/issues/978 and supersedes #980 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

mathetake: 🤞 
mathetake: ok cool passing now except seems flaky on macos with Go 1.17

```
--- FAIL: Test_stdio_large (1.74s)
    require.go:312: expected "stderr 2097152
        ", but was "fatal error: malloc deadlock
        
        runtime stack:
        runtime.throw({0x6b991, 0xf})
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/panic.go:1198 +0x7
        runtime.mallocgc(0x10, 0x42ca0, 0x1)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/malloc.go:972 +0xcf
        runtime.newobject(0x42ca0)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/malloc.go:12[34](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:35) +0x4
        runtime.handleEvent()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/lock_js.go:238 +0x2
        runtime.sysReserve(0x0, 0x800000)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mem_js.go:68 +0xe
        runtime.sysReserveAligned(0x0, 0x400000, 0x400000)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/malloc.go:805 +0x2
        runtime.(*mheap).sysAlloc(0x32a700, 0x400000)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/malloc.go:698 +0x4f
        runtime.(*mheap).grow(0x32a700, 0x122)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mheap.go:1347 +0xa
        runtime.(*mheap).allocSpan(0x32a700, 0x122, 0x0, 0x1)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mheap.go:1179 +0x3b
        runtime.(*mheap).alloc.func1()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mheap.go:913 +0x6
        runtime.systemstack()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/asm_wasm.s:170 +0x2
        runtime.mstart()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/asm_wasm.s:28
        
        goroutine 1 [running]:
        runtime.systemstack_switch()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/asm_wasm.s:181 fp=0x42ace8 sp=0x42ace0 pc=0x13ac0000
        runtime.(*mheap).alloc(0x32a700, 0x122, 0x1, 0x0)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mheap.go:907 +0x2 fp=0x42ad30 sp=0x42ace8 pc=0x11750002
        runtime.(*mcache).allocLarge(0x[35](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:36)0108, 0x24[40](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:41)00, 0x0, 0x1)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/mcache.go:227 +0x9 fp=0x[42](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:43)ad90 sp=0x42ad30 pc=0x10ed0009
        runtime.mallocgc(0x244000, 0x0, 0x0)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/malloc.go:1088 +0x8f fp=0x42ae18 sp=0x42ad90 pc=0x10a4008f
        runtime.growslice(0x274a0, {0x500000, 0x1d0000, 0x1d0000}, 0x1d0001)
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/slice.go:261 +0x83 fp=0x42ae68 sp=0x42ae18 pc=0x12c60083
        io.ReadAll({0xba720, 0x40c018})
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/io/io.go:631 +0x7 fp=0x42aee0 sp=0x42ae68 pc=0x15530007
        github.com/tetratelabs/wazero/internal/gojs/testdata/stdio.Main()
        	/Users/runner/work/wazero/wazero/internal/gojs/testdata/stdio/main.go:10 +0x2 fp=0x42af28 sp=0x42aee0 pc=0x1f2e0002
        main.main()
        	/Users/runner/work/wazero/wazero/internal/gojs/testdata/main.go:40 +0x2a fp=0x42af88 sp=0x42af28 pc=0x1f33002a
        runtime.main()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/proc.go:255 +0x35 fp=0x42afe0 sp=0x42af88 pc=0x121f0035
        runtime.goexit()
        	/Users/runner/hostedtoolcache/go/1.17.13/x64/src/runtime/asm_wasm.s:[43](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:44)1 +0x1 fp=0x42afe8 sp=0x42afe0 pc=0x13d[50](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:51)001
        stderr 20971[52](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:53)
        "
        /Users/runner/work/wazero/wazero/internal/gojs/misc_test.go:[56](https://github.com/tetratelabs/wazero/actions/runs/3810497727/jobs/6482610815#step:4:57)
```
mathetake: ok added the regression engine test case, and feeling good 😉 

Patch:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 60163b6d..c7b13152 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -186,6 +186,12 @@ func (c *amd64Compiler) compileGoDefinedHostFunction() error {
 
 	// Initializes the reserved stack base pointer which is used to retrieve the call frame stack.
 	c.compileReservedStackBasePointerInitialization()
+
+	// Go function can change the module state in arbitrary way, so we have to force
+	// the callEngine.moduleContext initialization on the function return. To do so,
+	// we zero-out callEngine.moduleInstanceAddress.
+	c.assembler.CompileConstToMemory(amd64.MOVQ,
+		0, amd64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset)
 	return c.compileReturnFunction()
 }
 
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index 98793c56..692fff04 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -398,6 +398,14 @@ func (c *arm64Compiler) compileGoDefinedHostFunction() error {
 
 	// Initializes the reserved stack base pointer which is used to retrieve the call frame stack.
 	c.compileReservedStackBasePointerRegisterInitialization()
+
+	// Go function can change the module state in arbitrary way, so we have to force
+	// the callEngine.moduleContext initialization on the function return. To do so,
+	// we zero-out callEngine.moduleInstanceAddress.
+	c.assembler.CompileRegisterToMemory(arm64.STRD,
+		arm64.RegRZR,
+		arm64ReservedRegisterForCallEngine, callEngineModuleContextModuleInstanceAddressOffset)
+
 	return c.compileReturnFunction()
 }
 


Test patch:
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index 6001d794..d9499343 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -58,6 +58,11 @@ func TestCompiler_Engine_NewModuleEngine(t *testing.T) {
 	enginetest.RunTestEngine_NewModuleEngine(t, et)
 }
 
+func TestCompiler_MemoryGrowInRecursiveCall(t *testing.T) {
+	defer functionLog.Reset()
+	enginetest.RunTestEngine_MemoryGrowInRecursiveCall(t, et)
+}
+
 func TestCompiler_ModuleEngine_LookupFunction(t *testing.T) {
 	defer functionLog.Reset()
 	enginetest.RunTestModuleEngine_LookupFunction(t, et)
diff --git a/internal/engine/interpreter/interpreter_test.go b/internal/engine/interpreter/interpreter_test.go
index 55e95372..ff58ac08 100644
--- a/internal/engine/interpreter/interpreter_test.go
+++ b/internal/engine/interpreter/interpreter_test.go
@@ -95,6 +95,11 @@ func (e engineTester) CompiledFunctionPointerValue(me wasm.ModuleEngine, funcInd
 	return uint64(uintptr(unsafe.Pointer(internal.functions[funcIndex])))
 }
 
+func TestInterpreter_MemoryGrowInRecursiveCall(t *testing.T) {
+	defer functionLog.Reset()
+	enginetest.RunTestEngine_MemoryGrowInRecursiveCall(t, et)
+}
+
 func TestInterpreter_Engine_NewModuleEngine(t *testing.T) {
 	enginetest.RunTestEngine_NewModuleEngine(t, et)
 }
diff --git a/internal/gojs/misc_test.go b/internal/gojs/misc_test.go
index 510567aa..b48c1113 100644
--- a/internal/gojs/misc_test.go
+++ b/internal/gojs/misc_test.go
@@ -45,7 +45,6 @@ func Test_stdio(t *testing.T) {
 }
 
 func Test_stdio_large(t *testing.T) {
-	t.Skip("TODO: #980 memory out of bounds when run with compiler")
 	t.Parallel()
 
 	size := 2 * 1024 * 1024 // 2MB
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index f735d256..9abe56d7 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -55,6 +55,64 @@ type EngineTester interface {
 	CompiledFunctionPointerValue(tme wasm.ModuleEngine, funcIndex wasm.Index) uint64
 }
 
+// RunTestEngine_MemoryGrowInRecursiveCall ensures that it's safe to grow memory in the recursive Wasm calls.
+func RunTestEngine_MemoryGrowInRecursiveCall(t *testing.T, et EngineTester) {
+	enabledFeatures := api.CoreFeaturesV1
+	e := et.NewEngine(enabledFeatures)
+	s, ns := wasm.NewStore(enabledFeatures, e)
+
+	const hostModuleName = "env"
+	const hostFnName = "grow_memory"
+	var growFn api.Function
+	hm, err := wasm.NewHostModule(hostModuleName, map[string]interface{}{hostFnName: func() {
+		// Does the recursive call into Wasm, which grows memory.
+		_, err := growFn.Call(context.Background())
+		require.NoError(t, err)
+	}}, map[string]*wasm.HostFuncNames{hostFnName: {}}, enabledFeatures)
+	require.NoError(t, err)
+
+	err = s.Engine.CompileModule(testCtx, hm, nil)
+	require.NoError(t, err)
+
+	_, err = s.Instantiate(testCtx, ns, hm, hostModuleName, nil)
+	require.NoError(t, err)
+
+	m := &wasm.Module{
+		TypeSection:     []*wasm.FunctionType{{Params: []wasm.ValueType{}, Results: []wasm.ValueType{}}},
+		FunctionSection: []wasm.Index{0, 0},
+		CodeSection: []*wasm.Code{
+			{
+				Body: []byte{
+					// Calls the imported host function, which in turn calls the next in-Wasm function recursively.
+					wasm.OpcodeCall, 0,
+					// Access the memory and this should succeed as we already had memory grown at this point.
+					wasm.OpcodeI32Const, 0,
+					wasm.OpcodeI32Load, 0x2, 0x0,
+					wasm.OpcodeDrop,
+					wasm.OpcodeEnd,
+				},
+			},
+			{
+				// Grows memory by 1 page.
+				Body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeMemoryGrow, wasm.OpcodeDrop, wasm.OpcodeEnd},
+			},
+		},
+		MemorySection: &wasm.Memory{Max: 1000},
+		ImportSection: []*wasm.Import{{Module: hostModuleName, Name: hostFnName, DescFunc: 0}},
+	}
+	m.BuildFunctionDefinitions()
+
+	err = s.Engine.CompileModule(testCtx, m, nil)
+	require.NoError(t, err)
+
+	inst, err := s.Instantiate(testCtx, ns, m, t.Name(), nil)
+	require.NoError(t, err)
+
+	growFn = inst.Function(2)
+	_, err = inst.Function(1).Call(context.Background())
+	require.NoError(t, err)
+}
+
 func RunTestEngine_NewModuleEngine(t *testing.T, et EngineTester) {
 	e := et.NewEngine(api.CoreFeaturesV1)