Repository: tetratelabs/wazero
Pull request #2175 (2024-04-06 00:49:01)

experimental: cleanup context keys
We follow 2 different models in wazero for context keys:
1. one where we make the key type public and the user is supposed to do `context.WithValue(ctx, key, value)`
2. another where we expose `experimental.WithBlargh(ctx, blargh)`

The second doesn't expose the key, and is more typesafe because it ensures `blargh` is of the expected type.

This PR refactors to use *2* everywhere.
Old code left deprecated, but we can remove (since it's experimental).


Patch:
diff --git a/cmd/wazero/wazero.go b/cmd/wazero/wazero.go
index d4667049..0c64bc90 100644
--- a/cmd/wazero/wazero.go
+++ b/cmd/wazero/wazero.go
@@ -451,7 +451,7 @@ func detectImports(imports []api.FunctionDefinition) importMode {
 
 func maybeHostLogging(ctx context.Context, scopes logging.LogScopes, stdErr logging.Writer) context.Context {
 	if scopes != 0 {
-		return context.WithValue(ctx, experimental.FunctionListenerFactoryKey{}, logging.NewHostLoggingListenerFactory(stdErr, scopes))
+		return experimental.WithFunctionListenerFactory(ctx, logging.NewHostLoggingListenerFactory(stdErr, scopes))
 	}
 	return ctx
 }
diff --git a/experimental/checkpoint.go b/experimental/checkpoint.go
index 1560e6ba..eb423305 100644
--- a/experimental/checkpoint.go
+++ b/experimental/checkpoint.go
@@ -1,5 +1,11 @@
 package experimental
 
+import (
+	"context"
+
+	"github.com/tetratelabs/wazero/internal/ctxkey"
+)
+
 // Snapshot holds the execution state at the time of a Snapshotter.Snapshot call.
 type Snapshot interface {
 	// Restore sets the Wasm execution state to the capture. Because a host function
@@ -18,8 +24,25 @@ type Snapshotter interface {
 // EnableSnapshotterKey is a context key to indicate that snapshotting should be enabled.
 // The context.Context passed to a exported function invocation should have this key set
 // to a non-nil value, and host functions will be able to retrieve it using SnapshotterKey.
-type EnableSnapshotterKey struct{}
+//
+// Deprecated: use WithSnapshotter to enable snapshots.
+type EnableSnapshotterKey = ctxkey.EnableSnapshotterKey
+
+// WithSnapshotter enables snapshots.
+// Passing the returned context to a exported function invocation enables snapshots,
+// and allows host functions to retrieve the Snapshotter using SnapshotterKey.
+func WithSnapshotter(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ctxkey.EnableSnapshotterKey{}, struct{}{})
+}
 
 // SnapshotterKey is a context key to access a Snapshotter from a host function.
 // It is only present if EnableSnapshotter was set in the function invocation context.
-type SnapshotterKey struct{}
+//
+// Deprecated: use GetSnapshotter to get the snapshotter.
+type SnapshotterKey = ctxkey.SnapshotterKey
+
+// GetSnapshotter gets the Snapshotter from a host function.
+// It is only present if WithSnapshotter was called with the function invocation context.
+func GetSnapshotter(ctx context.Context) Snapshotter {
+	return ctx.Value(ctxkey.SnapshotterKey{}).(Snapshotter)
+}
diff --git a/experimental/close.go b/experimental/close.go
index 1c37e8c0..0e654a61 100644
--- a/experimental/close.go
+++ b/experimental/close.go
@@ -3,7 +3,7 @@ package experimental
 import (
 	"context"
 
-	"github.com/tetratelabs/wazero/internal/close"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 )
 
 // CloseNotifier is a notification hook, invoked when a module is closed.
@@ -57,7 +57,7 @@ func (f CloseNotifyFunc) CloseNotify(ctx context.Context, exitCode uint32) {
 // context.Context.
 func WithCloseNotifier(ctx context.Context, notifier CloseNotifier) context.Context {
 	if notifier != nil {
-		return context.WithValue(ctx, close.NotifierKey{}, notifier)
+		return context.WithValue(ctx, ctxkey.CloseNotifierKey{}, notifier)
 	}
 	return ctx
 }
diff --git a/experimental/listener.go b/experimental/listener.go
index 64d98e90..d5a6274b 100644
--- a/experimental/listener.go
+++ b/experimental/listener.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 )
 
 // StackIterator allows iterating on each function of the call stack, starting
@@ -23,10 +24,17 @@ type StackIterator interface {
 	ProgramCounter() ProgramCounter
 }
 
-// FunctionListenerFactoryKey is a context.Context Value key. Its associated value should be a FunctionListenerFactory.
+// FunctionListenerFactoryKey is a context.Context Value key.
+// Its associated value should be a FunctionListenerFactory.
 //
-// See https://github.com/tetratelabs/wazero/issues/451
-type FunctionListenerFactoryKey struct{}
+// Deprecated: use WithFunctionListenerFactory to enable snapshots.
+type FunctionListenerFactoryKey = ctxkey.FunctionListenerFactoryKey
+
+// WithFunctionListenerFactory registers a FunctionListenerFactory
+// with the context.
+func WithFunctionListenerFactory(ctx context.Context, factory FunctionListenerFactory) context.Context {
+	return context.WithValue(ctx, ctxkey.FunctionListenerFactoryKey{}, factory)
+}
 
 // FunctionListenerFactory returns FunctionListeners to be notified when a
 // function is called.
diff --git a/internal/ctxkey/checkpoint.go b/internal/ctxkey/checkpoint.go
new file mode 100644
index 00000000..85672aa6
--- /dev/null
+++ b/internal/ctxkey/checkpoint.go
@@ -0,0 +1,10 @@
+package ctxkey
+
+// EnableSnapshotterKey is a context key to indicate that snapshotting should be enabled.
+// The context.Context passed to a exported function invocation should have this key set
+// to a non-nil value, and host functions will be able to retrieve it using SnapshotterKey.
+type EnableSnapshotterKey struct{}
+
+// SnapshotterKey is a context key to access a Snapshotter from a host function.
+// It is only present if EnableSnapshotter was set in the function invocation context.
+type SnapshotterKey struct{}
diff --git a/internal/close/close.go b/internal/ctxkey/close.go
similarity index 60%
rename from internal/close/close.go
rename to internal/ctxkey/close.go
index 33b39b1e..cc0c4726 100644
--- a/internal/close/close.go
+++ b/internal/ctxkey/close.go
@@ -1,12 +1,12 @@
 // Package close allows experimental.CloseNotifier without introducing a
 // package cycle.
-package close
+package ctxkey
 
 import "context"
 
-// NotifierKey is a context.Context Value key. Its associated value should be a
+// CloseNotifierKey is a context.Context Value key. Its associated value should be a
 // Notifier.
-type NotifierKey struct{}
+type CloseNotifierKey struct{}
 
 type Notifier interface {
 	CloseNotify(ctx context.Context, exitCode uint32)
diff --git a/internal/ctxkey/listener.go b/internal/ctxkey/listener.go
new file mode 100644
index 00000000..209741f8
--- /dev/null
+++ b/internal/ctxkey/listener.go
@@ -0,0 +1,7 @@
+package ctxkey
+
+// FunctionListenerFactoryKey is a context.Context Value key.
+// Its associated value should be a FunctionListenerFactory.
+//
+// See https://github.com/tetratelabs/wazero/issues/451
+type FunctionListenerFactoryKey struct{}
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index 87280754..695bed95 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	"github.com/tetratelabs/wazero/internal/filecache"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/moremath"
@@ -572,8 +573,8 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u
 		}
 	}
 
-	if ctx.Value(experimental.EnableSnapshotterKey{}) != nil {
-		ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, ce)
+	if ctx.Value(ctxkey.EnableSnapshotterKey{}) != nil {
+		ctx = context.WithValue(ctx, ctxkey.SnapshotterKey{}, ce)
 	}
 
 	defer func() {
@@ -743,7 +744,7 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance
 			frame.pc = op.Us[v]
 		case wazeroir.OperationKindCall:
 			func() {
-				if ctx.Value(experimental.EnableSnapshotterKey{}) != nil {
+				if ctx.Value(ctxkey.EnableSnapshotterKey{}) != nil {
 					defer func() {
 						if r := recover(); r != nil {
 							if s, ok := r.(*snapshot); ok && s.ce == ce {
diff --git a/internal/engine/wazevo/call_engine.go b/internal/engine/wazevo/call_engine.go
index e5a41985..ed1bb467 100644
--- a/internal/engine/wazevo/call_engine.go
+++ b/internal/engine/wazevo/call_engine.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	"github.com/tetratelabs/wazero/internal/engine/wazevo/wazevoapi"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -193,9 +194,9 @@ func (c *callEngine) CallWithStack(ctx context.Context, paramResultStack []uint6
 
 // CallWithStack implements api.Function.
 func (c *callEngine) callWithStack(ctx context.Context, paramResultStack []uint64) (err error) {
-	snapshotEnabled := ctx.Value(experimental.EnableSnapshotterKey{}) != nil
+	snapshotEnabled := ctx.Value(ctxkey.EnableSnapshotterKey{}) != nil
 	if snapshotEnabled {
-		ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, c)
+		ctx = context.WithValue(ctx, ctxkey.SnapshotterKey{}, c)
 	}
 
 	if wazevoapi.StackGuardCheckEnabled {
diff --git a/internal/logging/logging.go b/internal/logging/logging.go
index ecec61f7..577efff1 100644
--- a/internal/logging/logging.go
+++ b/internal/logging/logging.go
@@ -94,9 +94,6 @@ func (f LogScopes) String() string {
 	return builder.String()
 }
 
-// LoggerKey is a context.Context Value key with a FunctionLogger value.
-type LoggerKey struct{}
-
 type ParamLogger func(ctx context.Context, mod api.Module, w Writer, params []uint64)
 
 type ParamSampler func(ctx context.Context, mod api.Module, params []uint64) bool
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index bf0fa364..5ec672fe 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -8,7 +8,7 @@ import (
 	"sync/atomic"
 
 	"github.com/tetratelabs/wazero/api"
-	"github.com/tetratelabs/wazero/internal/close"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/leb128"
 	internalsys "github.com/tetratelabs/wazero/internal/sys"
@@ -124,7 +124,7 @@ type (
 		Source *Module
 
 		// CloseNotifier is an experimental hook called once on close.
-		CloseNotifier close.Notifier
+		CloseNotifier ctxkey.Notifier
 	}
 
 	// DataInstance holds bytes corresponding to the data segment in a module.
diff --git a/runtime.go b/runtime.go
index 220324ed..021c47f9 100644
--- a/runtime.go
+++ b/runtime.go
@@ -7,7 +7,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	experimentalapi "github.com/tetratelabs/wazero/experimental"
-	internalclose "github.com/tetratelabs/wazero/internal/close"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	internalsock "github.com/tetratelabs/wazero/internal/sock"
 	internalsys "github.com/tetratelabs/wazero/internal/sys"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -242,7 +242,7 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 
 func buildFunctionListeners(ctx context.Context, internal *wasm.Module) ([]experimentalapi.FunctionListener, error) {
 	// Test to see if internal code are using an experimental feature.
-	fnlf := ctx.Value(experimentalapi.FunctionListenerFactoryKey{})
+	fnlf := ctx.Value(ctxkey.FunctionListenerFactoryKey{})
 	if fnlf == nil {
 		return nil, nil
 	}
@@ -318,7 +318,7 @@ func (r *runtime) InstantiateModule(
 		return
 	}
 
-	if closeNotifier, ok := ctx.Value(internalclose.NotifierKey{}).(internalclose.Notifier); ok {
+	if closeNotifier, ok := ctx.Value(ctxkey.CloseNotifierKey{}).(ctxkey.Notifier); ok {
 		mod.(*wasm.ModuleInstance).CloseNotifier = closeNotifier
 	}
 


Test patch:
diff --git a/experimental/checkpoint_example_test.go b/experimental/checkpoint_example_test.go
index d995befe..7535535b 100644
--- a/experimental/checkpoint_example_test.go
+++ b/experimental/checkpoint_example_test.go
@@ -28,7 +28,7 @@ func Example_enableSnapshotterKey() {
 
 	// Enable experimental snapshotting functionality by setting it to context. We use this
 	// context when invoking functions, indicating to wazero to enable it.
-	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+	ctx = experimental.WithSnapshotter(ctx)
 
 	// Also place a mutable holder of snapshots to be referenced during restore.
 	var snapshots []experimental.Snapshot
@@ -39,8 +39,8 @@ func Example_enableSnapshotterKey() {
 	_, err := rt.NewHostModuleBuilder("example").
 		NewFunctionBuilder().
 		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
-			// Because we set EnableSnapshotterKey to context, this is non-nil.
-			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
+			// Because we enabled snapshots with WithSnapshotter, this is non-nil.
+			snapshot := experimental.GetSnapshotter(ctx).Snapshot()
 
 			// Get our mutable snapshots holder to be able to add to it. Our example only calls snapshot
 			// and restore once but real programs will often call them at multiple layers within a call
diff --git a/experimental/checkpoint_test.go b/experimental/checkpoint_test.go
index 3abf921f..ce5b57d6 100644
--- a/experimental/checkpoint_test.go
+++ b/experimental/checkpoint_test.go
@@ -24,7 +24,7 @@ func TestSnapshotNestedWasmInvocation(t *testing.T) {
 			defer func() {
 				sidechannel = 10
 			}()
-			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
+			snapshot := experimental.GetSnapshotter(ctx).Snapshot()
 			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
 			idx := len(*snapshots)
 			*snapshots = append(*snapshots, snapshot)
@@ -55,7 +55,7 @@ func TestSnapshotNestedWasmInvocation(t *testing.T) {
 
 	var snapshots []experimental.Snapshot
 	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)
-	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+	ctx = experimental.WithSnapshotter(ctx)
 
 	snapshotPtr := uint64(0)
 	res, err := mod.ExportedFunction("snapshot").Call(ctx, snapshotPtr)
@@ -75,7 +75,7 @@ func TestSnapshotMultipleWasmInvocations(t *testing.T) {
 	_, err := rt.NewHostModuleBuilder("example").
 		NewFunctionBuilder().
 		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
-			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
+			snapshot := experimental.GetSnapshotter(ctx).Snapshot()
 			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
 			idx := len(*snapshots)
 			*snapshots = append(*snapshots, snapshot)
@@ -103,7 +103,7 @@ func TestSnapshotMultipleWasmInvocations(t *testing.T) {
 
 	var snapshots []experimental.Snapshot
 	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)
-	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+	ctx = experimental.WithSnapshotter(ctx)
 
 	snapshotPtr := uint64(0)
 	res, err := mod.ExportedFunction("snapshot").Call(ctx, snapshotPtr)
diff --git a/experimental/close_test.go b/experimental/close_test.go
index 1fe9d8e3..814c444c 100644
--- a/experimental/close_test.go
+++ b/experimental/close_test.go
@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/experimental"
-	"github.com/tetratelabs/wazero/internal/close"
+	"github.com/tetratelabs/wazero/internal/ctxkey"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
@@ -33,7 +33,7 @@ func TestWithCloseNotifier(t *testing.T) {
 		tc := tt
 		t.Run(tc.name, func(t *testing.T) {
 			if decorated := experimental.WithCloseNotifier(testCtx, tc.notification); tc.expected {
-				require.NotNil(t, decorated.Value(close.NotifierKey{}))
+				require.NotNil(t, decorated.Value(ctxkey.CloseNotifierKey{}))
 			} else {
 				require.Same(t, testCtx, decorated)
 			}
diff --git a/experimental/listener_example_test.go b/experimental/listener_example_test.go
index 0d2ece41..28faae4f 100644
--- a/experimental/listener_example_test.go
+++ b/experimental/listener_example_test.go
@@ -59,7 +59,7 @@ func Example_customListenerFactory() {
 	u := uniqGoFuncs{}
 
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, u)
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), u)
 
 	r := wazero.NewRuntime(ctx)
 	defer r.Close(ctx) // This closes everything this Runtime created.
diff --git a/experimental/listener_test.go b/experimental/listener_test.go
index c634451b..f65fb8dc 100644
--- a/experimental/listener_test.go
+++ b/experimental/listener_test.go
@@ -44,7 +44,7 @@ func (r *recorder) NewFunctionListener(definition api.FunctionDefinition) experi
 func TestFunctionListenerFactory(t *testing.T) {
 	// Set context to one that has an experimental listener
 	factory := &recorder{m: map[string]struct{}{}}
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, factory)
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), factory)
 
 	// Define a module with two functions
 	bin := binaryencoding.EncodeModule(&wasm.Module{
diff --git a/experimental/logging/log_listener_example_test.go b/experimental/logging/log_listener_example_test.go
index 5109c15e..f1662c33 100644
--- a/experimental/logging/log_listener_example_test.go
+++ b/experimental/logging/log_listener_example_test.go
@@ -23,7 +23,7 @@ var listenerWasm []byte
 // it is configured.
 func Example_newHostLoggingListenerFactory() {
 	// Set context to one that has an experimental listener that logs all host functions.
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{},
+	ctx := experimental.WithFunctionListenerFactory(context.Background(),
 		logging.NewHostLoggingListenerFactory(os.Stdout, logging.LogScopeAll))
 
 	r := wazero.NewRuntime(ctx)
@@ -57,7 +57,7 @@ func Example_newHostLoggingListenerFactory() {
 // functions.
 func Example_newLoggingListenerFactory() {
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(os.Stdout))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(os.Stdout))
 
 	r := wazero.NewRuntime(ctx)
 	defer r.Close(ctx) // This closes everything this Runtime created.
diff --git a/imports/assemblyscript/assemblyscript_test.go b/imports/assemblyscript/assemblyscript_test.go
index f72a7aa7..705026a7 100644
--- a/imports/assemblyscript/assemblyscript_test.go
+++ b/imports/assemblyscript/assemblyscript_test.go
@@ -426,7 +426,7 @@ func requireProxyModule(t *testing.T, fns FunctionExporter, config wazero.Module
 	var log bytes.Buffer
 
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(testCtx, FunctionListenerFactoryKey{},
+	ctx := WithFunctionListenerFactory(testCtx,
 		proxy.NewLoggingListenerFactory(&log, scopes))
 
 	r := wazero.NewRuntime(ctx)
diff --git a/imports/emscripten/emscripten_test.go b/imports/emscripten/emscripten_test.go
index dd5eb95b..a68aa937 100644
--- a/imports/emscripten/emscripten_test.go
+++ b/imports/emscripten/emscripten_test.go
@@ -43,7 +43,7 @@ func TestGrow(t *testing.T) {
 	var log bytes.Buffer
 
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{},
+	ctx := experimental.WithFunctionListenerFactory(testCtx,
 		logging.NewHostLoggingListenerFactory(&log, logging.LogScopeMemory))
 
 	r := wazero.NewRuntime(ctx)
@@ -344,7 +344,7 @@ func TestInstantiateForModule(t *testing.T) {
 	var log bytes.Buffer
 
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&log))
+	ctx := experimental.WithFunctionListenerFactory(testCtx, logging.NewLoggingListenerFactory(&log))
 
 	r := wazero.NewRuntime(ctx)
 	defer r.Close(ctx)
diff --git a/imports/wasi_snapshot_preview1/wasi_test.go b/imports/wasi_snapshot_preview1/wasi_test.go
index 5cc50ba9..7b250b29 100644
--- a/imports/wasi_snapshot_preview1/wasi_test.go
+++ b/imports/wasi_snapshot_preview1/wasi_test.go
@@ -92,7 +92,7 @@ func requireProxyModuleWithContext(ctx context.Context, t *testing.T, config waz
 	var log bytes.Buffer
 
 	// Set context to one that has an experimental listener
-	ctx = context.WithValue(ctx, experimental.FunctionListenerFactoryKey{},
+	ctx = experimental.WithFunctionListenerFactory(ctx,
 		proxy.NewLoggingListenerFactory(&log, logging.LogScopeAll))
 
 	r := wazero.NewRuntime(ctx)
@@ -121,7 +121,7 @@ func requireErrnoNosys(t *testing.T, funcName string, params ...uint64) string {
 	var log bytes.Buffer
 
 	// Set context to one that has an experimental listener
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{},
+	ctx := experimental.WithFunctionListenerFactory(testCtx,
 		proxy.NewLoggingListenerFactory(&log, logging.LogScopeAll))
 
 	r := wazero.NewRuntime(ctx)
diff --git a/internal/engine/wazevo/e2e_test.go b/internal/engine/wazevo/e2e_test.go
index 98ba2ce7..078f4e77 100644
--- a/internal/engine/wazevo/e2e_test.go
+++ b/internal/engine/wazevo/e2e_test.go
@@ -912,7 +912,7 @@ func TestE2E(t *testing.T) {
 
 func TestE2E_host_functions(t *testing.T) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	for _, tc := range []struct {
 		name string
@@ -1216,7 +1216,7 @@ wasm stack trace:
 func TestListener_local(t *testing.T) {
 	var buf bytes.Buffer
 	config := wazero.NewRuntimeConfigCompiler()
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	r := wazero.NewRuntimeWithConfig(ctx, config)
 	defer func() {
@@ -1244,7 +1244,7 @@ func TestListener_local(t *testing.T) {
 func TestListener_imported(t *testing.T) {
 	var buf bytes.Buffer
 	config := wazero.NewRuntimeConfigCompiler()
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	r := wazero.NewRuntimeWithConfig(ctx, config)
 	defer func() {
@@ -1295,7 +1295,7 @@ func TestListener_long(t *testing.T) {
 
 	var buf bytes.Buffer
 	config := wazero.NewRuntimeConfigCompiler()
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	r := wazero.NewRuntimeWithConfig(ctx, config)
 	defer func() {
@@ -1345,7 +1345,7 @@ func TestListener_long_as_is(t *testing.T) {
 
 	var buf bytes.Buffer
 	config := wazero.NewRuntimeConfigCompiler()
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	r := wazero.NewRuntimeWithConfig(ctx, config)
 	defer func() {
@@ -1394,7 +1394,7 @@ func TestListener_long_many_consts(t *testing.T) {
 
 	var buf bytes.Buffer
 	config := wazero.NewRuntimeConfigCompiler()
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	r := wazero.NewRuntimeWithConfig(ctx, config)
 	defer func() {
diff --git a/internal/integration_test/engine/adhoc_test.go b/internal/integration_test/engine/adhoc_test.go
index 37adfdac..f0936480 100644
--- a/internal/integration_test/engine/adhoc_test.go
+++ b/internal/integration_test/engine/adhoc_test.go
@@ -1144,7 +1144,7 @@ func testModuleMemory(t *testing.T, r wazero.Runtime) {
 
 func testTwoIndirection(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(testCtx, logging.NewLoggingListenerFactory(&buf))
 	_, err := r.NewHostModuleBuilder("host").NewFunctionBuilder().WithFunc(func(
 		_ context.Context, m api.Module, d uint32,
 	) uint32 {
@@ -1328,7 +1328,7 @@ func testBeforeListenerGlobals(t *testing.T, r wazero.Runtime) {
 		ExportSection: []wasm.Export{{Name: "f", Type: wasm.ExternTypeFunc, Index: 0}},
 	})
 
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{}, fnListener)
+	ctx := experimental.WithFunctionListenerFactory(testCtx, fnListener)
 	inst, err := r.Instantiate(ctx, buf)
 	require.NoError(t, err)
 
@@ -1382,7 +1382,7 @@ func testBeforeListenerStackIterator(t *testing.T, r wazero.Runtime) {
 		},
 	}
 
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{}, fnListener)
+	ctx := experimental.WithFunctionListenerFactory(testCtx, fnListener)
 	_, err := r.NewHostModuleBuilder("host").NewFunctionBuilder().WithFunc(func(x int32) int32 {
 		return x + 100
 	}).Export("f4").Instantiate(ctx)
@@ -1484,7 +1484,7 @@ func testListenerStackIteratorOffset(t *testing.T, r wazero.Runtime) {
 			tape = append(tape, stack)
 		},
 	}
-	ctx := context.WithValue(testCtx, experimental.FunctionListenerFactoryKey{}, fnListener)
+	ctx := experimental.WithFunctionListenerFactory(testCtx, fnListener)
 
 	// Minimal DWARF info section to make debug/dwarf.New() happy.
 	// Necessary to make the compiler emit source offset maps.
@@ -1818,7 +1818,7 @@ func testManyParamsResultsCallManyConsts(t *testing.T, r wazero.Runtime) {
 
 func testManyParamsResultsCallManyConstsListener(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	bin, _ := manyParamsResultsMod()
 	mod, err := r.Instantiate(ctx, bin)
@@ -1887,7 +1887,7 @@ func testManyParamsResultsDoubler(t *testing.T, r wazero.Runtime) {
 
 func testManyParamsResultsDoublerListener(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	bin, params := manyParamsResultsMod()
 	mod, err := r.Instantiate(ctx, bin)
@@ -1947,7 +1947,7 @@ func testManyParamsResultsSwapper(t *testing.T, r wazero.Runtime) {
 
 func testManyParamsResultsSwapperListener(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	bin, params := manyParamsResultsMod()
 	mod, err := r.Instantiate(ctx, bin)
@@ -2002,7 +2002,7 @@ func testManyParamsResultsMain(t *testing.T, r wazero.Runtime) {
 
 func testManyParamsResultsMainListener(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	bin, params := manyParamsResultsMod()
 	mod, err := r.Instantiate(ctx, bin)
@@ -2051,7 +2051,7 @@ func testManyParamsResultsCallManyConstsAndPickLastVector(t *testing.T, r wazero
 
 func testManyParamsResultsCallManyConstsAndPickLastVectorListener(t *testing.T, r wazero.Runtime) {
 	var buf bytes.Buffer
-	ctx := context.WithValue(context.Background(), experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&buf))
+	ctx := experimental.WithFunctionListenerFactory(context.Background(), logging.NewLoggingListenerFactory(&buf))
 
 	bin, _ := manyParamsResultsMod()
 	mod, err := r.Instantiate(ctx, bin)
diff --git a/internal/integration_test/filecache/filecache_test.go b/internal/integration_test/filecache/filecache_test.go
index b7d7d8ee..58ff4fa6 100644
--- a/internal/integration_test/filecache/filecache_test.go
+++ b/internal/integration_test/filecache/filecache_test.go
@@ -117,8 +117,8 @@ func testListeners(t *testing.T, config wazero.RuntimeConfig) {
 		dir := t.TempDir()
 
 		out := bytes.NewBuffer(nil)
-		ctxWithListener := context.WithValue(context.Background(),
-			experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+		ctxWithListener := experimental.WithFunctionListenerFactory(
+			context.Background(), logging.NewLoggingListenerFactory(out))
 
 		{
 			cc, err := wazero.NewCompilationCacheWithDir(dir)
@@ -163,8 +163,8 @@ func testListeners(t *testing.T, config wazero.RuntimeConfig) {
 			rc := config.WithCompilationCache(cc)
 
 			out := bytes.NewBuffer(nil)
-			ctxWithListener := context.WithValue(context.Background(),
-				experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+			ctxWithListener := experimental.WithFunctionListenerFactory(
+				context.Background(), logging.NewLoggingListenerFactory(out))
 			r := wazero.NewRuntimeWithConfig(ctxWithListener, rc)
 			_, err = r.CompileModule(ctxWithListener, wasmBin)
 			require.NoError(t, err)
@@ -200,8 +200,8 @@ func testListeners(t *testing.T, config wazero.RuntimeConfig) {
 
 		// Then compile with listeners -> run it.
 		out := bytes.NewBuffer(nil)
-		ctxWithListener := context.WithValue(context.Background(),
-			experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+		ctxWithListener := experimental.WithFunctionListenerFactory(
+			context.Background(), logging.NewLoggingListenerFactory(out))
 
 		cc, err := wazero.NewCompilationCacheWithDir(dir)
 		require.NoError(t, err)
diff --git a/internal/testing/nodiff/nodiff.go b/internal/testing/nodiff/nodiff.go
index 8364fcbb..9a6c2ad7 100644
--- a/internal/testing/nodiff/nodiff.go
+++ b/internal/testing/nodiff/nodiff.go
@@ -60,8 +60,8 @@ func RequireNoDiff(wasmBin []byte, checkMemory, loggingCheck bool, requireNoErro
 	var interPreterLoggingBuf, compilerLoggingBuf bytes.Buffer
 	var errorDuringInvocation bool
 	if loggingCheck {
-		interpreterCtx = context.WithValue(interpreterCtx, experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&interPreterLoggingBuf))
-		compilerCtx = context.WithValue(compilerCtx, experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(&compilerLoggingBuf))
+		interpreterCtx = experimental.WithFunctionListenerFactory(interpreterCtx, logging.NewLoggingListenerFactory(&interPreterLoggingBuf))
+		compilerCtx = experimental.WithFunctionListenerFactory(compilerCtx, logging.NewLoggingListenerFactory(&compilerLoggingBuf))
 		defer func() {
 			if !errorDuringInvocation {
 				if !bytes.Equal(compilerLoggingBuf.Bytes(), interPreterLoggingBuf.Bytes()) {