From 0300f4b3c1bd1532b76b476a2e242fbf8953786f Mon Sep 17 00:00:00 2001 From: Crypt Keeper <64215+codefromthecrypt@users.noreply.github.com> Date: Tue, 11 Jul 2023 09:27:43 +0800 Subject: [PATCH] experimental: adds close notification hook (#1574) Signed-off-by: Adrian Cole --- experimental/close.go | 63 ++++++++++++++++++++++++++++++ experimental/close_example_test.go | 27 +++++++++++++ experimental/close_test.go | 42 ++++++++++++++++++++ experimental/gojs/gojs.go | 57 +++++++++++++-------------- internal/close/close.go | 13 ++++++ internal/gojs/compiler_test.go | 10 ++--- internal/gojs/run/gojs.go | 19 +++------ internal/wasm/module_instance.go | 8 +++- internal/wasm/store.go | 4 ++ runtime.go | 8 +++- 10 files changed, 201 insertions(+), 50 deletions(-) create mode 100644 experimental/close.go create mode 100644 experimental/close_example_test.go create mode 100644 experimental/close_test.go create mode 100644 internal/close/close.go diff --git a/experimental/close.go b/experimental/close.go new file mode 100644 index 00000000..1c37e8c0 --- /dev/null +++ b/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 +} diff --git a/experimental/close_example_test.go b/experimental/close_example_test.go new file mode 100644 index 00000000..59eb08e8 --- /dev/null +++ b/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 + } +} diff --git a/experimental/close_test.go b/experimental/close_test.go new file mode 100644 index 00000000..1fe9d8e3 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/experimental/gojs/gojs.go b/experimental/gojs/gojs.go index ebaabd78..4bbef6aa 100644 --- a/experimental/gojs/gojs.go +++ b/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) } diff --git a/internal/close/close.go b/internal/close/close.go new file mode 100644 index 00000000..33b39b1e --- /dev/null +++ b/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) +} diff --git a/internal/gojs/compiler_test.go b/internal/gojs/compiler_test.go index f50ef895..ff83a9d3 100644 --- a/internal/gojs/compiler_test.go +++ b/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 diff --git a/internal/gojs/run/gojs.go b/internal/gojs/run/gojs.go index 687c7666..651debb8 100644 --- a/internal/gojs/run/gojs.go +++ b/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 } diff --git a/internal/wasm/module_instance.go b/internal/wasm/module_instance.go index 7527f253..3e2e21dc 100644 --- a/internal/wasm/module_instance.go +++ b/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 diff --git a/internal/wasm/store.go b/internal/wasm/store.go index 70af5e6f..e853d7fa 100644 --- a/internal/wasm/store.go +++ b/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. diff --git a/runtime.go b/runtime.go index ce873416..70c63cda 100644 --- a/runtime.go +++ b/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 {