Adds failing test for cache closed while function that panics is in-flight

internal/engine/compiler/engine.go

@@ -1049,9 +1049,10 @@ entry:
 			case builtinFunctionIndexFunctionListenerAfter:
 				ce.builtinFunctionFunctionListenerAfter(ctx, m, caller)
 			case builtinFunctionIndexCheckExitCode:
-				// Note: this operation must be done in Go, not native code. The reason is that
-				// native code cannot be preempted and that means it can block forever if there are not
-				// enough OS threads (which we don't have control over).
+				// Note: this operation must be done in Go, not native code.
+				// The reason is that native code cannot be preempted: it can
+				// block forever if there are not enough OS threads (we don't
+				// have control over OS threads).
 				if err := m.FailIfClosed(); err != nil {
 					panic(err)
 				}

internal/engine/compiler/engine_cache.go

@@ -20,7 +20,7 @@ func (e *engine) deleteCompiledModule(module *wasm.Module) {
 	delete(e.codes, module.ID)
 
 	// Note: we do not call e.Cache.Delete, as the lifetime of
-	// the content is up to the implementation of extencache.Cache interface.
+	// the content is up to the implementation of filecache.Cache interface.
 }
 
 func (e *engine) addCompiledModule(module *wasm.Module, cm *compiledModule, withGoFunc bool) (err error) {

internal/integration_test/engine/adhoc_test.go

@@ -11,10 +11,12 @@ import (
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/testing/binaryencoding"
 	"github.com/tetratelabs/wazero/internal/testing/proxy"
 	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasip1"
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/sys"
 )
@@ -69,6 +71,10 @@ func runAllTests(t *testing.T, tests map[string]func(t *testing.T, r wazero.Runt
 			testf(t, wazero.NewRuntimeWithConfig(testCtx, config))
 		})
 	}
+	t.Run("close cache before runtime", func(t *testing.T) {
+		t.Parallel()
+		testCloseCacheBeforeRuntime(t, config)
+	})
 }
 
 var (
@@ -90,6 +96,8 @@ var (
 	globalExtendWasm []byte
 	//go:embed testdata/infinite_loop.wasm
 	infiniteLoopWasm []byte
+	//go:embed testdata/exit_on_start.wasm
+	exitOnStartWasm []byte
 )
 
 func testEnsureTerminationOnClose(t *testing.T, r wazero.Runtime) {
@@ -142,11 +150,48 @@ func testEnsureTerminationOnClose(t *testing.T, r wazero.Runtime) {
 			require.NoError(t, module.CloseWithExitCode(context.Background(), 2))
 		}()
 		_, err = infinite.Call(context.Background())
-		require.Error(t, err)
-		require.Contains(t, err.Error(), "module closed with exit_code(2)")
+		require.EqualError(t, err, "module closed with exit_code(2)")
 	})
 }
 
+func testCloseCacheBeforeRuntime(t *testing.T, cfg wazero.RuntimeConfig) {
+	cache, err := wazero.NewCompilationCacheWithDir(t.TempDir())
+	require.NoError(t, err)
+	defer cache.Close(testCtx)
+
+	// Use the compilation cache
+	r := wazero.NewRuntimeWithConfig(testCtx, cfg.WithCompilationCache(cache))
+	defer cache.Close(testCtx)
+
+	wasiBuilder := r.NewHostModuleBuilder(wasi_snapshot_preview1.ModuleName)
+	wasi_snapshot_preview1.NewFunctionExporter().ExportFunctions(wasiBuilder)
+
+	// Add a hacked proc_exit which simulates an out-of-order concern.
+	//
+	// Currently, locks apply to api.Module and api.CompiledModule. However,
+	// the file cache doesn't use locks and also call sites aren't prepared to
+	// check nil on everything that could be invalidated if the cache was
+	// ripped out from underneath the runtime. It may be a better idea to make
+	// cache.Close fail if it is still in use somehow.
+	wasiBuilder.NewFunctionBuilder().
+		WithGoModuleFunction(api.GoModuleFunc(func(ctx context.Context, mod api.Module, params []uint64) {
+			// Close the cache while inside this function.
+			require.NoError(t, cache.Close(testCtx))
+
+			// Below is the normal proc_exit code:
+			exitCode := uint32(params[0])
+			_ = mod.CloseWithExitCode(ctx, exitCode)
+			panic(sys.NewExitError(exitCode))
+			// ^-- results in a host panic looking up the function definition!
+		}), []api.ValueType{i32}, nil).Export(wasip1.ProcExitName)
+	_, err = wasiBuilder.Instantiate(testCtx)
+	require.NoError(t, err)
+
+	// We expect a normal exit with code, not a panic.
+	_, err = r.Instantiate(testCtx, exitOnStartWasm)
+	require.EqualError(t, err, "module closed with exit_code(2)")
+}
+
 func testUserDefinedPrimitiveHostFunc(t *testing.T, r wazero.Runtime) {
 	type u32 uint32
 	type u64 uint64

internal/integration_test/engine/testdata/exit_on_start.wat

@@ -0,0 +1,9 @@
+(module $exit_on_start
+  (import "wasi_snapshot_preview1" "proc_exit"
+    (func $wasi.proc_exit (param $rval i32)))
+
+  (func (export "_start")
+     i32.const 2           ;; push $rval onto the stack
+     call $wasi.proc_exit  ;; return a sys.ExitError to the caller
+  )
+)

internal/wasm/module_instance.go

@@ -16,8 +16,9 @@ func (m *ModuleInstance) FailIfClosed() (err error) {
 		switch closed & exitCodeFlagMask {
 		case exitCodeFlagResourceClosed:
 		case exitCodeFlagResourceNotClosed:
-			// This happens when this module is closed asynchronously in CloseModuleOnCanceledOrTimeout,
-			// and the closure of resources have been deferred here.
+			// This happens when this module is closed asynchronously in
+			// CloseModuleOnCanceledOrTimeout, and the closure of resources
+			// have been deferred here.
 			_ = m.ensureResourcesClosed(context.Background())
 		}
 		return sys.NewExitError(uint32(closed >> 32)) // Unpack the high order bits as the exit code.
@@ -148,13 +149,13 @@ func (m *ModuleInstance) ensureResourcesClosed(ctx context.Context) (err error)
 		m.Sys = nil
 	}
 
-	if m.CodeCloser == nil {
+	if codeCloser := m.CodeCloser; codeCloser != nil {
+		if err = codeCloser.Close(ctx); err != nil {
+			return err
+		}
+		m.CodeCloser = nil
 		return
 	}
-	if e := m.CodeCloser.Close(ctx); e != nil && err == nil {
-		err = e
-	}
-	m.CodeCloser = nil
 	return
 }
 

runtime.go

@@ -319,7 +319,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
 	}