Propagates api.Memory.Grow by users to ModuleEngine (#2216)

Since the introduction of optimizing compiler, the Memory.Grow by the users, which has the different call path than memory.grow instruction, didn't propagate the growth result onto compiler's ModuleInstance. Fixes #2215 Signed-off-by: Takeshi Yoneda <[email protected]>
tetratelabs · May 27, 2024 · 4dfb1f5 · 4dfb1f5
1 parent eb24363
commit 4dfb1f5
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 42 deletions.
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
@@ -98,6 +98,9 @@ func (e *moduleEngine) SetGlobalValue(idx wasm.Index, lo, hi uint64) {
 // OwnsGlobals implements the same method as documented on wasm.ModuleEngine.
 func (e *moduleEngine) OwnsGlobals() bool { return false }
 
+// MemoryGrown implements wasm.ModuleEngine.
+func (e *moduleEngine) MemoryGrown() {}
+
 // callEngine holds context per moduleEngine.Call, and shared across all the
 // function calls originating from the same moduleEngine.Call execution.
 //

diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
@@ -2,7 +2,6 @@ package wazevo
 
 import (
 	"context"
-	"encoding/binary"
 	"fmt"
 	"reflect"
 	"runtime"
@@ -310,15 +309,6 @@ func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint6
 				*argRes = uint64(0xffffffff) // = -1 in signed 32-bit integer.
 			} else {
 				*argRes = uint64(res)
-				calleeOpaque := opaqueViewFromPtr(uintptr(unsafe.Pointer(c.execCtx.callerModuleContextPtr)))
-				if mod.Source.MemorySection != nil { // Local memory.
-					putLocalMemory(calleeOpaque, 8 /* local memory begins at 8 */, mem)
-				} else {
-					// Imported memory's owner at offset 16 of the callerModuleContextPtr.
-					opaquePtr := uintptr(binary.LittleEndian.Uint64(calleeOpaque[16:]))
-					importedMemOwner := opaqueViewFromPtr(opaquePtr)
-					putLocalMemory(importedMemOwner, 8 /* local memory begins at 8 */, mem)
-				}
 			}
 			c.execCtx.exitCode = wazevoapi.ExitCodeOK
 			afterGoFunctionCallEntrypoint(c.execCtx.goCallReturnAddress, c.execCtxPtr, uintptr(unsafe.Pointer(c.execCtx.stackPointerBeforeGoCall)), c.execCtx.framePointerBeforeGoCall)
@@ -525,14 +515,6 @@ func (c *callEngine) callerModuleInstance() *wasm.ModuleInstance {
 	return moduleInstanceFromOpaquePtr(c.execCtx.callerModuleContextPtr)
 }
 
-func opaqueViewFromPtr(ptr uintptr) []byte {
-	var opaque []byte
-	sh := (*reflect.SliceHeader)(unsafe.Pointer(&opaque))
-	sh.Data = ptr
-	setSliceLimits(sh, 24, 24)
-	return opaque
-}
-
 const callStackCeiling = uintptr(50000000) // in uint64 (8 bytes) == 400000000 bytes in total == 400mb.
 
 func (c *callEngine) growStackWithGuarded() (newSP uintptr, newFP uintptr, err error) {

diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
@@ -86,16 +86,6 @@ func newAlignedOpaque(size int) moduleContextOpaque {
 	return buf
 }
 
-func putLocalMemory(opaque []byte, offset wazevoapi.Offset, mem *wasm.MemoryInstance) {
-	s := uint64(len(mem.Buffer))
-	var b uint64
-	if len(mem.Buffer) > 0 {
-		b = uint64(uintptr(unsafe.Pointer(&mem.Buffer[0])))
-	}
-	binary.LittleEndian.PutUint64(opaque[offset:], b)
-	binary.LittleEndian.PutUint64(opaque[offset+8:], s)
-}
-
 func (m *moduleEngine) setupOpaque() {
 	inst := m.module
 	offsets := &m.parent.offsets
@@ -106,7 +96,7 @@ func (m *moduleEngine) setupOpaque() {
 	)
 
 	if lm := offsets.LocalMemoryBegin; lm >= 0 {
-		putLocalMemory(opaque, lm, inst.MemoryInstance)
+		m.putLocalMemory()
 	}
 
 	// Note: imported memory is resolved in ResolveImportedFunction.
@@ -227,6 +217,25 @@ func (m *moduleEngine) SetGlobalValue(i wasm.Index, lo, hi uint64) {
 // OwnsGlobals implements the same method as documented on wasm.ModuleEngine.
 func (m *moduleEngine) OwnsGlobals() bool { return true }
 
+// MemoryGrown implements wasm.ModuleEngine.
+func (m *moduleEngine) MemoryGrown() {
+	m.putLocalMemory()
+}
+
+// putLocalMemory writes the local memory buffer pointer and length to the opaque buffer.
+func (m *moduleEngine) putLocalMemory() {
+	mem := m.module.MemoryInstance
+	offset := m.parent.offsets.LocalMemoryBegin
+
+	s := uint64(len(mem.Buffer))
+	var b uint64
+	if len(mem.Buffer) > 0 {
+		b = uint64(uintptr(unsafe.Pointer(&mem.Buffer[0])))
+	}
+	binary.LittleEndian.PutUint64(m.opaque[offset:], b)
+	binary.LittleEndian.PutUint64(m.opaque[offset+8:], s)
+}
+
 // ResolveImportedFunction implements wasm.ModuleEngine.
 func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	executableOffset, moduleCtxOffset, typeIDOffset := m.parent.offsets.ImportedFunctionOffset(index)

diff --git a/internal/integration_test/engine/adhoc_test.go b/internal/integration_test/engine/adhoc_test.go
@@ -1075,7 +1075,7 @@ func testModuleMemory(t *testing.T, r wazero.Runtime) {
 	bin := binaryencoding.EncodeModule(&wasm.Module{
 		TypeSection:     []wasm.FunctionType{{Params: []api.ValueType{api.ValueTypeI32}, ParamNumInUint64: 1}, {}},
 		FunctionSection: []wasm.Index{0, 1},
-		MemorySection:   &wasm.Memory{Min: 1, Cap: 1, Max: 2},
+		MemorySection:   &wasm.Memory{Min: 1, Cap: 1, Max: 20},
 		DataSection: []wasm.DataSegment{
 			{
 				Passive: true,
@@ -1109,6 +1109,10 @@ func testModuleMemory(t *testing.T, r wazero.Runtime) {
 
 	memory := inst.Memory()
 
+	// Ensures that memory grow by users is working.
+	_, ok := memory.Grow(10)
+	require.True(t, ok)
+
 	buf, ok := memory.Read(0, wasmPhraseSize)
 	require.True(t, ok)
 	require.Equal(t, make([]byte, wasmPhraseSize), buf)

diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
@@ -69,4 +69,7 @@ type ModuleEngine interface {
 	// FunctionInstanceReference returns Reference for the given Index for a FunctionInstance. The returned values are used by
 	// the initialization via ElementSegment.
 	FunctionInstanceReference(funcIndex Index) Reference
+
+	// MemoryGrown notifies the engine that the memory has grown.
+	MemoryGrown()
 }
diff --git a/internal/wasm/memory.go b/internal/wasm/memory.go
@@ -59,11 +59,14 @@ type MemoryInstance struct {
 	// with a fixed weight of 1 and no spurious notifications.
 	waiters sync.Map
 
+	// ownerModuleEngine is the module engine that owns this memory instance.
+	ownerModuleEngine ModuleEngine
+
 	expBuffer experimental.LinearMemory
 }
 
 // NewMemoryInstance creates a new instance based on the parameters in the SectionIDMemory.
-func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator) *MemoryInstance {
+func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator, moduleEngine ModuleEngine) *MemoryInstance {
 	minBytes := MemoryPagesToBytesNum(memSec.Min)
 	capBytes := MemoryPagesToBytesNum(memSec.Cap)
 	maxBytes := MemoryPagesToBytesNum(memSec.Max)
@@ -89,12 +92,13 @@ func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator) *
 		buffer = make([]byte, minBytes, capBytes)
 	}
 	return &MemoryInstance{
-		Buffer:    buffer,
-		Min:       memSec.Min,
-		Cap:       memoryBytesNumToPages(uint64(cap(buffer))),
-		Max:       memSec.Max,
-		Shared:    memSec.IsShared,
-		expBuffer: expBuffer,
+		Buffer:            buffer,
+		Min:               memSec.Min,
+		Cap:               memoryBytesNumToPages(uint64(cap(buffer))),
+		Max:               memSec.Max,
+		Shared:            memSec.IsShared,
+		expBuffer:         expBuffer,
+		ownerModuleEngine: moduleEngine,
 	}
 }
 
@@ -247,14 +251,12 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) {
 			m.Buffer = buffer
 			m.Cap = newPages
 		}
-		return currentPages, true
 	} else if newPages > m.Cap { // grow the memory.
 		if m.Shared {
 			panic("shared memory cannot be grown, this is a bug in wazero")
 		}
 		m.Buffer = append(m.Buffer, make([]byte, MemoryPagesToBytesNum(delta))...)
 		m.Cap = newPages
-		return currentPages, true
 	} else { // We already have the capacity we need.
 		if m.Shared {
 			// We assume grow is called under a guest lock.
@@ -264,8 +266,9 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) {
 		} else {
 			m.Buffer = m.Buffer[:MemoryPagesToBytesNum(newPages)]
 		}
-		return currentPages, true
 	}
+	m.ownerModuleEngine.MemoryGrown()
+	return currentPages, true
 }
 
 // Pages implements the same method as documented on api.Memory.

diff --git a/internal/wasm/memory_test.go b/internal/wasm/memory_test.go
@@ -48,6 +48,7 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			max := uint32(10)
 			maxBytes := MemoryPagesToBytesNum(max)
+			me := &mockModuleEngine{}
 			var m *MemoryInstance
 			switch {
 			default:
@@ -58,6 +59,7 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 				expBuffer := sliceAllocator(0, maxBytes)
 				m = &MemoryInstance{Max: max, Buffer: expBuffer.Reallocate(0), expBuffer: expBuffer}
 			}
+			m.ownerModuleEngine = me
 
 			res, ok := m.Grow(5)
 			require.True(t, ok)
@@ -94,6 +96,10 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 			// Ensure that the current page size equals the max.
 			require.Equal(t, max, m.Pages())
 
+			// Growing zero and beyond max won't notify the module engine.
+			// So in total, the memoryGrown should be called 3 times.
+			require.Equal(t, 3, me.memoryGrown)
+
 			if tc.capEqualsMax { // Ensure the capacity isn't more than max.
 				require.Equal(t, maxBytes, uint64(cap(m.Buffer)))
 			} else { // Slice doubles, so it should have a higher capacity than max.
@@ -851,12 +857,14 @@ func TestNewMemoryInstance_Shared(t *testing.T) {
 		},
 	}
 
+	me := &mockModuleEngine{}
 	for _, tc := range tests {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			m := NewMemoryInstance(tc.mem, nil)
+			m := NewMemoryInstance(tc.mem, nil, me)
 			require.Equal(t, tc.mem.Min, m.Min)
 			require.Equal(t, tc.mem.Max, m.Max)
+			require.Equal(t, me, m.ownerModuleEngine)
 			require.True(t, m.Shared)
 		})
 	}

diff --git a/internal/wasm/module.go b/internal/wasm/module.go
@@ -655,7 +655,7 @@ func paramNames(localNames IndirectNameMap, funcIdx uint32, paramLen int) []stri
 func (m *ModuleInstance) buildMemory(module *Module, allocator experimental.MemoryAllocator) {
 	memSec := module.MemorySection
 	if memSec != nil {
-		m.MemoryInstance = NewMemoryInstance(memSec, allocator)
+		m.MemoryInstance = NewMemoryInstance(memSec, allocator, m.Engine)
 		m.MemoryInstance.definition = &module.MemoryDefinitionSection[0]
 	}
 }

diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
@@ -414,6 +414,7 @@ type mockModuleEngine struct {
 	resolveImportsCalled map[Index]Index
 	importedMemModEngine ModuleEngine
 	lookupEntries        map[Index]mockModuleEngineLookupEntry
+	memoryGrown          int
 }
 
 type mockModuleEngineLookupEntry struct {
@@ -472,6 +473,9 @@ func (e *mockModuleEngine) SetGlobalValue(idx Index, lo, hi uint64) { panic("BUG
 // OwnsGlobals implements the same method as documented on wasm.ModuleEngine.
 func (e *mockModuleEngine) OwnsGlobals() bool { return false }
 
+// MemoryGrown implements the same method as documented on wasm.ModuleEngine.
+func (e *mockModuleEngine) MemoryGrown() { e.memoryGrown++ }
+
 // DoneInstantiation implements the same method as documented on wasm.ModuleEngine.
 func (e *mockModuleEngine) DoneInstantiation() {}