Skip to content

Commit

Permalink
experimental: adds close notification hook (#1574)
Browse files Browse the repository at this point in the history
Signed-off-by: Adrian Cole <adrian@tetrate.io>
  • Loading branch information
codefromthecrypt committed Jul 11, 2023
1 parent 326c267 commit 0300f4b
Show file tree
Hide file tree
Showing 10 changed files with 201 additions and 50 deletions.
63 changes: 63 additions & 0 deletions experimental/close.go
Original file line number Diff line number Diff line change
@@ -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
}
27 changes: 27 additions & 0 deletions experimental/close_example_test.go
Original file line number Diff line number Diff line change
@@ -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
}
}
42 changes: 42 additions & 0 deletions experimental/close_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
})
}
}
57 changes: 28 additions & 29 deletions experimental/gojs/gojs.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
}
13 changes: 13 additions & 0 deletions internal/close/close.go
Original file line number Diff line number Diff line change
@@ -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)
}
10 changes: 5 additions & 5 deletions internal/gojs/compiler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
19 changes: 6 additions & 13 deletions internal/gojs/run/gojs.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -46,5 +39,5 @@ func RunAndReturnState(
err = nil
}
}
return s, err
return err
}
8 changes: 7 additions & 1 deletion internal/wasm/module_instance.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions internal/wasm/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand Down
8 changes: 6 additions & 2 deletions runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand Down

0 comments on commit 0300f4b

Please sign in to comment.