compiler: zero-initializes moduleInstanceAddress of call engine (#783)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

internal/engine/compiler/engine.go

@@ -101,6 +101,9 @@ type (
 	moduleContext struct {
 		// moduleInstanceAddress is the address of module instance from which we initialize
 		// the following fields. This is set whenever we enter a function or return from function calls.
+		//
+		// On the entry to the native code, this must be initialized to zero to let native code preamble know
+		// that this is the initial function call (which leads to moduleContext initialization pass).
 		moduleInstanceAddress uintptr //lint:ignore U1000 This is only used by Compiler code.
 
 		// globalElement0Address is the address of the first element in the global slice,
@@ -554,14 +557,11 @@ func (ce *callEngine) Call(ctx context.Context, callCtx *wasm.CallContext, param
 	// and we have to make sure that all the runtime errors, including the one happening inside
 	// host functions, will be captured as errors, not panics.
 	defer func() {
-		// If the module closed during the call, and the call didn't err for another reason, set an ExitError.
+		err = ce.deferredOnCall(recover())
 		if err == nil {
+			// If the module closed during the call, and the call didn't err for another reason, set an ExitError.
 			err = callCtx.FailIfClosed()
-		}
-		// TODO: ^^ Will not fail if the function was imported from a closed module.
-
-		if v := recover(); v != nil {
-			err = ce.recoverOnCall(v)
+			// TODO: ^^ Will not fail if the function was imported from a closed module.
 		}
 	}()
 
@@ -573,19 +573,23 @@ func (ce *callEngine) Call(ctx context.Context, callCtx *wasm.CallContext, param
 	return
 }
 
-// recoverOnCall takes the recovered value `recoverOnCall`, and wraps it
-// with the call frame stack traces. Also, reset the state of callEngine
-// so that it can be used for the subsequent calls.
-func (ce *callEngine) recoverOnCall(v interface{}) (err error) {
-	builder := wasmdebug.NewErrorBuilder()
-	for i := uint64(0); i < ce.callFrameStackPointer; i++ {
-		def := ce.callFrameStack[ce.callFrameStackPointer-1-i].function.source.FunctionDefinition
-		builder.AddFrame(def.DebugName(), def.ParamTypes(), def.ResultTypes())
+// deferredOnCall takes the recovered value `recovered`, and wraps it
+// with the call frame stack traces when not nil. This also resets
+// the state of callEngine so that it can be used for the subsequent calls.
+//
+// This is defined for testability.
+func (ce *callEngine) deferredOnCall(recovered interface{}) (err error) {
+	if recovered != nil {
+		builder := wasmdebug.NewErrorBuilder()
+		for i := uint64(0); i < ce.callFrameStackPointer; i++ {
+			def := ce.callFrameStack[ce.callFrameStackPointer-1-i].function.source.FunctionDefinition
+			builder.AddFrame(def.DebugName(), def.ParamTypes(), def.ResultTypes())
+		}
+		err = builder.FromRecovered(recovered)
 	}
-	err = builder.FromRecovered(v)
 
 	// Allows the reuse of CallEngine.
-	ce.stackPointer, ce.callFrameStackPointer = 0, 0
+	ce.stackBasePointerInBytes, ce.stackPointer, ce.callFrameStackPointer, ce.moduleInstanceAddress = 0, 0, 0, 0
 	return
 }
 

internal/engine/compiler/engine_test.go

@@ -305,7 +305,7 @@ func TestCallEngine_builtinFunctionTableGrow(t *testing.T) {
 	require.Equal(t, uintptr(0xff), table.References[0])
 }
 
-func TestCallEngine_recoverOnCall(t *testing.T) {
+func TestCallEngine_deferredOnCall(t *testing.T) {
 	ce := &callEngine{
 		valueStack:        make([]uint64, 100),
 		valueStackContext: valueStackContext{stackPointer: 3},
@@ -317,11 +317,12 @@ func TestCallEngine_recoverOnCall(t *testing.T) {
 			{function: &function{source: &wasm.FunctionInstance{FunctionDefinition: newMockFunctionDefinition("4")}}},
 			{function: &function{source: &wasm.FunctionInstance{FunctionDefinition: newMockFunctionDefinition("5")}}},
 		},
+		moduleContext: moduleContext{moduleInstanceAddress: 0xdeafbeaf},
 	}
 
 	beforeRecoverValueStack, beforeRecoverCallFrameStack := ce.valueStack, ce.callFrameStack
 
-	err := ce.recoverOnCall(errors.New("some error"))
+	err := ce.deferredOnCall(errors.New("some error"))
 	require.EqualError(t, err, `some error (recovered by wazero)
 wasm stack trace:
 	5()
@@ -330,10 +331,12 @@ wasm stack trace:
 	2()
 	1()`)
 
-	// After recover, the stack pointers must be reset, but the underlying slices must be intact
-	// for the subsequent calls.
+	// After recover, the state of callEngine must be reset except that the underlying slices must be intact
+	// for the subsequent calls to avoid additional allocations on each call.
 	require.Equal(t, uint64(0), ce.stackBasePointerInBytes)
+	require.Equal(t, uint64(0), ce.stackPointer)
 	require.Equal(t, uint64(0), ce.callFrameStackPointer)
+	require.Equal(t, uintptr(0), ce.moduleInstanceAddress)
 	require.Equal(t, beforeRecoverValueStack, ce.valueStack)
 	require.Equal(t, beforeRecoverCallFrameStack, ce.callFrameStack)
 }

internal/integration_test/engine/adhoc_test.go

@@ -558,50 +558,53 @@ func testCloseInFlight(t *testing.T, r wazero.Runtime) {
 
 func testMemOps(t *testing.T, r wazero.Runtime) {
 	// Instantiate a module that manages its memory
-	memory, err := r.InstantiateModuleFromBinary(testCtx, memoryWasm)
+	mod, err := r.InstantiateModuleFromBinary(testCtx, memoryWasm)
 	require.NoError(t, err)
-	defer memory.Close(testCtx)
+	defer mod.Close(testCtx)
 
 	// Check the export worked
-	require.Equal(t, memory.Memory(), memory.ExportedMemory("memory"))
+	require.Equal(t, mod.Memory(), mod.ExportedMemory("memory"))
+	memory := mod.Memory()
+
+	sizeFn, storeFn, growFn := mod.ExportedFunction("size"), mod.ExportedFunction("store"), mod.ExportedFunction("grow")
 
 	// Check the size command worked
-	results, err := memory.ExportedFunction("size").Call(testCtx)
+	results, err := sizeFn.Call(testCtx)
 	require.NoError(t, err)
 	require.Zero(t, results[0])
-	require.Zero(t, memory.ExportedMemory("memory").Size(testCtx))
+	require.Zero(t, memory.Size(testCtx))
 
 	// Any offset should be out of bounds error even when it is less than memory capacity(=memoryCapacityPages).
-	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
+	_, err = storeFn.Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
 	require.Error(t, err) // Out of bounds error.
 
 	// Try to grow the memory by one page
-	results, err = memory.ExportedFunction("grow").Call(testCtx, 1)
+	results, err = growFn.Call(testCtx, 1)
 	require.NoError(t, err)
 	require.Zero(t, results[0]) // should succeed and return the old size in pages.
 
 	// Any offset larger than the current size should be out of bounds error even when it is less than memory capacity.
-	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
+	_, err = storeFn.Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
 	require.Error(t, err) // Out of bounds error.
 
 	// Check the size command works!
-	results, err = memory.ExportedFunction("size").Call(testCtx)
+	results, err = sizeFn.Call(testCtx)
 	require.NoError(t, err)
-	require.Equal(t, uint64(1), results[0])                        // 1 page
-	require.Equal(t, uint32(65536), memory.Memory().Size(testCtx)) // 64KB
+	require.Equal(t, uint64(1), results[0])               // 1 page
+	require.Equal(t, uint32(65536), memory.Size(testCtx)) // 64KB
 
 	// Grow again so that the memory size matches memory capacity.
-	results, err = memory.ExportedFunction("grow").Call(testCtx, 1)
+	results, err = growFn.Call(testCtx, 1)
 	require.NoError(t, err)
 	require.Equal(t, uint64(1), results[0])
 
 	// Verify the size matches cap.
-	results, err = memory.ExportedFunction("size").Call(testCtx)
+	results, err = sizeFn.Call(testCtx)
 	require.NoError(t, err)
 	require.Equal(t, uint64(memoryCapacityPages), results[0])
 
 	// Now the store instruction at the memory capcity bound should succeed.
-	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8) // i64.store needs 8 bytes from offset.
+	_, err = storeFn.Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8) // i64.store needs 8 bytes from offset.
 	require.NoError(t, err)
 }