Repository: tetratelabs/wazero
Pull request #505 (2022-04-25 07:04:29)

experimental: FunctionListenerFactory can decide whether to intercept function start/finish
For #451 (maybe fixes if this is enough functionality for now)

This allows registering a `FunctionListenerFactory` to a runtime. Anytime a module is instantiated, the factory is invoked while building the functions within the module. If the factory returns a listener, it is invoked anytime the defined function is called with `context`. The `context` can be used to propagate stack across function invocations in a stack.

We'll keep the hooks internal for now to let us bake within our own tests, so registration is defined using an internal interface implemented by `runtime`.

codefromthecrypt: I would recommend squashing and trying to contain changes as much as possible to new files. since a lot is in the same package, that should be possible. That way this low priority change doesn't have so much drift when we try to focus on things important to complete in the next 2 weeks
codefromthecrypt: This is great, @anuraaga!

As discussed offline, I squashed this draft into one commit then will add another to make it completely internal (ex no leak to wazero package) so we can iterate without effecting anyone.
codefromthecrypt: @anuraaga if you are keen, mark ready for review

Patch:
diff --git a/internal/experimental/README.md b/internal/experimental/README.md
new file mode 100644
index 00000000..c8d9a458
--- /dev/null
+++ b/internal/experimental/README.md
@@ -0,0 +1,4 @@
+# Experimental APIs
+
+wazero cares about compatability, and also recognizes that practice makes perfect. This area allows us to practice APIs
+without leaking them to exported non-internal interfaces. Notably, we sneak configuration through Go context values.
diff --git a/internal/experimental/api/listener.go b/internal/experimental/api/listener.go
new file mode 100644
index 00000000..3e8a041f
--- /dev/null
+++ b/internal/experimental/api/listener.go
@@ -0,0 +1,64 @@
+package experimentalapi
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/tetratelabs/wazero/api"
+)
+
+// FunctionListenerFactoryKey is a context.Context Value key. Its associated value should be a FunctionListenerFactory.
+type FunctionListenerFactoryKey struct{}
+
+// FunctionListenerFactory returns FunctionListeners to be notified when a function is called.
+type FunctionListenerFactory interface {
+	// NewListener returns a FunctionListener for a defined function. If nil is returned, no
+	// listener will be notified.
+	NewListener(info FunctionInfo) FunctionListener
+}
+
+// FunctionListener can be registered for any function via FunctionListenerFactory to
+// be notified when the function is called.
+type FunctionListener interface {
+	// Before is invoked before a function is called. ctx is the context of the caller function.
+	// The returned context will be used as the context of this function call. To add context
+	// information for this function call, add it to ctx and return the updated context. If
+	// no context information is needed, return ctx as is.
+	Before(ctx context.Context) context.Context
+
+	// After is invoked after a function is called. ctx is the context of this function call.
+	After(ctx context.Context)
+}
+
+// ValueInfo are information about the definition of a parameter or return value.
+type ValueInfo struct {
+	// The name of the value. Empty if name is not available.
+	Name string
+
+	// The type of the value.
+	Type api.ValueType
+}
+
+func (info ValueInfo) String() string {
+	n := info.Name
+	if len(n) == 0 {
+		n = "<unknown>"
+	}
+	return fmt.Sprintf("%v: %v", n, api.ValueTypeName(info.Type))
+}
+
+// FunctionInfo are information about a function available pre-instantiation.
+type FunctionInfo struct {
+	// The name of the module the function is defined in.
+	ModuleName string
+
+	// The name of the function. This will be the name of the export for an exported function.
+	// Empty if name is not available.
+	Name string
+
+	// The function parameters.
+	Params []ValueInfo
+
+	// The function return values.
+	Returns []ValueInfo
+}
diff --git a/internal/experimental/examples/function-listener/.gitignore b/internal/experimental/examples/function-listener/.gitignore
new file mode 100644
index 00000000..43c25e66
--- /dev/null
+++ b/internal/experimental/examples/function-listener/.gitignore
@@ -0,0 +1,1 @@
+function-listener
diff --git a/internal/experimental/examples/function-listener/README.md b/internal/experimental/examples/function-listener/README.md
new file mode 100644
index 00000000..6fe92b25
--- /dev/null
+++ b/internal/experimental/examples/function-listener/README.md
@@ -0,0 +1,13 @@
+## Function listener
+
+This example shows how to define a function listener to trace function calls. Function listeners are currently
+only implemented for interpreter mode.
+
+### Background
+
+As WebAssembly has become a target bytecode for many different languages and runtimes, we end up with binaries
+that encode the semantics and nuances for various frontends such as TinyGo and Rust, all in the single large
+binary. One line of code in a frontend language may result in working through many functions, possibly through
+host functions exposed with WASI. [print-trace.go](print-trace.go) shows how to use a FunctionListenerFactory to
+listen to all function invocations in the program. This can be used to find details about the execution of a
+wasm program, which can otherwise be a blackbox cobbled together by a frontend compiler.
diff --git a/internal/experimental/examples/function-listener/print-trace.go b/internal/experimental/examples/function-listener/print-trace.go
new file mode 100644
index 00000000..65a89190
--- /dev/null
+++ b/internal/experimental/examples/function-listener/print-trace.go
@@ -0,0 +1,136 @@
+package print_trace
+
+import (
+	"context"
+	_ "embed"
+	"fmt"
+	"log"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	experimentalapi "github.com/tetratelabs/wazero/internal/experimental/api"
+)
+
+type stackKey struct{}
+
+type callListener struct {
+	message string
+}
+
+func (l *callListener) Before(ctx context.Context) context.Context {
+	currStack, _ := ctx.Value(stackKey{}).([]string)
+	return context.WithValue(ctx, stackKey{}, append(currStack, l.message))
+}
+
+func (l *callListener) After(_ context.Context) {
+}
+
+type callListenerFactory struct {
+}
+
+func (f *callListenerFactory) NewListener(info experimentalapi.FunctionInfo) experimentalapi.FunctionListener {
+	return &callListener{
+		message: fmt.Sprintf("%v(%v) %v", info.Name, info.Params, info.Returns),
+	}
+}
+
+// main shows how to define, import and call a Go-defined function from a
+// WebAssembly-defined function.
+//
+// See README.md for a full description.
+func main() {
+	// FunctionListenerFactory is an experimental API, so the only way to configure it is via context key.
+	ctx := context.WithValue(context.Background(),
+		experimentalapi.FunctionListenerFactoryKey{}, &callListenerFactory{})
+
+	// Create a new WebAssembly Runtime.
+	// Note: This is interpreter-only for now!
+	r := wazero.NewRuntimeWithConfig(wazero.NewRuntimeConfigInterpreter())
+
+	// Instantiate a Go-defined module named "env" that exports functions to
+	// get the current year and log to the console.
+	//
+	// Note: As noted on ExportFunction documentation, function signatures are
+	// constrained to a subset of numeric types.
+	// Note: "env" is a module name conventionally used for arbitrary
+	// host-defined functions, but any name would do.
+	env, err := r.NewModuleBuilder("env").
+		ExportFunction("host1", host1).
+		ExportFunction("print_trace", print_trace).
+		Instantiate(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer env.Close(ctx)
+
+	// Instantiate a WebAssembly module named "listener" that imports
+	// functions defined in "env".
+	//
+	// Note: The import syntax in both Text and Binary format is the same
+	// regardless of if the function was defined in Go or WebAssembly.
+	listener, err := r.InstantiateModuleFromCode(ctx, []byte(`
+;; Define the optional module name. '$' prefixing is a part of the text format.
+(module $listener
+
+  ;; In WebAssembly, you don't import an entire module, rather each function.
+  ;; This imports the functions and gives them names which are easier to read
+  ;; than the alternative (zero-based index).
+  ;;
+  ;; Note: Importing unused functions is not an error in WebAssembly.
+  (import "env" "host1" (func $host1 (param i32)))
+  (import "env" "print_trace" (func $print_trace))
+
+  ;; wasm1 calls host1.
+  (func $wasm1 (param $val1 i32)
+                 ;; stack: []
+    local.get 0  ;; stack: [$value]
+    call $host1  ;; stack: []
+  )
+  ;; export allows api.Module to return this via ExportedFunction("wasm1")
+  (export "wasm1" (func $wasm1))
+
+  ;; wasm1 calls print_trace.
+  (func $wasm2 (param $val2 i32)
+                        ;; stack: []
+    call $print_trace   ;; stack: []
+  )
+  ;; export allows api.Module to return this via ExportedFunction("wasm2")
+  (export "wasm2" (func $wasm2))
+)`))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer listener.Close(ctx)
+
+	// First, try calling the "get_age" function and printing to the console externally.
+	_, err = listener.ExportedFunction("wasm1").Call(ctx, 100)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func host1(ctx context.Context, m api.Module, val uint32) {
+	hostonly(ctx, m, val)
+}
+
+// Wazero cannot intercept host->host calls as it is precompiled by Go. But since
+// ctx is propagated, such calls can still participate in the trace manually if
+// they want.
+func hostonly(ctx context.Context, m api.Module, val uint32) {
+	ctx = (&callListener{message: "hostonly"}).Before(ctx)
+	host2(ctx, m, val)
+}
+
+func host2(ctx context.Context, m api.Module, val uint32) {
+	_, err := m.ExportedFunction("wasm2").Call(ctx, uint64(val))
+	if err != nil {
+		log.Fatalf("Could not invoke wasm2: %v", err)
+	}
+}
+
+func print_trace(ctx context.Context) {
+	stack := ctx.Value(stackKey{}).([]string)
+	for _, f := range stack {
+		fmt.Println(f)
+	}
+}
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index 0899460a..fd5b0b1a 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -6,6 +6,7 @@ import (
 	"sort"
 	"strings"
 
+	experimental_api "github.com/tetratelabs/wazero/internal/experimental/api"
 	"github.com/tetratelabs/wazero/internal/wasmdebug"
 )
 
@@ -163,7 +164,10 @@ func (m *Module) maybeAddType(ft *FunctionType) Index {
 	return result
 }
 
-func (m *Module) buildHostFunctions(moduleName string) (functions []*FunctionInstance) {
+func (m *Module) buildHostFunctions(
+	moduleName string,
+	functionListenerFactory experimental_api.FunctionListenerFactory,
+) (functions []*FunctionInstance) {
 	// ModuleBuilder has no imports, which means the FunctionSection index is the same as the position in the function
 	// index namespace. Also, it ensures every function has a name. That's why there is less error checking here.
 	var functionNames = m.NameSection.FunctionNames
@@ -176,6 +180,25 @@ func (m *Module) buildHostFunctions(moduleName string) (functions []*FunctionIns
 			Index:  Index(idx),
 		}
 		f.DebugName = wasmdebug.FuncName(moduleName, functionNames[f.Index].Name, f.Index)
+		if functionListenerFactory != nil {
+			params := make([]experimental_api.ValueInfo, len(f.ParamTypes()))
+			for i, p := range f.ParamTypes() {
+				params[i].Type = p
+			}
+			results := make([]experimental_api.ValueInfo, len(f.ResultTypes()))
+			for i, r := range f.ResultTypes() {
+				results[i].Type = r
+			}
+
+			// TODO: add parameter names for host functions (vararg strings that must match arity with param length)
+
+			f.FunctionListener = functionListenerFactory.NewListener(experimental_api.FunctionInfo{
+				ModuleName: moduleName,
+				Name:       f.DebugName,
+				Params:     params,
+				Returns:    results,
+			})
+		}
 		functions = append(functions, f)
 	}
 	return
diff --git a/internal/wasm/interpreter/interpreter.go b/internal/wasm/interpreter/interpreter.go
index 45c9259d..56b4a25e 100644
--- a/internal/wasm/interpreter/interpreter.go
+++ b/internal/wasm/interpreter/interpreter.go
@@ -576,14 +576,23 @@ func (ce *callEngine) callGoFunc(ctx context.Context, callCtx *wasm.CallContext,
 		// Use the caller's memory, which might be different from the defining module on an imported function.
 		callCtx = callCtx.WithMemory(ce.frames[len(ce.frames)-1].f.source.Module.Memory)
 	}
+	if f.source.FunctionListener != nil {
+		ctx = f.source.FunctionListener.Before(ctx)
+	}
 	frame := &callFrame{f: f}
 	ce.pushFrame(frame)
 	results = wasm.CallGoFunc(ctx, callCtx, f.source, params)
 	ce.popFrame()
+	if f.source.FunctionListener != nil {
+		f.source.FunctionListener.After(ctx)
+	}
 	return
 }
 
 func (ce *callEngine) callNativeFunc(ctx context.Context, callCtx *wasm.CallContext, f *function) {
+	if f.source.FunctionListener != nil {
+		ctx = f.source.FunctionListener.Before(ctx)
+	}
 	frame := &callFrame{f: f}
 	moduleInst := f.source.Module
 	memoryInst := moduleInst.Memory
@@ -1674,6 +1683,9 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, callCtx *wasm.CallCont
 		}
 	}
 	ce.popFrame()
+	if f.source.FunctionListener != nil {
+		f.source.FunctionListener.After(ctx)
+	}
 }
 
 // popMemoryOffset takes a memory offset off the stack for use in load and store instructions.
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 02586c42..13922aac 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/tetratelabs/wazero/api"
+	experimental_api "github.com/tetratelabs/wazero/internal/experimental/api"
 	"github.com/tetratelabs/wazero/internal/ieee754"
 	"github.com/tetratelabs/wazero/internal/leb128"
 	"github.com/tetratelabs/wazero/internal/wasmdebug"
