experimental: adds close notification hook (#1574)

Signed-off-by: Adrian Cole <adrian@tetrate.io>

experimental/close.go

@@ -0,0 +1,63 @@
+package experimental
+
+import (
+	"context"
+
+	"github.com/tetratelabs/wazero/internal/close"
+)
+
+// CloseNotifier is a notification hook, invoked when a module is closed.
+//
+// Note: This is experimental progress towards #1197, and likely to change. Do
+// not expose this in shared libraries as it can cause version locks.
+type CloseNotifier interface {
+	// CloseNotify is a notification that occurs *before* an api.Module is
+	// closed. `exitCode` is zero on success or in the case there was no exit
+	// code.
+	//
+	// Notes:
+	//   - This does not return an error because the module will be closed
+	//     unconditionally.
+	//   - Do not panic from this function as it doing so could cause resource
+	//     leaks.
+	//   - While this is only called once per module, if configured for
+	//     multiple modules, it will be called for each, e.g. on runtime close.
+	CloseNotify(ctx context.Context, exitCode uint32)
+}
+
+// ^-- Note: This might need to be a part of the listener or become a part of
+// host state implementation. For example, if this is used to implement state
+// cleanup for host modules, possibly something like below would be better, as
+// it could be implemented in a way that allows concurrent module use.
+//
+//	// key is like a context key, stateFactory is invoked per instantiate and
+//	// is associated with the key (exposed as `Module.State` similar to go
+//	// context). Using a key is better than the module name because we can
+//	// de-dupe it for host modules that can be instantiated into different
+//	// names. Also, you can make the key package private.
+//	HostModuleBuilder.WithState(key any, stateFactory func() Cleanup)`
+//
+// Such a design could work to isolate state only needed for wasip1, for
+// example the dirent cache. However, if end users use this for different
+// things, we may need separate designs.
+//
+// In summary, the purpose of this iteration is to identify projects that
+// would use something like this, and then we can figure out which way it
+// should go.
+
+// CloseNotifyFunc is a convenience for defining inlining a CloseNotifier.
+type CloseNotifyFunc func(ctx context.Context, exitCode uint32)
+
+// CloseNotify implements CloseNotifier.CloseNotify.
+func (f CloseNotifyFunc) CloseNotify(ctx context.Context, exitCode uint32) {
+	f(ctx, exitCode)
+}
+
+// WithCloseNotifier registers the given CloseNotifier into the given
+// context.Context.
+func WithCloseNotifier(ctx context.Context, notifier CloseNotifier) context.Context {
+	if notifier != nil {
+		return context.WithValue(ctx, close.NotifierKey{}, notifier)
+	}
+	return ctx
+}

experimental/close_example_test.go

@@ -0,0 +1,27 @@
+package experimental_test
+
+import (
+	"context"
+
+	"github.com/tetratelabs/wazero/experimental"
+)
+
+var ctx context.Context
+
+// This shows how to implement a custom cleanup task on close.
+func Example_closeNotifier() {
+	closeCh := make(chan struct{})
+	ctx = experimental.WithCloseNotifier(
+		ctx,
+		experimental.CloseNotifyFunc(func(context.Context, uint32) { close(closeCh) }),
+	)
+
+	// ... create module, do some work. Sometime later in another goroutine:
+
+	select {
+	case <-closeCh:
+		// do some cleanup
+	default:
+		// do some more work with the module
+	}
+}

experimental/close_test.go

@@ -0,0 +1,42 @@
+package experimental_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/close"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+// testCtx is an arbitrary, non-default context. Non-nil also prevents linter errors.
+var testCtx = context.WithValue(context.Background(), struct{}{}, "arbitrary")
+
+func TestWithCloseNotifier(t *testing.T) {
+	tests := []struct {
+		name         string
+		notification experimental.CloseNotifier
+		expected     bool
+	}{
+		{
+			name:     "returns input when notification nil",
+			expected: false,
+		},
+		{
+			name:         "decorates with notification",
+			notification: experimental.CloseNotifyFunc(func(context.Context, uint32) {}),
+			expected:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		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{}))
+			} else {
+				require.Same(t, testCtx, decorated)
+			}
+		})
+	}
+}

experimental/gojs/gojs.go

@@ -19,9 +19,9 @@ import (
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
-	. "github.com/tetratelabs/wazero/internal/gojs"
+	"github.com/tetratelabs/wazero/internal/gojs"
 	internalconfig "github.com/tetratelabs/wazero/internal/gojs/config"
-	. "github.com/tetratelabs/wazero/internal/gojs/run"
+	"github.com/tetratelabs/wazero/internal/gojs/run"
 	"github.com/tetratelabs/wazero/internal/wasm"
 )
 
