diff --git a/experimental/memory.go b/experimental/memory.go new file mode 100644 index 0000000000..1fb58f4320 --- /dev/null +++ b/experimental/memory.go @@ -0,0 +1,35 @@ +package experimental + +import ( + "context" + + "github.com/tetratelabs/wazero/internal/ctxkey" +) + +// MemoryAllocator is a memory allocation hook which is invoked +// to create a new MemoryBuffer, with the given specification: +// min is the initial and minimum length (in bytes) of the backing []byte, +// cap a suggested initial capacity, and max the maximum length +// that will ever be requested. +type MemoryAllocator func(min, cap, max uint64) MemoryBuffer + +// MemoryBuffer is a memory buffer that backs a Wasm memory. +type MemoryBuffer interface { + // Return the backing []byte for the memory buffer. + Buffer() []byte + // Grow the backing memory buffer to size bytes in length. + // To back a shared memory, Grow can't change the address + // of the backing []byte (only its length/capacity may change). + Grow(size uint64) []byte + // Free the backing memory buffer. + Free() +} + +// WithMemoryAllocator registers the given MemoryAllocator into the given +// context.Context. +func WithMemoryAllocator(ctx context.Context, allocator MemoryAllocator) context.Context { + if allocator != nil { + return context.WithValue(ctx, ctxkey.MemoryAllocatorKey{}, allocator) + } + return ctx +} diff --git a/imports/assemblyscript/assemblyscript_test.go b/imports/assemblyscript/assemblyscript_test.go index 705026a7dc..04bb74db1e 100644 --- a/imports/assemblyscript/assemblyscript_test.go +++ b/imports/assemblyscript/assemblyscript_test.go @@ -16,11 +16,11 @@ import ( "github.com/tetratelabs/wazero/api" . "github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/experimental/logging" + "github.com/tetratelabs/wazero/experimental/wazerotest" . "github.com/tetratelabs/wazero/internal/assemblyscript" "github.com/tetratelabs/wazero/internal/testing/proxy" "github.com/tetratelabs/wazero/internal/testing/require" "github.com/tetratelabs/wazero/internal/u64" - "github.com/tetratelabs/wazero/internal/wasm" "github.com/tetratelabs/wazero/sys" ) @@ -376,7 +376,7 @@ func Test_readAssemblyScriptString(t *testing.T) { tc := tt t.Run(tc.name, func(t *testing.T) { - mem := wasm.NewMemoryInstance(&wasm.Memory{Min: 1, Cap: 1, Max: 1}) + mem := wazerotest.NewFixedMemory(wazerotest.PageSize) tc.memory(mem) s, ok := readAssemblyScriptString(mem, uint32(tc.offset)) diff --git a/internal/ctxkey/memory.go b/internal/ctxkey/memory.go new file mode 100644 index 0000000000..bd07010988 --- /dev/null +++ b/internal/ctxkey/memory.go @@ -0,0 +1,3 @@ +package ctxkey + +type MemoryAllocatorKey struct{} diff --git a/internal/wasm/memory.go b/internal/wasm/memory.go index b390183b71..a1a34d4740 100644 --- a/internal/wasm/memory.go +++ b/internal/wasm/memory.go @@ -12,6 +12,7 @@ import ( "unsafe" "github.com/tetratelabs/wazero/api" + "github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/internal/internalapi" "github.com/tetratelabs/wazero/internal/wasmruntime" ) @@ -57,12 +58,22 @@ type MemoryInstance struct { // waiters implements atomic wait and notify. It is implemented similarly to golang.org/x/sync/semaphore, // with a fixed weight of 1 and no spurious notifications. waiters sync.Map + + expBuffer experimental.MemoryBuffer } // NewMemoryInstance creates a new instance based on the parameters in the SectionIDMemory. -func NewMemoryInstance(memSec *Memory) *MemoryInstance { - var size uint64 - if memSec.IsShared { +func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator) *MemoryInstance { + minBytes := MemoryPagesToBytesNum(memSec.Min) + capBytes := MemoryPagesToBytesNum(memSec.Cap) + maxBytes := MemoryPagesToBytesNum(memSec.Max) + + var buffer []byte + var expBuffer experimental.MemoryBuffer + if allocator != nil { + expBuffer = allocator(minBytes, capBytes, maxBytes) + buffer = expBuffer.Buffer() + } else if memSec.IsShared { // Shared memory needs a fixed buffer, so allocate with the maximum size. // // The rationale as to why we can simply use make([]byte) to a fixed buffer is that Go's GC is non-relocating. @@ -73,18 +84,17 @@ func NewMemoryInstance(memSec *Memory) *MemoryInstance { // the memory buffer allocation here is virtual and doesn't consume physical memory until it's used. // * https://github.com/golang/go/blob/8121604559035734c9677d5281bbdac8b1c17a1e/src/runtime/malloc.go#L1059 // * https://github.com/golang/go/blob/8121604559035734c9677d5281bbdac8b1c17a1e/src/runtime/malloc.go#L1165 - size = MemoryPagesToBytesNum(memSec.Max) + buffer = make([]byte, minBytes, maxBytes) } else { - size = MemoryPagesToBytesNum(memSec.Cap) + buffer = make([]byte, minBytes, capBytes) } - - buffer := make([]byte, MemoryPagesToBytesNum(memSec.Min), size) return &MemoryInstance{ - Buffer: buffer, - Min: memSec.Min, - Cap: memoryBytesNumToPages(uint64(cap(buffer))), - Max: memSec.Max, - Shared: memSec.IsShared, + Buffer: buffer, + Min: memSec.Min, + Cap: memoryBytesNumToPages(uint64(cap(buffer))), + Max: memSec.Max, + Shared: memSec.IsShared, + expBuffer: expBuffer, } } @@ -222,6 +232,22 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) { newPages := currentPages + delta if newPages > m.Max || int32(delta) < 0 { return 0, false + } else if m.expBuffer != nil { + buffer := m.expBuffer.Grow(MemoryPagesToBytesNum(newPages)) + if m.Shared { + if unsafe.SliceData(buffer) != unsafe.SliceData(m.Buffer) { + panic("shared memory cannot move, this is a bug in the memory allocator") + } + // We assume grow is called under a guest lock. + // But the memory length is accessed elsewhere, + // so use atomic to make the new length visible across threads. + atomicStoreLength(&m.Buffer, uintptr(len(buffer))) + m.Cap = memoryBytesNumToPages(uint64(cap(buffer))) + } else { + 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") @@ -231,9 +257,10 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) { return currentPages, true } else { // We already have the capacity we need. if m.Shared { - sp := (*reflect.SliceHeader)(unsafe.Pointer(&m.Buffer)) - // Use atomic write to ensure new length is visible across threads. - atomic.StoreUintptr((*uintptr)(unsafe.Pointer(&sp.Len)), uintptr(MemoryPagesToBytesNum(newPages))) + // We assume grow is called under a guest lock. + // But the memory length is accessed elsewhere, + // so use atomic to make the new length visible across threads. + atomicStoreLength(&m.Buffer, uintptr(MemoryPagesToBytesNum(newPages))) } else { m.Buffer = m.Buffer[:MemoryPagesToBytesNum(newPages)] } @@ -267,6 +294,13 @@ func PagesToUnitOfBytes(pages uint32) string { // Below are raw functions used to implement the api.Memory API: +// Uses atomic write to update the length of a slice. +func atomicStoreLength(slice *[]byte, length uintptr) { + slicePtr := (*reflect.SliceHeader)(unsafe.Pointer(slice)) + lenPtr := (*uintptr)(unsafe.Pointer(&slicePtr.Len)) + atomic.StoreUintptr(lenPtr, length) +} + // memoryBytesNumToPages converts the given number of bytes into the number of pages. func memoryBytesNumToPages(bytesNum uint64) (pages uint32) { return uint32(bytesNum >> MemoryPageSizeInBits) diff --git a/internal/wasm/memory_test.go b/internal/wasm/memory_test.go index 2aed432ff7..b08e8e4a89 100644 --- a/internal/wasm/memory_test.go +++ b/internal/wasm/memory_test.go @@ -9,6 +9,7 @@ import ( "unsafe" "github.com/tetratelabs/wazero/api" + "github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/internal/testing/require" ) @@ -34,9 +35,11 @@ func TestMemoryInstance_Grow_Size(t *testing.T) { tests := []struct { name string capEqualsMax bool + expAllocator bool }{ {name: ""}, {name: "capEqualsMax", capEqualsMax: true}, + {name: "expAllocator", expAllocator: true}, } for _, tt := range tests { @@ -46,10 +49,14 @@ func TestMemoryInstance_Grow_Size(t *testing.T) { max := uint32(10) maxBytes := MemoryPagesToBytesNum(max) var m *MemoryInstance - if tc.capEqualsMax { - m = &MemoryInstance{Cap: max, Max: max, Buffer: make([]byte, 0, maxBytes)} - } else { + switch { + default: m = &MemoryInstance{Max: max, Buffer: make([]byte, 0)} + case tc.capEqualsMax: + m = &MemoryInstance{Cap: max, Max: max, Buffer: make([]byte, 0, maxBytes)} + case tc.expAllocator: + expBuffer := sliceAllocator(0, 0, maxBytes) + m = &MemoryInstance{Max: max, Buffer: expBuffer.Buffer(), expBuffer: expBuffer} } res, ok := m.Grow(5) @@ -814,6 +821,13 @@ func BenchmarkWriteString(b *testing.B) { } } +func Test_atomicStoreLength(t *testing.T) { + // Doesn't verify atomicity, but at least we're updating the correct thing. + slice := make([]byte, 10, 20) + atomicStoreLength(&slice, 15) + require.Equal(t, 15, len(slice)) +} + func TestNewMemoryInstance_Shared(t *testing.T) { tests := []struct { name string @@ -832,7 +846,7 @@ func TestNewMemoryInstance_Shared(t *testing.T) { for _, tc := range tests { tc := tc t.Run(tc.name, func(t *testing.T) { - m := NewMemoryInstance(tc.mem) + m := NewMemoryInstance(tc.mem, nil) require.Equal(t, tc.mem.Min, m.Min) require.Equal(t, tc.mem.Max, m.Max) require.True(t, m.Shared) @@ -979,3 +993,25 @@ func requireChannelEmpty(t *testing.T, ch chan string) { // fallthrough } } + +func sliceAllocator(min, cap, max uint64) experimental.MemoryBuffer { + return &sliceBuffer{make([]byte, min, cap), max} +} + +type sliceBuffer struct { + buf []byte + max uint64 +} + +func (b *sliceBuffer) Free() {} + +func (b *sliceBuffer) Buffer() []byte { return b.buf } + +func (b *sliceBuffer) Grow(size uint64) []byte { + if cap := uint64(cap(b.buf)); size > cap { + b.buf = append(b.buf[:cap], make([]byte, size-cap)...) + } else { + b.buf = b.buf[:size] + } + return b.buf +} diff --git a/internal/wasm/module.go b/internal/wasm/module.go index 6ef68d508b..68573b918e 100644 --- a/internal/wasm/module.go +++ b/internal/wasm/module.go @@ -652,10 +652,10 @@ func paramNames(localNames IndirectNameMap, funcIdx uint32, paramLen int) []stri return nil } -func (m *ModuleInstance) buildMemory(module *Module) { +func (m *ModuleInstance) buildMemory(module *Module, allocator experimental.MemoryAllocator) { memSec := module.MemorySection if memSec != nil { - m.MemoryInstance = NewMemoryInstance(memSec) + m.MemoryInstance = NewMemoryInstance(memSec, allocator) m.MemoryInstance.definition = &module.MemoryDefinitionSection[0] } } diff --git a/internal/wasm/module_instance.go b/internal/wasm/module_instance.go index ea313e1348..20c733e6f7 100644 --- a/internal/wasm/module_instance.go +++ b/internal/wasm/module_instance.go @@ -151,20 +151,24 @@ func (m *ModuleInstance) ensureResourcesClosed(ctx context.Context) (err error) } if sysCtx := m.Sys; sysCtx != nil { // nil if from HostModuleBuilder - if err = sysCtx.FS().Close(); err != nil { - return err - } + err = sysCtx.FS().Close() m.Sys = nil } - if m.CodeCloser == nil { - return + if mem := m.MemoryInstance; mem != nil { + if mem.expBuffer != nil { + mem.expBuffer.Free() + mem.expBuffer = nil + } } - if e := m.CodeCloser.Close(ctx); e != nil && err == nil { - err = e + + if m.CodeCloser != nil { + if e := m.CodeCloser.Close(ctx); err == nil { + err = e + } + m.CodeCloser = nil } - m.CodeCloser = nil - return + return err } // Memory implements the same method as documented on api.Module. diff --git a/internal/wasm/module_test.go b/internal/wasm/module_test.go index 6d8efe26ac..9c860d4388 100644 --- a/internal/wasm/module_test.go +++ b/internal/wasm/module_test.go @@ -839,7 +839,7 @@ func TestModule_buildGlobals(t *testing.T) { func TestModule_buildMemoryInstance(t *testing.T) { t.Run("nil", func(t *testing.T) { m := ModuleInstance{} - m.buildMemory(&Module{}) + m.buildMemory(&Module{}, nil) require.Nil(t, m.MemoryInstance) }) t.Run("non-nil", func(t *testing.T) { @@ -850,7 +850,7 @@ func TestModule_buildMemoryInstance(t *testing.T) { m.buildMemory(&Module{ MemorySection: &Memory{Min: min, Cap: min, Max: max}, MemoryDefinitionSection: []MemoryDefinition{mDef}, - }) + }, nil) mem := m.MemoryInstance require.Equal(t, min, mem.Min) require.Equal(t, max, mem.Max) diff --git a/internal/wasm/store.go b/internal/wasm/store.go index 5ec672fe34..ece30aa2e8 100644 --- a/internal/wasm/store.go +++ b/internal/wasm/store.go @@ -8,6 +8,7 @@ import ( "sync/atomic" "github.com/tetratelabs/wazero/api" + "github.com/tetratelabs/wazero/experimental" "github.com/tetratelabs/wazero/internal/ctxkey" "github.com/tetratelabs/wazero/internal/internalapi" "github.com/tetratelabs/wazero/internal/leb128" @@ -362,8 +363,13 @@ func (s *Store) instantiate( return nil, err } + var allocator experimental.MemoryAllocator + if ctx != nil { + allocator, _ = ctx.Value(ctxkey.MemoryAllocatorKey{}).(experimental.MemoryAllocator) + } + m.buildGlobals(module, m.Engine.FunctionInstanceReference) - m.buildMemory(module) + m.buildMemory(module, allocator) m.Exports = module.Exports for _, exp := range m.Exports { if exp.Type == ExternTypeTable {