Cleanups ctxkey package (#2184)

Signed-off-by: Takeshi Yoneda <[email protected]>

experimental/checkpoint.go

@@ -3,7 +3,7 @@ package experimental
 import (
 	"context"
 
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 )
 
 // Snapshot holds the execution state at the time of a Snapshotter.Snapshot call.
@@ -26,23 +26,23 @@ type Snapshotter interface {
 // to a non-nil value, and host functions will be able to retrieve it using SnapshotterKey.
 //
 // Deprecated: use WithSnapshotter to enable snapshots.
-type EnableSnapshotterKey = ctxkey.EnableSnapshotterKey
+type EnableSnapshotterKey = expctxkeys.EnableSnapshotterKey
 
 // WithSnapshotter enables snapshots.
 // Passing the returned context to a exported function invocation enables snapshots,
 // and allows host functions to retrieve the Snapshotter using GetSnapshotter.
 func WithSnapshotter(ctx context.Context) context.Context {
-	return context.WithValue(ctx, ctxkey.EnableSnapshotterKey{}, struct{}{})
+	return context.WithValue(ctx, expctxkeys.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.
 //
 // Deprecated: use GetSnapshotter to get the snapshotter.
-type SnapshotterKey = ctxkey.SnapshotterKey
+type SnapshotterKey = expctxkeys.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)
+	return ctx.Value(expctxkeys.SnapshotterKey{}).(Snapshotter)
 }

experimental/close.go

@@ -3,7 +3,7 @@ package experimental
 import (
 	"context"
 
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 )
 
 // 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, ctxkey.CloseNotifierKey{}, notifier)
+		return context.WithValue(ctx, expctxkeys.CloseNotifierKey{}, notifier)
 	}
 	return ctx
 }

experimental/close_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	"github.com/tetratelabs/wazero/experimental"
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 	"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(ctxkey.CloseNotifierKey{}))
+				require.NotNil(t, decorated.Value(expctxkeys.CloseNotifierKey{}))
 			} else {
 				require.Same(t, testCtx, decorated)
 			}

experimental/listener.go

@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/tetratelabs/wazero/api"
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 )
 
 // StackIterator allows iterating on each function of the call stack, starting
