Repository: tetratelabs/wazero (go)

Issue:
Add Runtime.Close
I noticed that users of wazero including ourselves often have a single imported module (WASI) and the one they made. To close requires them closing those in reverse order. I wonder if we can't have `Runtime.Close` which can do the same?

An alternative would be if WASI could be special cased to never need closing. I don't know off the top of my head if we could pre-compile any glue needed and/or use offsets based on a go:embed FS or any other tricks to make WASI be safe to always be around.

Pull request description:
Adds Runtime.Close
Fixes #382

It can be tedious to close many modules that have been instantiated, particularly when they are not instantiated in a scope where it is appropriate to use `defer`. `Runtime.Close` closes all the modules in a runtime's store, in reverse order of instantiation. I am wondering whether we should provide a version of this that takes an exit code like the `Module` version or if that would be overkill. It wouldn't be hard to add that if we need it here.

To keep track of order, `moduleNames` changes from a map to a slice - with a huge number of modules this could reduce instantiation performance. If this seems like it needs to be avoided, I can use a separate field instead, which would be something like inlining a linked map of modules into the store.

Also because Go locks aren't reentrant, I had to provide an internal module close codepath that doesn't remove from the store. There is a bit of circularity in that a module calls into store to remove itself, but when closing the store it is the other way around. Perhaps we may consider changing `Module.Close` to `Runtime.CloseModule(Module)` as a random idea.

Patch:
diff --git a/internal/wasm/call_context.go b/internal/wasm/call_context.go
index 130d6375..b9ee95b8 100644
--- a/internal/wasm/call_context.go
+++ b/internal/wasm/call_context.go
@@ -79,18 +79,28 @@ func (m *CallContext) Close(ctx context.Context) (err error) {
 }
 
 // CloseWithExitCode implements the same method as documented on api.Module.
