experimental: custom memory allocator API tweaks (#2186)

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

Author: ncruces

experimental/memory.go

@@ -6,21 +6,36 @@ import (
 	"github.com/tetratelabs/wazero/internal/expctxkeys"
 )
 
-// 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
+// MemoryAllocator is a memory allocation hook,
+// invoked to create a LinearMemory.
+type MemoryAllocator interface {
+	// Allocate should create a new LinearMemory with the given specification:
+	// cap is the suggested initial capacity for the backing []byte,
+	// and max the maximum length that will ever be requested.
+	//
+	// Notes:
+	//   - To back a shared memory, the address of the backing []byte cannot
+	//     change. This is checked at runtime. Implementations should document
+	//     if the returned LinearMemory meets this requirement.
+	Allocate(cap, max uint64) LinearMemory
+}
+
+// MemoryAllocatorFunc is a convenience for defining inlining a MemoryAllocator.
+type MemoryAllocatorFunc func(cap, max uint64) LinearMemory
+
+// Allocate implements MemoryAllocator.Allocate.
+func (f MemoryAllocatorFunc) Allocate(cap, max uint64) LinearMemory {
+	return f(cap, max)
+}
 
-// MemoryBuffer is a memory buffer that backs a Wasm memory.
-type MemoryBuffer interface {
-	// Buffer returns 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
+// LinearMemory is an expandable []byte that backs a Wasm linear memory.
+type LinearMemory interface {
+	// Reallocates the linear memory to size bytes in length.
+	//
+	// Notes:
+	//   - To back a shared memory, Reallocate can't change the address of the
+	//     backing []byte (only its length/capacity may change).
+	Reallocate(size uint64) []byte
 	// Free the backing memory buffer.
 	Free()
 }

internal/wasm/memory.go

@@ -59,7 +59,7 @@ type MemoryInstance struct {
 	// with a fixed weight of 1 and no spurious notifications.
 	waiters sync.Map
 
-	expBuffer experimental.MemoryBuffer
+	expBuffer experimental.LinearMemory
 }
 
 // NewMemoryInstance creates a new instance based on the parameters in the SectionIDMemory.
@@ -69,10 +69,10 @@ func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator) *
 	maxBytes := MemoryPagesToBytesNum(memSec.Max)
 
 	var buffer []byte
-	var expBuffer experimental.MemoryBuffer
+	var expBuffer experimental.LinearMemory
 	if allocator != nil {
-		expBuffer = allocator(minBytes, capBytes, maxBytes)
-		buffer = expBuffer.Buffer()
+		expBuffer = allocator.Allocate(capBytes, maxBytes)
+		buffer = expBuffer.Reallocate(minBytes)
 	} else if memSec.IsShared {
 		// Shared memory needs a fixed buffer, so allocate with the maximum size.
 		//
@@ -233,7 +233,7 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) {
 	if newPages > m.Max || int32(delta) < 0 {
 		return 0, false
 	} else if m.expBuffer != nil {
-		buffer := m.expBuffer.Grow(MemoryPagesToBytesNum(newPages))
+		buffer := m.expBuffer.Reallocate(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")

internal/wasm/memory_test.go

@@ -55,8 +55,8 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 			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}
+				expBuffer := sliceAllocator(0, maxBytes)
+				m = &MemoryInstance{Max: max, Buffer: expBuffer.Reallocate(0), expBuffer: expBuffer}
 			}
 
 			res, ok := m.Grow(5)
@@ -994,8 +994,8 @@ func requireChannelEmpty(t *testing.T, ch chan string) {
 	}
 }
 
-func sliceAllocator(min, cap, max uint64) experimental.MemoryBuffer {
-	return &sliceBuffer{make([]byte, min, cap), max}
+func sliceAllocator(cap, max uint64) experimental.LinearMemory {
+	return &sliceBuffer{make([]byte, cap), max}
 }
 
 type sliceBuffer struct {
@@ -1005,9 +1005,7 @@ type sliceBuffer struct {
 
 func (b *sliceBuffer) Free() {}
 
-func (b *sliceBuffer) Buffer() []byte { return b.buf }
-
-func (b *sliceBuffer) Grow(size uint64) []byte {
+func (b *sliceBuffer) Reallocate(size uint64) []byte {
 	if cap := uint64(cap(b.buf)); size > cap {
 		b.buf = append(b.buf[:cap], make([]byte, size-cap)...)
 	} else {