@@ -92,31 +92,31 @@ type functionExporter struct{}
 func (e *functionExporter) ExportFunctions(builder wazero.HostModuleBuilder) {
 	hfExporter := builder.(wasm.HostFuncExporter)
 
-	hfExporter.ExportHostFunc(GetRandomData)
-	hfExporter.ExportHostFunc(Nanotime1)
-	hfExporter.ExportHostFunc(WasmExit)
-	hfExporter.ExportHostFunc(CopyBytesToJS)
-	hfExporter.ExportHostFunc(ValueCall)
-	hfExporter.ExportHostFunc(ValueGet)
-	hfExporter.ExportHostFunc(ValueIndex)
-	hfExporter.ExportHostFunc(ValueLength)
-	hfExporter.ExportHostFunc(ValueNew)
-	hfExporter.ExportHostFunc(ValueSet)
-	hfExporter.ExportHostFunc(WasmWrite)
-	hfExporter.ExportHostFunc(ResetMemoryDataView)
-	hfExporter.ExportHostFunc(Walltime)
-	hfExporter.ExportHostFunc(ScheduleTimeoutEvent)
-	hfExporter.ExportHostFunc(ClearTimeoutEvent)
-	hfExporter.ExportHostFunc(FinalizeRef)
-	hfExporter.ExportHostFunc(StringVal)
-	hfExporter.ExportHostFunc(ValueDelete)
-	hfExporter.ExportHostFunc(ValueSetIndex)
-	hfExporter.ExportHostFunc(ValueInvoke)
-	hfExporter.ExportHostFunc(ValuePrepareString)
-	hfExporter.ExportHostFunc(ValueInstanceOf)
-	hfExporter.ExportHostFunc(ValueLoadString)
-	hfExporter.ExportHostFunc(CopyBytesToGo)
-	hfExporter.ExportHostFunc(Debug)
+	hfExporter.ExportHostFunc(gojs.GetRandomData)
+	hfExporter.ExportHostFunc(gojs.Nanotime1)
+	hfExporter.ExportHostFunc(gojs.WasmExit)
+	hfExporter.ExportHostFunc(gojs.CopyBytesToJS)
+	hfExporter.ExportHostFunc(gojs.ValueCall)
+	hfExporter.ExportHostFunc(gojs.ValueGet)
+	hfExporter.ExportHostFunc(gojs.ValueIndex)
+	hfExporter.ExportHostFunc(gojs.ValueLength)
+	hfExporter.ExportHostFunc(gojs.ValueNew)
+	hfExporter.ExportHostFunc(gojs.ValueSet)
+	hfExporter.ExportHostFunc(gojs.WasmWrite)
+	hfExporter.ExportHostFunc(gojs.ResetMemoryDataView)
+	hfExporter.ExportHostFunc(gojs.Walltime)
+	hfExporter.ExportHostFunc(gojs.ScheduleTimeoutEvent)
+	hfExporter.ExportHostFunc(gojs.ClearTimeoutEvent)
+	hfExporter.ExportHostFunc(gojs.FinalizeRef)
+	hfExporter.ExportHostFunc(gojs.StringVal)
+	hfExporter.ExportHostFunc(gojs.ValueDelete)
+	hfExporter.ExportHostFunc(gojs.ValueSetIndex)
+	hfExporter.ExportHostFunc(gojs.ValueInvoke)
+	hfExporter.ExportHostFunc(gojs.ValuePrepareString)
+	hfExporter.ExportHostFunc(gojs.ValueInstanceOf)
+	hfExporter.ExportHostFunc(gojs.ValueLoadString)
+	hfExporter.ExportHostFunc(gojs.CopyBytesToGo)
+	hfExporter.ExportHostFunc(gojs.Debug)
 }
 
 // Config extends wazero.ModuleConfig with GOOS=js specific extensions.