@@ -480,7 +481,7 @@ func (m *Module) buildGlobals(importedGlobals []*GlobalInstance) (globals []*Glo
 	return
 }
 
-func (m *Module) buildFunctions(moduleName string) (functions []*FunctionInstance) {
+func (m *Module) buildFunctions(moduleName string, functionListenerFactory experimental_api.FunctionListenerFactory) (functions []*FunctionInstance) {
 	var functionNames NameMap
 	if m.NameSection != nil {
 		functionNames = m.NameSection.FunctionNames
@@ -511,6 +512,10 @@ func (m *Module) buildFunctions(moduleName string) (functions []*FunctionInstanc
 			Index:      funcIdx,
 		}
 		f.DebugName = wasmdebug.FuncName(moduleName, funcName, funcIdx)
+		if functionListenerFactory != nil {
+			info := m.resolveFunction(moduleName, f, funcIdx)
+			f.FunctionListener = functionListenerFactory.NewListener(info)
+		}
 		functions = append(functions, f)
 	}
 	return
@@ -528,6 +533,39 @@ func (m *Module) buildMemory() (mem *MemoryInstance) {
 	return
 }
 
+func (m *Module) resolveFunction(moduleName string, f *FunctionInstance, funcIdx uint32) experimental_api.FunctionInfo {
+	params := make([]experimental_api.ValueInfo, len(f.ParamTypes()))
+	for i, p := range f.ParamTypes() {
+		params[i].Type = p
+	}
+	results := make([]experimental_api.ValueInfo, len(f.ResultTypes()))
+	for i, r := range f.ResultTypes() {
+		results[i].Type = r
+	}
+
+	for _, nm := range m.NameSection.LocalNames {
+		if nm.Index == funcIdx {
+			for _, n := range nm.NameMap {
+				if int(n.Index) < len(params) {
+					params[n.Index].Name = n.Name
+				} else {
+					resIdx := int(n.Index) - len(params)
+					// Only malformed program would have an index out of bounds here.
+					if resIdx < len(results) {
+						results[resIdx].Name = n.Name
+					}
+				}
+			}
+		}
+	}
+	return experimental_api.FunctionInfo{
+		ModuleName: moduleName,
+		Name:       f.DebugName,
+		Params:     params,
+		Returns:    results,
+	}
+}
+
 // Index is the offset in an index namespace, not necessarily an absolute position in a Module section. This is because
 // index namespaces are often preceded by a corresponding type in the Module.ImportSection.
 //
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 39ac5949..09b074cb 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -8,6 +8,7 @@ import (
 	"sync"
 
 	"github.com/tetratelabs/wazero/api"
+	experimentalapi "github.com/tetratelabs/wazero/internal/experimental/api"
 	"github.com/tetratelabs/wazero/internal/ieee754"
 	"github.com/tetratelabs/wazero/internal/leb128"
 )
@@ -121,6 +122,9 @@ type (
 
 		// Index holds the index of this function instance in the function index namespace (beginning with imports).
 		Index Index
+
+		// FunctionListener holds a listener to notify when this function is called.
+		FunctionListener experimentalapi.FunctionListener
 	}
 
 	// GlobalInstance represents a global instance in a store.
@@ -244,7 +248,13 @@ func NewStore(enabledFeatures Features, engine Engine) *Store {
 // * sys: the system context, which will be closed (SysContext.Close) on CallContext.Close.
 //
 // Note: Module.Validate must be called prior to instantiation.
-func (s *Store) Instantiate(ctx context.Context, module *Module, name string, sys *SysContext) (*CallContext, error) {
+func (s *Store) Instantiate(
+	ctx context.Context,
+	module *Module,
+	name string,
+	sys *SysContext,
+	functionListenerFactory experimentalapi.FunctionListenerFactory,
+) (*CallContext, error) {
 	if ctx == nil {
 		ctx = context.Background()
 	}
@@ -277,10 +287,10 @@ func (s *Store) Instantiate(ctx context.Context, module *Module, name string, sy
 	var funcSection SectionID
 	if module.HostFunctionSection == nil {
 		funcSection = SectionIDFunction
-		functions = module.buildFunctions(name)
+		functions = module.buildFunctions(name, functionListenerFactory)
 	} else {
 		funcSection = SectionIDHostFunction
-		functions = module.buildHostFunctions(name)
+		functions = module.buildHostFunctions(name, functionListenerFactory)
 	}
 
 	// Now we have all instances from imports and local ones, so ready to create a new ModuleInstance.
diff --git a/wasm.go b/wasm.go
index d583af9e..15978787 100644
--- a/wasm.go
+++ b/wasm.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 
 	"github.com/tetratelabs/wazero/api"
+	experimentalapi "github.com/tetratelabs/wazero/internal/experimental/api"
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wasm/binary"
 	"github.com/tetratelabs/wazero/internal/wasm/text"
@@ -224,7 +225,14 @@ func (r *runtime) InstantiateModuleWithConfig(ctx context.Context, compiled *Com
 
 	module := config.replaceImports(compiled.module)
 
-	mod, err = r.store.Instantiate(ctx, module, name, sysCtx)
+	var functionListenerFactory experimentalapi.FunctionListenerFactory
+	if ctx != nil { // Test to see if internal code are using an experimental feature.
+		if fnlf := ctx.Value(experimentalapi.FunctionListenerFactoryKey{}); fnlf != nil {
+			functionListenerFactory = fnlf.(experimentalapi.FunctionListenerFactory)
+		}
+	}
+
+	mod, err = r.store.Instantiate(ctx, module, name, sysCtx, functionListenerFactory)
 	if err != nil {
 		return
 	}


Test patch:
diff --git a/internal/experimental/examples/function-listener/print-trace_test.go b/internal/experimental/examples/function-listener/print-trace_test.go
new file mode 100644
index 00000000..a901fdea
--- /dev/null
+++ b/internal/experimental/examples/function-listener/print-trace_test.go
@@ -0,0 +1,21 @@
+package print_trace
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/maintester"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+// Test_main ensures the following will work:
+//
+//	go run print-trace.go
+func Test_main(t *testing.T) {
+	stdout, _ := maintester.TestMain(t, main)
+	require.Equal(t, `listener.wasm1([val1: i32]) []
+env.host1([<unknown>: i32]) []
+hostonly
+listener.wasm2([val2: i32]) []
+env.print_trace([]) []
+`, stdout)
+}
diff --git a/internal/integration_test/spectest/spec_test.go b/internal/integration_test/spectest/spec_test.go
index ddf30008..02ccd77c 100644
--- a/internal/integration_test/spectest/spec_test.go
+++ b/internal/integration_test/spectest/spec_test.go
@@ -273,7 +273,7 @@ func addSpectestModule(t *testing.T, store *wasm.Store) {
 	err = store.Engine.CompileModule(testCtx, mod)
 	require.NoError(t, err)
 
-	_, err = store.Instantiate(testCtx, mod, mod.NameSection.ModuleName, wasm.DefaultSysContext())
+	_, err = store.Instantiate(testCtx, mod, mod.NameSection.ModuleName, wasm.DefaultSysContext(), nil)
 	require.NoError(t, err)
 }
 
@@ -346,7 +346,7 @@ func runTest(t *testing.T, newEngine func(wasm.Features) wasm.Engine) {
 						require.NoError(t, err, msg)
 
 						moduleName = strings.TrimPrefix(moduleName, "$")
-						_, err = store.Instantiate(testCtx, mod, moduleName, nil)
+						_, err = store.Instantiate(testCtx, mod, moduleName, nil, nil)
 						lastInstantiatedModuleName = moduleName
 						require.NoError(t, err)
 					case "register":
@@ -476,7 +476,7 @@ func requireInstantiationError(t *testing.T, store *wasm.Store, buf []byte, msg
 		return
 	}
 
-	_, err = store.Instantiate(testCtx, mod, t.Name(), nil)
+	_, err = store.Instantiate(testCtx, mod, t.Name(), nil, nil)
 	require.Error(t, err, msg)
 }
 
diff --git a/internal/wasm/call_context_test.go b/internal/wasm/call_context_test.go
index f17c0f28..29640e0a 100644
--- a/internal/wasm/call_context_test.go
+++ b/internal/wasm/call_context_test.go
@@ -80,7 +80,7 @@ func TestCallContext_String(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			// Ensure paths that can create the host module can see the name.
-			m, err := s.Instantiate(context.Background(), &Module{}, tc.moduleName, nil)
+			m, err := s.Instantiate(context.Background(), &Module{}, tc.moduleName, nil, nil)
 			defer m.Close(testCtx) //nolint
 
 			require.NoError(t, err)
@@ -119,7 +119,7 @@ func TestCallContext_Close(t *testing.T) {
 		t.Run(fmt.Sprintf("%s calls store.CloseWithExitCode(module.name))", tc.name), func(t *testing.T) {
 			for _, ctx := range []context.Context{nil, testCtx} { // Ensure it doesn't crash on nil!
 				moduleName := t.Name()
-				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil)
+				m, err := s.Instantiate(ctx, &Module{}, moduleName, nil, nil)
 				require.NoError(t, err)
 
 				// We use side effects to see if Close called store.CloseWithExitCode (without repeating store_test.go).
@@ -160,7 +160,7 @@ func TestCallContext_Close(t *testing.T) {
 		require.NoError(t, err)
 
 		moduleName := t.Name()
-		m, err := s.Instantiate(context.Background(), &Module{}, moduleName, sys)
+		m, err := s.Instantiate(context.Background(), &Module{}, moduleName, sys, nil)
 		require.NoError(t, err)
 
 		// We use side effects to determine if Close in fact called SysContext.Close (without repeating sys_test.go).
diff --git a/internal/wasm/global_test.go b/internal/wasm/global_test.go
index 2a249f20..f56484dd 100644
--- a/internal/wasm/global_test.go
+++ b/internal/wasm/global_test.go
@@ -294,7 +294,7 @@ func TestPublicModule_Global(t *testing.T) {
 		s := newStore()
 		t.Run(tc.name, func(t *testing.T) {
 			// Instantiate the module and get the export of the above global
-			module, err := s.Instantiate(context.Background(), tc.module, t.Name(), nil)
+			module, err := s.Instantiate(context.Background(), tc.module, t.Name(), nil, nil)
 			require.NoError(t, err)
 
 			if global := module.ExportedGlobal("global"); tc.expected != nil {
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index b6353548..6c021767 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -262,7 +262,7 @@ func TestJIT_SliceAllocatedOnHeap(t *testing.T) {
 	err = store.Engine.CompileModule(testCtx, hm)
 	require.NoError(t, err)
 
-	_, err = store.Instantiate(testCtx, hm, hostModuleName, nil)
+	_, err = store.Instantiate(testCtx, hm, hostModuleName, nil, nil)
 	require.NoError(t, err)
 
 	const valueStackCorruption = "value_stack_corruption"
@@ -317,7 +317,7 @@ func TestJIT_SliceAllocatedOnHeap(t *testing.T) {
 	err = store.Engine.CompileModule(testCtx, m)
 	require.NoError(t, err)
 
-	mi, err := store.Instantiate(testCtx, m, t.Name(), nil)
+	mi, err := store.Instantiate(testCtx, m, t.Name(), nil, nil)
 	require.NoError(t, err)
 
 	for _, fnName := range []string{valueStackCorruption, callStackCorruption} {
diff --git a/internal/wasm/module_test.go b/internal/wasm/module_test.go
index 12bca650..6f60ec82 100644
--- a/internal/wasm/module_test.go
+++ b/internal/wasm/module_test.go
@@ -713,7 +713,7 @@ func TestModule_buildFunctionInstances(t *testing.T) {
 	}
 
 	// Note: This only returns module-defined functions, not imported ones. That's why the index starts with 1, not 0.
-	actual := m.buildFunctions("counter")
+	actual := m.buildFunctions("counter", nil)
 	expectedNames := []string{"counter.[1]", "counter.two", "counter.[3]", "counter.four", "counter.five"}
 	for i, f := range actual {
 		require.Equal(t, expectedNames[i], f.DebugName)
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 05773ce6..dee9a1dc 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -75,7 +75,7 @@ func TestModuleInstance_Memory(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			s := newStore()
 
-			instance, err := s.Instantiate(testCtx, tc.input, "test", nil)
+			instance, err := s.Instantiate(testCtx, tc.input, "test", nil, nil)
 			require.NoError(t, err)
 
 			mem := instance.ExportedMemory("memory")
@@ -100,7 +100,7 @@ func TestStore_Instantiate(t *testing.T) {
 	require.NoError(t, err)
 
 	sys := &SysContext{}
-	mod, err := s.Instantiate(testCtx, m, "", sys)
+	mod, err := s.Instantiate(testCtx, m, "", sys, nil)
 	require.NoError(t, err)
 	defer mod.Close(testCtx)
 
@@ -131,7 +131,7 @@ func TestStore_CloseModule(t *testing.T) {
 					Features20191205,
 				)
 				require.NoError(t, err)
-				_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+				_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 				require.NoError(t, err)
 			},
 		},
@@ -143,7 +143,7 @@ func TestStore_CloseModule(t *testing.T) {
 					FunctionSection: []uint32{0},
 					CodeSection:     []*Code{{Body: []byte{OpcodeEnd}}},
 					ExportSection:   []*Export{{Type: ExternTypeFunc, Index: 0, Name: "fn"}},
-				}, importedModuleName, nil)
+				}, importedModuleName, nil, nil)
 				require.NoError(t, err)
 			},
 		},
@@ -159,7 +159,7 @@ func TestStore_CloseModule(t *testing.T) {
 				MemorySection: &Memory{Min: 1},
 				GlobalSection: []*Global{{Type: &GlobalType{}, Init: &ConstantExpression{Opcode: OpcodeI32Const, Data: const1}}},
 				TableSection:  &Table{Min: 10},
-			}, importingModuleName, nil)
+			}, importingModuleName, nil, nil)
 			require.NoError(t, err)
 
 			imported, ok := s.modules[importedModuleName]
@@ -195,7 +195,7 @@ func TestStore_hammer(t *testing.T) {
 	require.NoError(t, err)
 
 	s := newStore()
-	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil)
+	imported, err := s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 	require.NoError(t, err)
 
 	_, ok := s.modules[imported.Name()]
@@ -222,7 +222,7 @@ func TestStore_hammer(t *testing.T) {
 		N = 100
 	}
 	hammer.NewHammer(t, P, N).Run(func(name string) {
-		mod, instantiateErr := s.Instantiate(testCtx, importingModule, name, DefaultSysContext())
+		mod, instantiateErr := s.Instantiate(testCtx, importingModule, name, DefaultSysContext(), nil)
 		require.NoError(t, instantiateErr)
 		require.NoError(t, mod.Close(testCtx))
 	}, nil)
@@ -252,17 +252,17 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 
 	t.Run("Fails if module name already in use", func(t *testing.T) {
 		s := newStore()
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 		require.NoError(t, err)
 
 		// Trying to register it again should fail
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 		require.EqualError(t, err, "module imported has already been instantiated")
 	})
 
 	t.Run("fail resolve import", func(t *testing.T) {
 		s := newStore()
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 		require.NoError(t, err)
 
 		hm := s.modules[importedModuleName]
@@ -276,14 +276,14 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				// But the second one tries to import uninitialized-module ->
 				{Type: ExternTypeFunc, Module: "non-exist", Name: "fn", DescFunc: 0},
 			},
-		}, importingModuleName, nil)
+		}, importingModuleName, nil, nil)
 		require.EqualError(t, err, "module[non-exist] not instantiated")
 	})
 
 	t.Run("compilation failed", func(t *testing.T) {
 		s := newStore()
 
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 		require.NoError(t, err)
 
 		hm := s.modules[importedModuleName]
@@ -302,7 +302,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 			ImportSection: []*Import{
 				{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 			},
-		}, importingModuleName, nil)
+		}, importingModuleName, nil, nil)
 		require.EqualError(t, err, "compilation failed: some compilation error")
 	})
 
@@ -311,7 +311,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 		engine := s.Engine.(*mockEngine)
 		engine.callFailIndex = 1
 
-		_, err = s.Instantiate(testCtx, m, importedModuleName, nil)
+		_, err = s.Instantiate(testCtx, m, importedModuleName, nil, nil)
 		require.NoError(t, err)
 
 		hm := s.modules[importedModuleName]
@@ -326,7 +326,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 			ImportSection: []*Import{
 				{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0},
 			},
-		}, importingModuleName, nil)
+		}, importingModuleName, nil, nil)
 		require.EqualError(t, err, "start function[1] failed: call failed")
 	})
 }
@@ -344,7 +344,7 @@ func TestCallContext_ExportedFunction(t *testing.T) {
 	s := newStore()
 
 	// Add the host module
-	imported, err := s.Instantiate(testCtx, host, host.NameSection.ModuleName, nil)
+	imported, err := s.Instantiate(testCtx, host, host.NameSection.ModuleName, nil, nil)
 	require.NoError(t, err)
 	defer imported.Close(testCtx)
 
@@ -354,7 +354,7 @@ func TestCallContext_ExportedFunction(t *testing.T) {
 			ImportSection: []*Import{{Type: ExternTypeFunc, Module: "host", Name: "host_fn", DescFunc: 0}},
 			MemorySection: &Memory{Min: 1},
 			ExportSection: []*Export{{Type: ExternTypeFunc, Name: "host.fn", Index: 0}},
-		}, "test", nil)
+		}, "test", nil, nil)
 		require.NoError(t, err)
 		defer importing.Close(testCtx)