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/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.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/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.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/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/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/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/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/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()) { 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 }