experimental: FunctionListenerFactory can decide whether to intercept function start/finish (#505)

Signed-off-by: Anuraag Agrawal <anuraaga@gmail.com>
Signed-off-by: Adrian Cole <adrian@tetrate.io>

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.

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

internal/experimental/examples/function-listener/.gitignore

@@ -0,0 +1 @@
+function-listener

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.

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

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

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

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

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 {

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

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.

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

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

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)

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.

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)
 

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
 	}