-func (m *CallContext) CloseWithExitCode(_ context.Context, exitCode uint32) (err error) {
+func (m *CallContext) CloseWithExitCode(ctx context.Context, exitCode uint32) (err error) {
+	closed, err := m.close(ctx, exitCode)
+	if !closed {
+		return nil
+	}
+	m.store.deleteModule(m.Name())
+	return err
+}
+
+// close marks this CallContext as closed and releases underlying system resources without removing
+// from the store.
+func (m *CallContext) close(_ context.Context, exitCode uint32) (c bool, err error) {
 	// Note: If you use the context.Context param, don't forget to coerce nil to context.Background()!
 
 	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
 	if !atomic.CompareAndSwapUint64(m.closed, 0, closed) {
-		return nil
+		return false, nil
 	}
-	m.store.deleteModule(m.Name())
 	if sys := m.Sys; sys != nil { // ex nil if from ModuleBuilder
-		return sys.Close()
+		return true, sys.Close()
 	}
-	return
+	return true, nil
 }
 
 // Memory implements the same method as documented on api.Module.
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 6b1faaac..37d200a1 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -33,7 +33,7 @@ type (
 		Engine Engine
 
 		// moduleNames ensures no race conditions instantiating two modules of the same name
-		moduleNames map[string]struct{} // guarded by mux
+		moduleNames []string // guarded by mux
 
 		// modules holds the instantiated Wasm modules by module name from Instantiate.
 		modules map[string]*ModuleInstance // guarded by mux
@@ -309,7 +309,7 @@ func NewStore(enabledFeatures Features, engine Engine) *Store {
 	return &Store{
 		EnabledFeatures:  enabledFeatures,
 		Engine:           engine,
-		moduleNames:      map[string]struct{}{},
+		moduleNames:      nil,
 		modules:          map[string]*ModuleInstance{},
 		typeIDs:          map[string]FunctionTypeID{},
 		functionMaxTypes: maximumFunctionTypes,
@@ -413,7 +413,13 @@ func (s *Store) deleteModule(moduleName string) {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 	delete(s.modules, moduleName)
-	delete(s.moduleNames, moduleName)
+	// remove this module name
+	for i, n := range s.moduleNames {
+		if n == moduleName {
+			s.moduleNames = append(s.moduleNames[:i], s.moduleNames[i+1:]...)
+			break
+		}
+	}
 }
 
 // requireModuleName is a pre-flight check to reserve a module.
@@ -421,10 +427,12 @@ func (s *Store) deleteModule(moduleName string) {
 func (s *Store) requireModuleName(moduleName string) error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
-	if _, ok := s.moduleNames[moduleName]; ok {
-		return fmt.Errorf("module %s has already been instantiated", moduleName)
+	for _, n := range s.moduleNames {
+		if n == moduleName {
+			return fmt.Errorf("module %s has already been instantiated", moduleName)
+		}
 	}
-	s.moduleNames[moduleName] = struct{}{}
+	s.moduleNames = append(s.moduleNames, moduleName)
 	return nil
 }
 
@@ -619,3 +627,19 @@ func (s *Store) getFunctionTypeID(t *FunctionType) (FunctionTypeID, error) {
 	}
 	return id, nil
 }
+
+// CloseWithExitCode implements the same method as documented on wazero.Runtime.
+func (s *Store) CloseWithExitCode(ctx context.Context, exitCode uint32) (err error) {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+	// Close modules in reverse initialization order.
+	for i := len(s.moduleNames) - 1; i >= 0; i-- {
+		// If closing this module errs, proceed anyway to close the others.
+		if _, e := s.modules[s.moduleNames[i]].CallCtx.close(ctx, exitCode); e != nil && err == nil {
+			err = e // first error
+		}
+	}
+	s.moduleNames = nil
+	s.modules = map[string]*ModuleInstance{}
+	return err
+}
diff --git a/wasm.go b/wasm.go
index 1696463f..905b6574 100644
--- a/wasm.go
+++ b/wasm.go
@@ -92,6 +92,28 @@ type Runtime interface {
 	// Note: Config is copied during instantiation: Later changes to config do not affect the instantiated result.
 	// Note: When the context is nil, it defaults to context.Background.
 	InstantiateModule(ctx context.Context, compiled CompiledModule, config ModuleConfig) (api.Module, error)
+
+	// Close closes all the modules that have been initialized in this Runtime with an exit code of 0.
+	// An error is returned if any module returns an error when closed.
+	//
+	// Ex.
+	//	ctx := context.Background()
+	//	r := wazero.NewRuntime()
+	//	defer r.Close(ctx)
+	//	_, _ = wasi.InstantiateSnapshotPreview1(ctx, r)
+	//	mod, _ := r.InstantiateModuleFromCode(ctx, source)
+	Close(context.Context) error
+
+	// CloseWithExitCode closes all the modules that have been initialized in this Runtime with the provided exit code.
+	// An error is returned if any module returns an error when closed.
+	//
+	// Ex.
+	//	ctx := context.Background()
+	//	r := wazero.NewRuntime()
+	//	defer r.CloseWithExitCode(ctx, 2)
+	//	_, _ = wasi.InstantiateSnapshotPreview1(ctx, r)
+	//	mod, _ := r.InstantiateModuleFromCode(ctx, source)
+	CloseWithExitCode(ctx context.Context, exitCode uint32) error
 }
 
 func NewRuntime() Runtime {
@@ -114,6 +136,7 @@ func NewRuntimeWithConfig(rConfig RuntimeConfig) Runtime {
 type runtime struct {
 	store           *wasm.Store
 	enabledFeatures wasm.Features
+	compiledModules []*compiledCode
 }
 
 // Module implements Runtime.Module
@@ -166,7 +189,9 @@ func (r *runtime) CompileModule(ctx context.Context, source []byte, cConfig Comp
 		return nil, err
 	}
 
-	return &compiledCode{module: internal, compiledEngine: r.store.Engine}, nil
+	c := &compiledCode{module: internal, compiledEngine: r.store.Engine}
+	r.compiledModules = append(r.compiledModules, c)
+	return c, nil
 }
 
 //
@@ -280,3 +305,19 @@ func (r *runtime) InstantiateModule(ctx context.Context, compiled CompiledModule
 	}
 	return
 }
+
+// Close implements Runtime.Close
+func (r *runtime) Close(ctx context.Context) error {
+	return r.CloseWithExitCode(ctx, 0)
+}
+
+// CloseWithExitCode implements Runtime.CloseWithExitCode
+func (r *runtime) CloseWithExitCode(ctx context.Context, exitCode uint32) error {
+	err := r.store.CloseWithExitCode(ctx, exitCode)
+	for _, c := range r.compiledModules {
+		if e := c.Close(ctx); e != nil && err == nil {
+			err = e
+		}
+	}
+	return err
+}


Test patch:
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 83d454fd..f1ccad3f 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -182,6 +182,60 @@ func TestStore_CloseModule(t *testing.T) {
 	}
 }
 
+func TestStore_CloseStore(t *testing.T) {
+	const importedModuleName = "imported"
+	const importingModuleName = "test"
+
+	for _, tc := range []struct {
+		name       string
+		testClosed bool
+	}{
+		{
+			name:       "nothing closed",
+			testClosed: false,
+		},
+		{
+			name:       "partially closed",
+			testClosed: true,
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			s := newStore()
+
+			m1, err := s.Instantiate(testCtx, &Module{
+				TypeSection:     []*FunctionType{{}},
+				FunctionSection: []uint32{0},
+				CodeSection:     []*Code{{Body: []byte{OpcodeEnd}}},
+				ExportSection:   []*Export{{Type: ExternTypeFunc, Index: 0, Name: "fn"}},
+			}, importedModuleName, nil, nil)
+			require.NoError(t, err)
+
+			m2, err := s.Instantiate(testCtx, &Module{
+				TypeSection:   []*FunctionType{{}},
+				ImportSection: []*Import{{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0}},
+				MemorySection: &Memory{Min: 1, Cap: 1},
+				GlobalSection: []*Global{{Type: &GlobalType{}, Init: &ConstantExpression{Opcode: OpcodeI32Const, Data: const1}}},
+				TableSection:  []*Table{{Min: 10}},
+			}, importingModuleName, nil, nil)
+			require.NoError(t, err)
+
+			if tc.testClosed {
+				err = m2.CloseWithExitCode(testCtx, 2)
+				require.NoError(t, err)
+			}
+
+			err = s.CloseWithExitCode(testCtx, 2)
+			require.NoError(t, err)
+			require.Nil(t, s.modules[importedModuleName])
+			require.Nil(t, s.modules[importingModuleName])
+
+			require.Equal(t, uint64(1)+uint64(2)<<32, *m1.closed)
+			require.Equal(t, uint64(1)+uint64(2)<<32, *m2.closed)
+		})
+	}
+}
+
 func TestStore_hammer(t *testing.T) {
 	const importedModuleName = "imported"
 
diff --git a/wasm_test.go b/wasm_test.go
index a141a143..38d3d012 100644
--- a/wasm_test.go
+++ b/wasm_test.go
@@ -529,6 +529,73 @@ func TestInstantiateModule_ExitError(t *testing.T) {
 	require.Equal(t, err, sys.NewExitError("env", 2))
 }
 
+func TestClose(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		exitCode uint32
+	}{
+		{
+			name:     "exit code 0",
+			exitCode: uint32(0),
+		},
+		{
+			name:     "exit code 2",
+			exitCode: uint32(2),
+		},
+	} {
+
+		r := NewRuntime()
+
+		m1, err := r.NewModuleBuilder("mod1").ExportFunction("func1", func() {}).Instantiate(testCtx)
+		require.NoError(t, err)
+		m2, err := r.NewModuleBuilder("mod2").ExportFunction("func2", func() {}).Instantiate(testCtx)
+		require.NoError(t, err)
+
+		func1 := m1.ExportedFunction("func1")
+		func2 := m2.ExportedFunction("func2")
+
+		// Modules not closed so calls succeed
+
+		_, err = func1.Call(testCtx)
+		require.NoError(t, err)
+
+		_, err = func2.Call(testCtx)
+		require.NoError(t, err)
+
+		if tc.exitCode == 0 {
+			err = r.Close(testCtx)
+		} else {
+			err = r.CloseWithExitCode(testCtx, tc.exitCode)
+		}
+		require.NoError(t, err)
+
+		// Modules closed so calls fail
+		_, err = func1.Call(testCtx)
+		// TODO: This could be neater with ErrorIs
+		require.Error(t, err, sys.NewExitError("mod1", tc.exitCode).Error())
+
+		_, err = func2.Call(testCtx)
+		require.Error(t, err, sys.NewExitError("mod2", tc.exitCode).Error())
+	}
+}
+
+func TestClose_ClosesCompiledModules(t *testing.T) {
+	engine := &mockEngine{name: "mock", cachedModules: map[*wasm.Module]struct{}{}}
+	conf := *engineLessConfig
+	conf.newEngine = func(_ wasm.Features) wasm.Engine {
+		return engine
+	}
+	r := NewRuntimeWithConfig(&conf)
+	// Normally compiled modules are closed when instantiated but this is never instantiated.
+	_, err := r.CompileModule(testCtx, []byte(`(module $0 (memory 1))`), NewCompileConfig())
+	require.NoError(t, err)
+	require.Equal(t, 1, len(engine.cachedModules))
+
+	err = r.Close(testCtx)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(engine.cachedModules))
+}
+
 // requireImportAndExportFunction re-exports a host function because only host functions can see the propagated context.
 func requireImportAndExportFunction(t *testing.T, r Runtime, hostFn func(ctx context.Context) uint64, functionName string) ([]byte, func(context.Context) error) {
 	mod, err := r.NewModuleBuilder("host").ExportFunction(functionName, hostFn).Instantiate(testCtx)