wazevo: sets finalizers on executables (#1696)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/engine.go

@@ -27,6 +27,8 @@ type (
 		refToBinaryOffset map[ssa.FuncRef]int
 		// builtinFunctions is hods compiled builtin function trampolines.
 		builtinFunctions *builtinFunctions
+		// setFinalizer defaults to runtime.SetFinalizer, but overridable for tests.
+		setFinalizer func(obj interface{}, finalizer interface{})
 	}
 
 	builtinFunctions struct {
@@ -66,7 +68,10 @@ var _ wasm.Engine = (*engine)(nil)
 
 // NewEngine returns the implementation of wasm.Engine.
 func NewEngine(ctx context.Context, _ api.CoreFeatures, _ filecache.Cache) wasm.Engine {
-	e := &engine{compiledModules: make(map[wasm.ModuleID]*compiledModule), refToBinaryOffset: make(map[ssa.FuncRef]int)}
+	e := &engine{
+		compiledModules: make(map[wasm.ModuleID]*compiledModule), refToBinaryOffset: make(map[ssa.FuncRef]int),
+		setFinalizer: runtime.SetFinalizer,
+	}
 	e.compileBuiltinFunctions(ctx)
 	return e
 }
@@ -205,9 +210,7 @@ func (e *engine) compileModule(ctx context.Context, module *wasm.Module, listene
 	}
 	e.compiledModules[module.ID] = cm
 	cm.builtinFunctions = e.builtinFunctions
-
-	// TODO: finalizer.
-
+	e.setFinalizer(cm, compiledModuleFinalizer)
 	return cm, nil
 }
 
@@ -449,7 +452,28 @@ func (e *engine) compileBuiltinFunctions(ctx context.Context) {
 		e.builtinFunctions.stackGrowExecutable = mmapExecutable(src)
 	}
 
-	// TODO: finalizer.
+	e.setFinalizer(e.builtinFunctions, builtinFunctionFinalizer)
+}
+
+func builtinFunctionFinalizer(bf *builtinFunctions) {
+	if err := platform.MunmapCodeSegment(bf.memoryGrowExecutable); err != nil {
+		panic(err)
+	}
+	if err := platform.MunmapCodeSegment(bf.stackGrowExecutable); err != nil {
+		panic(err)
+	}
+
+	bf.memoryGrowExecutable = nil
+	bf.stackGrowExecutable = nil
+}
+
+func compiledModuleFinalizer(cm *compiledModule) {
+	if len(cm.executable) > 0 {
+		if err := platform.MunmapCodeSegment(cm.executable); err != nil {
+			panic(err)
+		}
+	}
+	cm.executable = nil
 }
 
 func mmapExecutable(src []byte) []byte {

internal/engine/wazevo/engine_test.go

@@ -0,0 +1,78 @@
+package wazevo
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/platform"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasm"
+)
+
+func Test_builtinFunctionFinalizer(t *testing.T) {
+	bf := &builtinFunctions{}
+
+	b1, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
+
+	b2, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
+	bf.memoryGrowExecutable = b1
+	bf.stackGrowExecutable = b2
+
+	builtinFunctionFinalizer(bf)
+	require.Nil(t, bf.memoryGrowExecutable)
+	require.Nil(t, bf.stackGrowExecutable)
+}
+
+func Test_compiledModuleFinalizer(t *testing.T) {
+	cm := &compiledModule{}
+
+	b, err := platform.MmapCodeSegment(100)
+	require.NoError(t, err)
+	cm.executable = b
+	compiledModuleFinalizer(cm)
+	require.Nil(t, cm.executable)
+}
+
+type fakeFinalizer map[*compiledModule]func(module *compiledModule)
+
+func (f fakeFinalizer) setFinalizer(obj interface{}, finalizer interface{}) {
+	cf := obj.(*compiledModule)
+	if _, ok := f[cf]; ok { // easier than adding a field for testing.T
+		panic(fmt.Sprintf("BUG: %v already had its finalizer set", cf))
+	}
+	f[cf] = finalizer.(func(*compiledModule))
+}
+
+func TestEngine_CompileModule(t *testing.T) {
+	ctx := context.Background()
+	e := NewEngine(ctx, 0, nil).(*engine)
+	ff := fakeFinalizer{}
+	e.setFinalizer = ff.setFinalizer
+
+	okModule := &wasm.Module{
+		TypeSection:     []wasm.FunctionType{{}},
+		FunctionSection: []wasm.Index{0, 0, 0, 0},
+		CodeSection: []wasm.Code{
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeEnd}},
+		},
+		ID: wasm.ModuleID{},
+	}
+
+	err := e.CompileModule(ctx, okModule, nil, false)
+	require.NoError(t, err)
+
+	// Compiling same module shouldn't be compiled again, but instead should be cached.
+	err = e.CompileModule(ctx, okModule, nil, false)
+	require.NoError(t, err)
+
+	// Pretend the finalizer executed, by invoking them one-by-one.
+	for k, v := range ff {
+		v(k)
+	}
+}