Exposes Module.IsClosed to prevent calling functions when closed (#1573)

Signed-off-by: Adrian Cole <adrian@tetrate.io>
Co-authored-by: Takeshi Yoneda <takeshi@tetrate.io>

RATIONALE.md

@@ -484,26 +484,26 @@ https://github.com/bytecodealliance/wasmtime/blob/v0.29.0/crates/lightbeam/src/m
 
 ## Exit
 
-### Why do we return a `sys.ExitError` on exit code zero?
+### Why do we only return a `sys.ExitError` on a non-zero exit code?
 
-It may be surprising to find an error returned on success (exit code zero).
-This can be explained easier when you think of function returns: When results
-aren't empty, then you must return an error. This is trickier to explain when
-results are empty, such as the case in the "_start" function in WASI.
+It is reasonable to think an exit error should be returned, even if the code is
+success (zero). Even on success, the module is no longer functional. For
+example, function exports would error later. However, wazero does not. The only
+time `sys.ExitError` is on error (non-zero).
 
-The main rationale for returning an exit error even if the code is success is
-that the module is no longer functional. For example, function exports would
-error later. In cases like these, it is better to handle errors where they
-occur.
+This decision was to improve performance and ergonomics for guests that both
+use WASI (have a `_start` function), and also allow custom exports.
+Specifically, Rust, TinyGo and normal wasi-libc, don't exit the module during
+`_start`. If they did, it would invalidate their function exports. This means
+it is unlikely most compilers will change this behavior.
 
-Luckily, it is not common to exit a module during the "_start" function. For
-example, the only known compilation target that does this is Emscripten. Most,
-such as Rust, TinyGo, or normal wasi-libc, don't. If they did, it would
-invalidate their function exports. This means it is unlikely most compilers
-will change this behavior.
+`GOOS=waspi1` from Go 1.21 does exit during `_start`. However, it doesn't
+support other exports besides `_start`, and `_start` is not defined to be
+called multiple times anyway.
 
-In summary, we return a `sys.ExitError` to the caller whenever we get it, as it
-properly reflects the state of the module, which would be closed on this error.
+Since `sys.ExitError` is not always returned, we added `Module.IsClosed` for
+defensive checks. This helps integrators avoid calling functions which will
+always fail.
 
 ### Why panic with `sys.ExitError` after a host function exits?
 

api/wasm.go

@@ -150,6 +150,10 @@ type Module interface {
 	Memory() Memory
 
 	// ExportedFunction returns a function exported from this module or nil if it wasn't.
+	//
+	// Note: The default wazero.ModuleConfig attempts to invoke `_start`, which
+	// in rare cases can close the module. When in doubt, check IsClosed prior
+	// to invoking a function export after instantiation.
 	ExportedFunction(name string) Function
 
 	// ExportedFunctionDefinitions returns all the exported function
@@ -190,6 +194,21 @@ type Module interface {
 	// Closer closes this module by delegating to CloseWithExitCode with an exit code of zero.
 	Closer
 
+	// IsClosed returns true if the module is closed, so no longer usable.
+	//
+	// This can happen for the following reasons:
+	//   - Closer was called directly.
+	//   - A guest function called Closer indirectly, such as `_start` calling
+	//     `proc_exit`, which internally closed the module.
+	//   - wazero.RuntimeConfig `WithCloseOnContextDone` was enabled and a
+	//     context completion closed the module.
+	//
+	// Where any of the above are possible, check this value before calling an
+	// ExportedFunction, even if you didn't formerly receive a sys.ExitError.
+	// sys.ExitError is only returned on non-zero code, something that closes
+	// the module successfully will not result it one.
+	IsClosed() bool
+
 	internalapi.WazeroOnly
 }
 

experimental/wazerotest/wazerotest.go

@@ -52,14 +52,17 @@ func NewModule(memory *Memory, functions ...*Function) *Module {
 	return &Module{Functions: functions, ExportMemory: memory}
 }
 
+// String implements fmt.Stringer.
 func (m *Module) String() string {
 	return "module[" + m.ModuleName + "]"
 }
 
+// Name implements the same method as documented on api.Module.
 func (m *Module) Name() string {
 	return m.ModuleName
 }
 
+// Memory implements the same method as documented on api.Module.
 func (m *Module) Memory() api.Memory {
 	if m.ExportMemory != nil {
 		m.once.Do(m.initialize)
@@ -68,16 +71,19 @@ func (m *Module) Memory() api.Memory {
 	return nil
 }
 
+// ExportedFunction implements the same method as documented on api.Module.
 func (m *Module) ExportedFunction(name string) api.Function {
 	m.once.Do(m.initialize)
 	return m.exportedFunctions[name]
 }
 
+// ExportedFunctionDefinitions implements the same method as documented on api.Module.
 func (m *Module) ExportedFunctionDefinitions() map[string]api.FunctionDefinition {
 	m.once.Do(m.initialize)
 	return m.exportedFunctionDefinitions
 }
 
+// ExportedMemory implements the same method as documented on api.Module.
 func (m *Module) ExportedMemory(name string) api.Memory {
 	if m.ExportMemory != nil && name == "memory" {
 		m.once.Do(m.initialize)
@@ -86,16 +92,48 @@ func (m *Module) ExportedMemory(name string) api.Memory {
 	return nil
 }
 
+// ExportedMemoryDefinitions implements the same method as documented on api.Module.
 func (m *Module) ExportedMemoryDefinitions() map[string]api.MemoryDefinition {
 	m.once.Do(m.initialize)
 	return m.exportedMemoryDefinitions
 }
 
+// ExportedGlobal implements the same method as documented on api.Module.
 func (m *Module) ExportedGlobal(name string) api.Global {
 	m.once.Do(m.initialize)
 	return m.exportedGlobals[name]
 }
 
+// Close implements the same method as documented on api.Closer.
+func (m *Module) Close(ctx context.Context) error {
+	return m.CloseWithExitCode(ctx, 0)
+}
+
+// CloseWithExitCode implements the same method as documented on api.Closer.
+func (m *Module) CloseWithExitCode(ctx context.Context, exitCode uint32) error {
+	atomic.CompareAndSwapUint64(&m.exitStatus, 0, exitStatusMarker|uint64(exitCode))
+	return nil
+}
+
+// IsClosed implements the same method as documented on api.Module.
+func (m *Module) IsClosed() bool {
+	_, exited := m.ExitStatus()
+	return exited
+}
+
+// NumGlobal implements the same method as documented on experimental.InternalModule.
+func (m *Module) NumGlobal() int {
+	return len(m.Globals)
+}
+
+// Global implements the same method as documented on experimental.InternalModule.
+func (m *Module) Global(i int) api.Global {
+	m.once.Do(m.initialize)
+	return m.Globals[i]
+}
+
+// Below are undocumented extensions
+
 func (m *Module) NumFunction() int {
 	return len(m.Functions)
 }
@@ -105,24 +143,6 @@ func (m *Module) Function(i int) api.Function {
 	return m.Functions[i]
 }
 
-func (m *Module) NumGlobal() int {
-	return len(m.Globals)
-}
-
-func (m *Module) Global(i int) api.Global {
-	m.once.Do(m.initialize)
-	return m.Globals[i]
-}
-
-func (m *Module) Close(ctx context.Context) error {
-	return m.CloseWithExitCode(ctx, 0)
-}
-
-func (m *Module) CloseWithExitCode(ctx context.Context, exitCode uint32) error {
-	atomic.CompareAndSwapUint64(&m.exitStatus, 0, exitStatusMarker|uint64(exitCode))
-	return nil
-}
-
 func (m *Module) ExitStatus() (exitCode uint32, exited bool) {
 	exitStatus := atomic.LoadUint64(&m.exitStatus)
 	return uint32(exitStatus), exitStatus != 0

internal/wasm/module_instance.go

@@ -106,6 +106,11 @@ func (m *ModuleInstance) CloseWithExitCode(ctx context.Context, exitCode uint32)
 	return m.ensureResourcesClosed(ctx)
 }
 
+// IsClosed implements the same method as documented on api.Module.
+func (m *ModuleInstance) IsClosed() bool {
+	return atomic.LoadUint64(&m.Closed) != 0
+}
+
 func (m *ModuleInstance) closeWithExitCodeWithoutClosingResource(exitCode uint32) (err error) {
 	if !m.setExitCode(exitCode, exitCodeFlagResourceNotClosed) {
 		return nil // not an error to have already closed

internal/wasm/module_instance_test.go

@@ -100,6 +100,9 @@ func TestModuleInstance_Close(t *testing.T) {
 
 				require.Equal(t, tc.expectedClosed, m.Closed)
 
+				// Outside callers should be able to know it was closed.
+				require.True(t, m.IsClosed())
+
 				// Verify our intended side-effect
 				require.Nil(t, s.Module(moduleName))
 

runtime.go

@@ -320,7 +320,8 @@ func (r *runtime) InstantiateModule(
 		return
 	}
 
-	// Attach the code closer so that anything afterwards closes the compiled code when closing the module.
+	// Attach the code closer so that anything afterward closes the compiled
+	// code when closing the module.
 	if code.closeWithModule {
 		mod.(*wasm.ModuleInstance).CodeCloser = code
 	}