experimental: configure custom memory allocator

experimental/memory.go

@@ -0,0 +1,31 @@
+package experimental
+
+import (
+	"context"
+
+	"github.com/tetratelabs/wazero/internal/ctxkey"
+)
+
+// MemoryAllocator is a memory allocation hook.
+type MemoryAllocator interface {
+	// Make is invoked to create a new memory, with the given specification.
+	// Implementations must return a []byte min bytes in length,
+	// should return a []byte with at least cap capacity,
+	// and be prepared to allocate up to max bytes of memory.
+	Make(min, cap, max uint64) []byte
+
+	// Grow is invoked to grow the memory to size bytes in length.
+	Grow(size uint64) []byte
+
+	// Free is invoked to free the memory.
+	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
+}

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))

internal/ctxkey/memory.go

@@ -0,0 +1,3 @@
+package ctxkey
+
+type MemoryAllocatorKey struct{}

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,11 +58,15 @@ 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
+
+	allocator experimental.MemoryAllocator
 }
 
 // NewMemoryInstance creates a new instance based on the parameters in the SectionIDMemory.
-func NewMemoryInstance(memSec *Memory) *MemoryInstance {
-	var size uint64
+func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator) *MemoryInstance {
+	min := MemoryPagesToBytesNum(memSec.Min)
+	cur := MemoryPagesToBytesNum(memSec.Cap)
+	max := MemoryPagesToBytesNum(memSec.Max)
 	if memSec.IsShared {
 		// Shared memory needs a fixed buffer, so allocate with the maximum size.
 		//
@@ -73,18 +78,22 @@ 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)
-	} else {
-		size = MemoryPagesToBytesNum(memSec.Cap)
+		cur = max
 	}
 
-	buffer := make([]byte, MemoryPagesToBytesNum(memSec.Min), size)
+	var buffer []byte
+	if allocator != nil {
+		buffer = allocator.Make(min, cur, max)
+	} else {
+		buffer = make([]byte, min, cur)
+	}
 	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,
+		allocator: allocator,
 	}
 }
 
@@ -222,6 +231,10 @@ 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.allocator != nil {
+		m.Buffer = m.allocator.Grow(MemoryPagesToBytesNum(newPages))
+		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")

internal/wasm/memory_test.go

@@ -832,7 +832,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)

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]
 	}
 }

internal/wasm/module_instance.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	"github.com/tetratelabs/wazero/sys"
 )
 
@@ -123,6 +125,11 @@ func (m *ModuleInstance) closeWithExitCode(ctx context.Context, exitCode uint32)
 	if !m.setExitCode(exitCode, exitCodeFlagResourceClosed) {
 		return nil // not an error to have already closed
 	}
+	if mem := m.MemoryInstance; mem != nil {
+		if allocator, ok := ctx.Value(ctxkey.MemoryAllocatorKey{}).(experimental.MemoryAllocator); ok {
+			allocator.Free()
+		}
+	}
 	return m.ensureResourcesClosed(ctx)
 }
 

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)

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 {