Repository: tetratelabs/wazero
Pull request #2315 (2024-09-14 18:06:18)

Allow custom memory allocator Reallocate to fail
Hey! 👋

This is just an idea, but the change was small enough that I decided to just try it.

How would we feel about allowing the custom allocator to return `nil` to signal it can't grow the linear memory?

In Go, a failure to allocate memory is a fatal condition, so it makes sense that we don't think about this when using `append` to grow the linear memory.

But custom allocators can handle this, it doesn't need to be a fatal condition. Also Wasm modules can deal with it: e.g. for `wasi-lib`, if `sbrk` fails, `malloc` just returns `NULL`, and (e.g.) SQLite can deal with that by freeing caches and trying again (if configured to do so).

cc: @anuraaga

mathetake: Sounds reasonable to me 
ncruces: OK. I'll add a test and document the behavior in the interface.

Patch:
diff --git a/experimental/memory.go b/experimental/memory.go
index e379bf05..26540648 100644
--- a/experimental/memory.go
+++ b/experimental/memory.go
@@ -35,6 +35,8 @@ type LinearMemory interface {
 	// Notes:
 	//   - To back a shared memory, Reallocate can't change the address of the
 	//     backing []byte (only its length/capacity may change).
+	//   - Reallocate may return nil if fails to grow the LinearMemory. This
+	//     condition may or may not be handled gracefully by the Wasm module.
 	Reallocate(size uint64) []byte
 	// Free the backing memory buffer.
 	Free()
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index ee0b453c..5b5e6e9d 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -487,7 +487,7 @@ func (e *engine) setLabelAddress(op *uint64, label label, labelAddressResolution
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (e *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	imported := importedModuleEngine.(*moduleEngine)
 	e.functions[index] = imported.functions[indexInImportedModule]
 }
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index efa1b9bb..8811feed 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -237,7 +237,7 @@ func (m *moduleEngine) putLocalMemory() {
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (m *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	executableOffset, moduleCtxOffset, typeIDOffset := m.parent.offsets.ImportedFunctionOffset(index)
 	importedME := importedModuleEngine.(*moduleEngine)
 
@@ -245,12 +245,12 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 		indexInImportedModule -= wasm.Index(len(importedME.importedFunctions))
 	} else {
 		imported := &importedME.importedFunctions[indexInImportedModule]
-		m.ResolveImportedFunction(index, imported.indexInModule, imported.me)
+		m.ResolveImportedFunction(index, descFunc, imported.indexInModule, imported.me)
 		return // Recursively resolve the imported function.
 	}
 
 	offset := importedME.parent.functionOffsets[indexInImportedModule]
-	typeID := getTypeIDOf(indexInImportedModule, importedME.module)
+	typeID := m.module.TypeIDs[descFunc]
 	executable := &importedME.parent.executable[offset]
 	// Write functionInstance.
 	binary.LittleEndian.PutUint64(m.opaque[executableOffset:], uint64(uintptr(unsafe.Pointer(executable))))
@@ -261,28 +261,6 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	m.importedFunctions[index] = importedFunction{me: importedME, indexInModule: indexInImportedModule}
 }
 
-func getTypeIDOf(funcIndex wasm.Index, m *wasm.ModuleInstance) wasm.FunctionTypeID {
-	source := m.Source
-
-	var typeIndex wasm.Index
-	if funcIndex >= source.ImportFunctionCount {
-		funcIndex -= source.ImportFunctionCount
-		typeIndex = source.FunctionSection[funcIndex]
-	} else {
-		var cnt wasm.Index
-		for i := range source.ImportSection {
-			if source.ImportSection[i].Type == wasm.ExternTypeFunc {
-				if cnt == funcIndex {
-					typeIndex = source.ImportSection[i].DescFunc
-					break
-				}
-				cnt++
-			}
-		}
-	}
-	return m.TypeIDs[typeIndex]
-}
-
 // ResolveImportedMemory implements wasm.ModuleEngine.
 func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEngine) {
 	importedME := importedModuleEngine.(*moduleEngine)
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index 61a342ef..8c387e9e 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -44,9 +44,10 @@ type ModuleEngine interface {
 
 	// ResolveImportedFunction is used to add imported functions needed to make this ModuleEngine fully functional.
 	// 	- `index` is the function Index of this imported function.
+	// 	- `descFunc` is the type Index in Module.TypeSection of this imported function. It corresponds to Import.DescFunc.
 	// 	- `indexInImportedModule` is the function Index of the imported function in the imported module.
 	//	- `importedModuleEngine` is the ModuleEngine for the imported ModuleInstance.
-	ResolveImportedFunction(index, indexInImportedModule Index, importedModuleEngine ModuleEngine)
+	ResolveImportedFunction(index, descFunc, indexInImportedModule Index, importedModuleEngine ModuleEngine)
 
 	// ResolveImportedMemory is called when this module imports a memory from another module.
 	ResolveImportedMemory(importedModuleEngine ModuleEngine)
diff --git a/internal/wasm/memory.go b/internal/wasm/memory.go
index 8e072fd1..9a1f0c6d 100644
--- a/internal/wasm/memory.go
+++ b/internal/wasm/memory.go
@@ -77,6 +77,7 @@ func NewMemoryInstance(memSec *Memory, allocator experimental.MemoryAllocator, m
 	if allocator != nil {
 		expBuffer = allocator.Allocate(capBytes, maxBytes)
 		buffer = expBuffer.Reallocate(minBytes)
+		_ = buffer[:minBytes] // Bounds check that the minimum was allocated.
 	} else if memSec.IsShared {
 		// Shared memory needs a fixed buffer, so allocate with the maximum size.
 		//
@@ -238,12 +239,15 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) {
 		return currentPages, true
 	}
 
-	// If exceeds the max of memory size, we push -1 according to the spec.
 	newPages := currentPages + delta
 	if newPages > m.Max || int32(delta) < 0 {
 		return 0, false
 	} else if m.expBuffer != nil {
 		buffer := m.expBuffer.Reallocate(MemoryPagesToBytesNum(newPages))
+		if buffer == nil {
+			// Allocator failed to grow.
+			return 0, false
+		}
 		if m.Shared {
 			if unsafe.SliceData(buffer) != unsafe.SliceData(m.Buffer) {
 				panic("shared memory cannot move, this is a bug in the memory allocator")
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index dda6e5b6..c7909c67 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -446,7 +446,7 @@ func (m *ModuleInstance) resolveImports(ctx context.Context, module *Module) (er
 					return
 				}
 
-				m.Engine.ResolveImportedFunction(i.IndexPerType, imported.Index, importedModule.Engine)
+				m.Engine.ResolveImportedFunction(i.IndexPerType, i.DescFunc, imported.Index, importedModule.Engine)
 			case ExternTypeTable:
 				expected := i.DescTable
 				importedTable := importedModule.Tables[imported.Index]


Test patch:
diff --git a/internal/engine/wazevo/module_engine_test.go b/internal/engine/wazevo/module_engine_test.go
index 45234e19..7e6bbf33 100644
--- a/internal/engine/wazevo/module_engine_test.go
+++ b/internal/engine/wazevo/module_engine_test.go
@@ -152,28 +152,19 @@ func TestModuleEngine_setupOpaque(t *testing.T) {
 }
 
 func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
-	const begin = 5000
-	m := &moduleEngine{
-		opaque:            make([]byte, 10000),
-		importedFunctions: make([]importedFunction, 4),
-		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
-			ImportedFunctionsBegin: begin,
-		}},
-	}
-
 	var op1, op2 byte = 0xaa, 0xbb
-	im1 := &moduleEngine{
+	importing := &moduleEngine{
 		opaquePtr: &op1,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
 			functionOffsets: []int{1, 5, 10},
 		},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 0, 111, 222, 333},
+			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 888, 111, 222, 333},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{4, 5, 6}},
 		},
 	}
-	im2 := &moduleEngine{
+	imported := &moduleEngine{
 		opaquePtr: &op2,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
@@ -185,10 +176,20 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, im1)
-	m.ResolveImportedFunction(1, 0, im2)
-	m.ResolveImportedFunction(2, 2, im1)
-	m.ResolveImportedFunction(3, 1, im1)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 4, 0, importing)
+	m.ResolveImportedFunction(1, 3, 0, imported)
+	m.ResolveImportedFunction(2, 6, 2, importing)
+	m.ResolveImportedFunction(3, 5, 1, importing)
 
 	for i, tc := range []struct {
 		index      int
@@ -196,10 +197,10 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{index: 0, op: &op1, executable: &im1.parent.executable[1], expTypeID: 111},
-		{index: 1, op: &op2, executable: &im2.parent.executable[50], expTypeID: 999},
-		{index: 2, op: &op1, executable: &im1.parent.executable[10], expTypeID: 333},
-		{index: 3, op: &op1, executable: &im1.parent.executable[5], expTypeID: 222},
+		{index: 0, op: &op1, executable: &importing.parent.executable[1], expTypeID: 111},
+		{index: 1, op: &op2, executable: &imported.parent.executable[50], expTypeID: 888},
+		{index: 2, op: &op1, executable: &importing.parent.executable[10], expTypeID: 333},
+		{index: 3, op: &op1, executable: &importing.parent.executable[5], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			buf := m.opaque[begin+wazevoapi.FunctionInstanceSize*tc.index:]
@@ -216,15 +217,6 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 }
 
 func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
-	const begin = 5000
-	m := &moduleEngine{
-		opaque:            make([]byte, 10000),
-		importedFunctions: make([]importedFunction, 4),
-		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
-			ImportedFunctionsBegin: begin,
-		}},
-	}
-
 	var importingOp, importedOp byte = 0xaa, 0xbb
 	imported := &moduleEngine{
 		opaquePtr: &importedOp,
@@ -245,13 +237,23 @@ func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
 		},
 		importedFunctions: []importedFunction{{me: imported, indexInModule: 0}},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 222, 0},
