From 4dfb1f5b26fecab6ee14d3e2f750cef5e067695b Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Mon, 27 May 2024 10:37:03 +0900 Subject: [PATCH] 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 --- internal/engine/interpreter/interpreter.go | 3 ++ internal/engine/wazevo/call_engine.go | 18 ----------- internal/engine/wazevo/module_engine.go | 31 ++++++++++++------- .../integration_test/engine/adhoc_test.go | 6 +++- internal/wasm/engine.go | 3 ++ internal/wasm/memory.go | 23 ++++++++------ internal/wasm/memory_test.go | 10 +++++- internal/wasm/module.go | 2 +- internal/wasm/store_test.go | 4 +++ 9 files changed, 58 insertions(+), 42 deletions(-) diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go index a89ddc4573..18c5f4252d 100644 --- 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 index 3379c4ddef..72ce44e26f 100644 --- 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 index ba8f546c0d..efa1b9bbaa 100644 --- 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 index 770362b6ed..2f4f5e2623 100644 --- 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 index 58a4582178..61a342ef23 100644 --- 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 index 5cc5012dae..947b16112d 100644 --- 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 index 542591f56e..1d04014548 100644 --- 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 index 68573b918e..31b9b530f7 100644 --- 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 index 68cd3a390b..8df289306e 100644 --- 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() {}