@@ -28,12 +28,12 @@ type StackIterator interface {
 // Its associated value should be a FunctionListenerFactory.
 //
 // Deprecated: use WithFunctionListenerFactory to enable snapshots.
-type FunctionListenerFactoryKey = ctxkey.FunctionListenerFactoryKey
+type FunctionListenerFactoryKey = expctxkeys.FunctionListenerFactoryKey
 
 // WithFunctionListenerFactory registers a FunctionListenerFactory
 // with the context.
 func WithFunctionListenerFactory(ctx context.Context, factory FunctionListenerFactory) context.Context {
-	return context.WithValue(ctx, ctxkey.FunctionListenerFactoryKey{}, factory)
+	return context.WithValue(ctx, expctxkeys.FunctionListenerFactoryKey{}, factory)
 }
 
 // FunctionListenerFactory returns FunctionListeners to be notified when a

experimental/memory.go

@@ -3,7 +3,7 @@ package experimental
 import (
 	"context"
 
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 )
 
 // MemoryAllocator is a memory allocation hook which is invoked
@@ -15,7 +15,7 @@ type MemoryAllocator func(min, cap, max uint64) MemoryBuffer
 
 // MemoryBuffer is a memory buffer that backs a Wasm memory.
 type MemoryBuffer interface {
-	// Return the backing []byte for the memory buffer.
+	// Buffer returns the backing []byte for the memory buffer.
 	Buffer() []byte
 	// Grow the backing memory buffer to size bytes in length.
 	// To back a shared memory, Grow can't change the address
@@ -29,7 +29,7 @@ type MemoryBuffer interface {
 // context.Context.
 func WithMemoryAllocator(ctx context.Context, allocator MemoryAllocator) context.Context {
 	if allocator != nil {
-		return context.WithValue(ctx, ctxkey.MemoryAllocatorKey{}, allocator)
+		return context.WithValue(ctx, expctxkeys.MemoryAllocatorKey{}, allocator)
 	}
 	return ctx
 }

internal/engine/interpreter/interpreter.go

@@ -12,7 +12,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 	"github.com/tetratelabs/wazero/internal/filecache"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/moremath"
@@ -573,8 +573,8 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u
 		}
 	}
 
-	if ctx.Value(ctxkey.EnableSnapshotterKey{}) != nil {
-		ctx = context.WithValue(ctx, ctxkey.SnapshotterKey{}, ce)
+	if ctx.Value(expctxkeys.EnableSnapshotterKey{}) != nil {
+		ctx = context.WithValue(ctx, expctxkeys.SnapshotterKey{}, ce)
 	}
 
 	defer func() {
@@ -744,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(ctxkey.EnableSnapshotterKey{}) != nil {
+				if ctx.Value(expctxkeys.EnableSnapshotterKey{}) != nil {
 					defer func() {
 						if r := recover(); r != nil {
 							if s, ok := r.(*snapshot); ok && s.ce == ce {

internal/engine/wazevo/call_engine.go

@@ -11,8 +11,8 @@ 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/expctxkeys"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wasmdebug"
@@ -194,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(ctxkey.EnableSnapshotterKey{}) != nil
+	snapshotEnabled := ctx.Value(expctxkeys.EnableSnapshotterKey{}) != nil
 	if snapshotEnabled {
-		ctx = context.WithValue(ctx, ctxkey.SnapshotterKey{}, c)
+		ctx = context.WithValue(ctx, expctxkeys.SnapshotterKey{}, c)
 	}
 
 	if wazevoapi.StackGuardCheckEnabled {

internal/ctxkey/checkpoint.go → internal/expctxkeys/checkpoint.go

@@ -1,4 +1,4 @@
-package ctxkey
+package expctxkeys
 
 // 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

internal/expctxkeys/close.go

@@ -0,0 +1,5 @@
+package expctxkeys
+
+// CloseNotifierKey is a context.Context Value key. Its associated value should be a
+// Notifier.
+type CloseNotifierKey struct{}

internal/expctxkeys/expctxkeys.go

@@ -0,0 +1,2 @@
+// Package expctxkeys provides keys for the context used to store the experimental APIs.
+package expctxkeys

internal/ctxkey/listener.go → internal/expctxkeys/listener.go

@@ -1,4 +1,4 @@
-package ctxkey
+package expctxkeys
 
 // FunctionListenerFactoryKey is a context.Context Value key.
 // Its associated value should be a FunctionListenerFactory.

internal/expctxkeys/memory.go

@@ -0,0 +1,4 @@
+package expctxkeys
+
+// MemoryAllocatorKey is a context.Context key for the experimental memory allocator.
+type MemoryAllocatorKey struct{}

internal/wasm/store.go

@@ -9,7 +9,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/experimental"
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/leb128"
 	internalsys "github.com/tetratelabs/wazero/internal/sys"
@@ -125,7 +125,7 @@ type (
 		Source *Module
 
 		// CloseNotifier is an experimental hook called once on close.
-		CloseNotifier ctxkey.Notifier
+		CloseNotifier experimental.CloseNotifier
 	}
 
 	// DataInstance holds bytes corresponding to the data segment in a module.
@@ -365,7 +365,7 @@ func (s *Store) instantiate(
 
 	var allocator experimental.MemoryAllocator
 	if ctx != nil {
-		allocator, _ = ctx.Value(ctxkey.MemoryAllocatorKey{}).(experimental.MemoryAllocator)
+		allocator, _ = ctx.Value(expctxkeys.MemoryAllocatorKey{}).(experimental.MemoryAllocator)
 	}
 
 	m.buildGlobals(module, m.Engine.FunctionInstanceReference)

runtime.go

@@ -7,7 +7,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	experimentalapi "github.com/tetratelabs/wazero/experimental"
-	"github.com/tetratelabs/wazero/internal/ctxkey"
+	"github.com/tetratelabs/wazero/internal/expctxkeys"
 	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(ctxkey.FunctionListenerFactoryKey{})
+	fnlf := ctx.Value(expctxkeys.FunctionListenerFactoryKey{})
 	if fnlf == nil {
 		return nil, nil
 	}
@@ -318,7 +318,7 @@ func (r *runtime) InstantiateModule(
 		return
 	}
 
-	if closeNotifier, ok := ctx.Value(ctxkey.CloseNotifierKey{}).(ctxkey.Notifier); ok {
+	if closeNotifier, ok := ctx.Value(expctxkeys.CloseNotifierKey{}).(experimentalapi.CloseNotifier); ok {
 		mod.(*wasm.ModuleInstance).CloseNotifier = closeNotifier
 	}