@@ -195,6 +195,5 @@ func (c *cfg) WithOSWorkdir() Config {
 //   - The guest module is closed after being run.
 func Run(ctx context.Context, r wazero.Runtime, compiled wazero.CompiledModule, moduleConfig Config) error {
 	c := moduleConfig.(*cfg)
-	_, err := RunAndReturnState(ctx, r, compiled, c.moduleConfig, c.internal)
-	return err
+	return run.Run(ctx, r, compiled, c.moduleConfig, c.internal)
 }

internal/close/close.go

@@ -0,0 +1,13 @@
+// Package close allows experimental.CloseNotifier without introducing a
+// package cycle.
+package close
+
+import "context"
+
+// NotifierKey is a context.Context Value key. Its associated value should be a
+// Notifier.
+type NotifierKey struct{}
+
+type Notifier interface {
+	CloseNotify(ctx context.Context, exitCode uint32)
+}

internal/gojs/compiler_test.go

@@ -17,6 +17,7 @@ import (
 	"time"
 
 	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/experimental/gojs"
 	"github.com/tetratelabs/wazero/internal/fstest"
 	internalgojs "github.com/tetratelabs/wazero/internal/gojs"
@@ -58,14 +59,13 @@ func compileAndRunWithRuntime(ctx context.Context, r wazero.Runtime, arg string,
 		WithStderr(&stderrBuf).
 		WithArgs("test", arg))
 
-	var s *internalgojs.State
-	s, err = run.RunAndReturnState(ctx, r, guest, mc, c)
-	if err == nil {
+	ctx = experimental.WithCloseNotifier(ctx, experimental.CloseNotifyFunc(func(ctx context.Context, exitCode uint32) {
+		s := ctx.Value(internalgojs.StateKey{})
 		if want, have := internalgojs.NewState(c), s; !reflect.DeepEqual(want, have) {
 			log.Panicf("unexpected state: want %#v, have %#v", want, have)
 		}
-	}
-
+	}))
+	err = run.Run(ctx, r, guest, mc, c)
 	stdout = stdoutBuf.String()
 	stderr = stderrBuf.String()
 	return

internal/gojs/run/gojs.go

@@ -11,33 +11,26 @@ import (
 	"github.com/tetratelabs/wazero/sys"
 )
 
-func RunAndReturnState(
-	ctx context.Context,
-	r wazero.Runtime,
-	compiled wazero.CompiledModule,
-	moduleConfig wazero.ModuleConfig,
-	config *config.Config,
-) (*gojs.State, error) {
+func Run(ctx context.Context, r wazero.Runtime, compiled wazero.CompiledModule, moduleConfig wazero.ModuleConfig, config *config.Config) error {
 	if err := config.Init(); err != nil {
-		return nil, err
+		return err
 	}
 
 	// Instantiate the module compiled by go, noting it has no init function.
 	mod, err := r.InstantiateModule(ctx, compiled, moduleConfig)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	defer mod.Close(ctx)
 
 	// Extract the args and env from the module Config and write it to memory.
 	argc, argv, err := gojs.WriteArgsAndEnviron(mod)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	// Create host-side state for JavaScript values and events.
-	s := gojs.NewState(config)
-	ctx = context.WithValue(ctx, gojs.StateKey{}, s)
+	ctx = context.WithValue(ctx, gojs.StateKey{}, gojs.NewState(config))
 
 	// Invoke the run function.
 	_, err = mod.ExportedFunction("run").Call(ctx, uint64(argc), uint64(argv))
@@ -46,5 +39,5 @@ func RunAndReturnState(
 			err = nil
 		}
 	}
-	return s, err
+	return err
 }

internal/wasm/module_instance.go

@@ -144,8 +144,14 @@ func (m *ModuleInstance) setExitCode(exitCode uint32, flag exitCodeFlag) bool {
 }
 
 // ensureResourcesClosed ensures that resources assigned to ModuleInstance is released.
-// Multiple calls to this function is safe.
+// Only one call will happen per module, due to external atomic guards on Closed.
 func (m *ModuleInstance) ensureResourcesClosed(ctx context.Context) (err error) {
+	if closeNotifier := m.CloseNotifier; closeNotifier != nil { // experimental
+		closed := atomic.LoadUint64(&m.Closed)
+		closeNotifier.CloseNotify(ctx, uint32(closed>>32))
+		m.CloseNotifier = nil
+	}
+
 	if sysCtx := m.Sys; sysCtx != nil { // nil if from HostModuleBuilder
 		if err = sysCtx.FS().Close(); err != nil {
 			return err

internal/wasm/store.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/close"
 	"github.com/tetratelabs/wazero/internal/internalapi"
 	"github.com/tetratelabs/wazero/internal/leb128"
 	internalsys "github.com/tetratelabs/wazero/internal/sys"
@@ -124,6 +125,9 @@ type (
 		prev, next *ModuleInstance
 		// Source is a pointer to the Module from which this ModuleInstance derives.
 		Source *Module
+
+		// CloseNotifier is an experimental hook called once on close.
+		CloseNotifier close.Notifier
 	}
 
 	// DataInstance holds bytes corresponding to the data segment in a module.

runtime.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/tetratelabs/wazero/api"
 	experimentalapi "github.com/tetratelabs/wazero/experimental"
+	internalclose "github.com/tetratelabs/wazero/internal/close"
 	internalsock "github.com/tetratelabs/wazero/internal/sock"
 	internalsys "github.com/tetratelabs/wazero/internal/sys"
 	"github.com/tetratelabs/wazero/internal/wasm"
@@ -292,8 +293,7 @@ func (r *runtime) InstantiateModule(
 	code := compiled.(*compiledModule)
 	config := mConfig.(*moduleConfig)
 
-	// Only build listeners on a guest module. A host module doesn't have
-	// memory, and a guest without memory can't use listeners anyway.
+	// Only add guest module configuration to guests.
 	if !code.module.IsHostModule {
 		if sockConfig, ok := ctx.Value(internalsock.ConfigKey{}).(*internalsock.Config); ok {
 			config.sockConfig = sockConfig
@@ -320,6 +320,10 @@ func (r *runtime) InstantiateModule(
 		return
 	}
 
+	if closeNotifier, ok := ctx.Value(internalclose.NotifierKey{}).(internalclose.Notifier); ok {
+		mod.(*wasm.ModuleInstance).CloseNotifier = closeNotifier
+	}
+
 	// Attach the code closer so that anything afterward closes the compiled
 	// code when closing the module.
 	if code.closeWithModule {