experimental: configure custom memory allocator (#2177)

Signed-off-by: Nuno Cruces <[email protected]>

Author: ncruces

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

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

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

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

@@ -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.

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 {