+			TypeIDs: []wasm.FunctionTypeID{999, 222, 0},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{1}},
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, importing)
-	m.ResolveImportedFunction(1, 1, importing)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 0, 0, importing)
+	m.ResolveImportedFunction(1, 1, 1, importing)
 
 	for i, tc := range []struct {
 		index      int
@@ -259,7 +261,7 @@ func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{index: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 111},
+		{index: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 999},
 		{index: 1, op: &importingOp, executable: &importing.parent.executable[500], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -312,26 +314,6 @@ func Test_functionInstance_offsets(t *testing.T) {
 	require.Equal(t, typeID, ptr+16)
 }
 
-func Test_getTypeIDOf(t *testing.T) {
-	m := &wasm.ModuleInstance{
-		TypeIDs: []wasm.FunctionTypeID{111, 222, 333, 444},
-		Source: &wasm.Module{
-			ImportFunctionCount: 1,
-			ImportSection: []wasm.Import{
-				{Type: wasm.ExternTypeMemory},
-				{Type: wasm.ExternTypeTable},
-				{Type: wasm.ExternTypeFunc, DescFunc: 3},
-			},
-			FunctionSection: []wasm.Index{2, 1, 0},
-		},
-	}
-
-	require.Equal(t, wasm.FunctionTypeID(444), getTypeIDOf(0, m))
-	require.Equal(t, wasm.FunctionTypeID(333), getTypeIDOf(1, m))
-	require.Equal(t, wasm.FunctionTypeID(222), getTypeIDOf(2, m))
-	require.Equal(t, wasm.FunctionTypeID(111), getTypeIDOf(3, m))
-}
-
 func Test_newAlignedOpaque(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		s := 16 * (i + 10)
diff --git a/internal/wasm/memory_test.go b/internal/wasm/memory_test.go
index d10c8276..16cb01b6 100644
--- a/internal/wasm/memory_test.go
+++ b/internal/wasm/memory_test.go
@@ -33,13 +33,15 @@ func TestMemoryBytesNumToPages(t *testing.T) {
 
 func TestMemoryInstance_Grow_Size(t *testing.T) {
 	tests := []struct {
-		name         string
-		capEqualsMax bool
-		expAllocator bool
+		name          string
+		capEqualsMax  bool
+		expAllocator  bool
+		failAllocator bool
 	}{
 		{name: ""},
 		{name: "capEqualsMax", capEqualsMax: true},
 		{name: "expAllocator", expAllocator: true},
+		{name: "failAllocator", failAllocator: true},
 	}
 
 	for _, tt := range tests {
@@ -58,6 +60,9 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 			case tc.expAllocator:
 				expBuffer := sliceAllocator(0, maxBytes)
 				m = &MemoryInstance{Max: max, Buffer: expBuffer.Reallocate(0), expBuffer: expBuffer}
+			case tc.failAllocator:
+				expBuffer := sliceAllocator(0, maxBytes)
+				m = &MemoryInstance{Max: max * 2, Buffer: expBuffer.Reallocate(0), expBuffer: expBuffer}
 			}
 			m.ownerModuleEngine = me
 
@@ -1023,6 +1028,9 @@ type sliceBuffer struct {
 func (b *sliceBuffer) Free() {}
 
 func (b *sliceBuffer) Reallocate(size uint64) []byte {
+	if size > b.max {
+		return nil
+	}
 	if cap := uint64(cap(b.buf)); size > cap {
 		b.buf = append(b.buf[:cap], make([]byte, size-cap)...)
 	} else {
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 08297591..a54ef182 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -341,7 +341,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				importedModuleName: {{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0}},
 				"non-exist":        {{Name: "fn", DescFunc: 0}},
 			},
-		}, importingModuleName, nil, nil)
+		}, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "module[non-exist] not instantiated")
 	})
 
@@ -485,7 +485,7 @@ func (e *mockModuleEngine) FunctionInstanceReference(i Index) Reference {
 }
 
 // ResolveImportedFunction implements the same method as documented on wasm.ModuleEngine.
-func (e *mockModuleEngine) ResolveImportedFunction(index, importedIndex Index, _ ModuleEngine) {
+func (e *mockModuleEngine) ResolveImportedFunction(index, _, importedIndex Index, _ ModuleEngine) {
 	e.resolveImportsCalled[index] = importedIndex
 }
 
@@ -742,7 +742,7 @@ func Test_resolveImports(t *testing.T) {
 				},
 			}
 
-			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module}
+			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module, TypeIDs: []FunctionTypeID{0, 1}}
 			err := m.resolveImports(context.Background(), module)
 			require.NoError(t, err)