Repository: tetratelabs/wazero
Pull request #1381 (2023-04-19 01:38:26)

cache: fixes consistency with CloseOnContextDone and listeners
This fixes the compilation cache consistency with respect to the 
listeners and CloseOnContextDone configurations by including them
in the wasm.Module.ID.


Patch:
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 07307bb3..7370824b 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -507,7 +507,7 @@ func (e *engine) Close() (err error) {
 
 // CompileModule implements the same method as documented on wasm.Engine.
 func (e *engine) CompileModule(_ context.Context, module *wasm.Module, listeners []experimental.FunctionListener, ensureTermination bool) error {
-	if _, ok, err := e.getCodes(module); ok { // cache hit!
+	if _, ok, err := e.getCodes(module, listeners); ok { // cache hit!
 		return nil
 	} else if err != nil {
 		return err
@@ -571,7 +571,10 @@ func (e *engine) NewModuleEngine(module *wasm.Module, instance *wasm.ModuleInsta
 
 	// Note: imported functions are resolved in moduleEngine.ResolveImportedFunction.
 
-	codes, ok, err := e.getCodes(module)
+	codes, ok, err := e.getCodes(module,
+		// listeners arg is not needed here since NewModuleEngine is called after CompileModule which
+		// ensures the association of listener with *code.
+		nil)
 	if !ok {
 		return nil, errors.New("source module must be compiled before instantiation")
 	} else if err != nil {
diff --git a/internal/engine/compiler/engine_cache.go b/internal/engine/compiler/engine_cache.go
index d9719751..1fda9598 100644
--- a/internal/engine/compiler/engine_cache.go
+++ b/internal/engine/compiler/engine_cache.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"runtime"
 
+	"github.com/tetratelabs/wazero/experimental"
 	"github.com/tetratelabs/wazero/internal/platform"
 	"github.com/tetratelabs/wazero/internal/u32"
 	"github.com/tetratelabs/wazero/internal/u64"
@@ -30,7 +31,7 @@ func (e *engine) addCodes(module *wasm.Module, codes []*code, withGoFunc bool) (
 	return
 }
 
-func (e *engine) getCodes(module *wasm.Module) (codes []*code, ok bool, err error) {
+func (e *engine) getCodes(module *wasm.Module, listeners []experimental.FunctionListener) (codes []*code, ok bool, err error) {
 	codes, ok = e.getCodesFromMemory(module)
 	if ok {
 		return
@@ -39,6 +40,13 @@ func (e *engine) getCodes(module *wasm.Module) (codes []*code, ok bool, err erro
 	if ok {
 		e.addCodesToMemory(module, codes)
 	}
+
+	if len(listeners) > 0 {
+		// Files do not contain the actual listener instances (it's impossible to cache them as files!), so assign each here.
+		for i := range codes {
+			codes[i].listener = listeners[i]
+		}
+	}
 	return
 }
 
@@ -104,6 +112,11 @@ func serializeCodes(wazeroVersion string, codes []*code) io.Reader {
 	buf.WriteByte(byte(len(wazeroVersion)))
 	// Version of wazero.
 	buf.WriteString(wazeroVersion)
+	if len(codes) > 0 && codes[0].withEnsureTermination {
+		buf.WriteByte(1)
+	} else {
+		buf.WriteByte(0)
+	}
 	// Number of *code (== locally defined functions in the module): 4 bytes.
 	buf.Write(u32.LeBytes(uint32(len(codes))))
 	for _, c := range codes {
@@ -119,7 +132,7 @@ func serializeCodes(wazeroVersion string, codes []*code) io.Reader {
 
 func deserializeCodes(wazeroVersion string, reader io.ReadCloser) (codes []*code, staleCache bool, err error) {
 	defer reader.Close()
-	cacheHeaderSize := len(wazeroMagic) + 1 /* version size */ + len(wazeroVersion) + 4 /* number of functions */
+	cacheHeaderSize := len(wazeroMagic) + 1 /* version size */ + len(wazeroVersion) + 1 /* ensure termination */ + 4 /* number of functions */
 
 	// Read the header before the native code.
 	header := make([]byte, cacheHeaderSize)
@@ -144,13 +157,15 @@ func deserializeCodes(wazeroVersion string, reader io.ReadCloser) (codes []*code
 		return
 	}
 
+	ensureTermination := header[cachedVersionEnd] != 0
+
 	functionsNum := binary.LittleEndian.Uint32(header[len(header)-4:])
 	codes = make([]*code, 0, functionsNum)
 
 	var eightBytes [8]byte
 	var nativeCodeLen uint64
 	for i := uint32(0); i < functionsNum; i++ {
-		c := &code{}
+		c := &code{withEnsureTermination: ensureTermination}
 
 		// Read the stack pointer ceil.
 		if c.stackPointerCeil, err = readUint64(reader, &eightBytes); err != nil {
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index cb4ca65e..400bcea3 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -71,7 +71,8 @@ func NewHostModule(
 	// Wasm codes for Wasm-implemented host functions) are not available and compiles each time. On the other hand,
 	// compilation of host modules is not costly as it's merely small trampolines vs the real-world native Wasm binary.
 	// TODO: refactor engines so that we can properly cache compiled machine codes for host modules.
-	m.AssignModuleID([]byte(fmt.Sprintf("@@@@@@@@%p", m))) // @@@@@@@@ = any 8 bytes different from Wasm header.
+	m.AssignModuleID([]byte(fmt.Sprintf("@@@@@@@@%p", m)), // @@@@@@@@ = any 8 bytes different from Wasm header.
+		false, false)
 	m.BuildFunctionDefinitions()
 	return
 }
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 8cbd8555..57063bec 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -162,7 +162,8 @@ type Module struct {
 	// See https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/appendix/changes.html#bulk-memory-and-table-instructions
 	DataCountSection *uint32
 
-	// ID is the sha256 value of the source wasm and is used for caching.
+	// ID is the sha256 value of the source wasm plus the configurations which affect the runtime representation of
+	// Wasm binary. This is only used for caching.
 	ID ModuleID
 
 	// IsHostModule true if this is the host module, false otherwise.
@@ -191,9 +192,24 @@ const (
 	MaximumTableIndex    = uint32(1 << 27)
 )
 
-// AssignModuleID calculates a sha256 checksum on `wasm` and set Module.ID to the result.
-func (m *Module) AssignModuleID(wasm []byte) {
-	m.ID = sha256.Sum256(wasm)
+// AssignModuleID calculates a sha256 checksum on `wasm` and other args, and set Module.ID to the result.
+// See the doc on Module.ID on what it's used for.
+func (m *Module) AssignModuleID(wasm []byte, withListener, withEnsureTermination bool) {
+	h := sha256.New()
+	h.Write(wasm)
+	// Use the pre-allocated space on m.ID to append the booleans to sha256 hash.
+	m.ID[0] = boolToByte(withListener)
+	m.ID[1] = boolToByte(withEnsureTermination)
+	h.Write(m.ID[:2])
+	// Get checksum by passing the slice underlying m.ID.
+	h.Sum(m.ID[:0])
+}
+
+func boolToByte(b bool) (ret byte) {
+	if b {
+		ret = 1
+	}
+	return
 }
 
 // TypeOfFunction returns the wasm.SectionIDType index for the given function space index or nil.
diff --git a/runtime.go b/runtime.go
index 475938f0..d2a72f60 100644
--- a/runtime.go
+++ b/runtime.go
@@ -2,7 +2,6 @@ package wazero
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"sync/atomic"
 
@@ -206,10 +205,6 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 		return nil, err
 	}
 
-	if binary == nil {
-		return nil, errors.New("binary == nil")
-	}
-
 	internal, err := binaryformat.DecodeModule(binary, r.enabledFeatures,
 		r.memoryLimitPages, r.memoryCapacityFromMax, !r.dwarfDisabled, r.storeCustomSections)
 	if err != nil {
@@ -220,8 +215,6 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 		return nil, err
 	}
 
-	internal.AssignModuleID(binary)
-
 	// Now that the module is validated, cache the function and memory definitions.
 	internal.BuildFunctionDefinitions()
 	internal.BuildMemoryDefinitions()
@@ -239,7 +232,7 @@ func (r *runtime) CompileModule(ctx context.Context, binary []byte) (CompiledMod
 	if err != nil {
 		return nil, err
 	}
-
+	internal.AssignModuleID(binary, len(listeners) > 0, r.ensureTermination)
 	if err = r.store.Engine.CompileModule(ctx, internal, listeners, r.ensureTermination); err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/internal/engine/compiler/engine_cache_test.go b/internal/engine/compiler/engine_cache_test.go
index 513badab..1fdb27a4 100644
--- a/internal/engine/compiler/engine_cache_test.go
+++ b/internal/engine/compiler/engine_cache_test.go
@@ -37,6 +37,20 @@ func TestSerializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},             // ensure termination.
+				u32.LeBytes(1),        // number of functions.
+				u64.LeBytes(12345),    // stack pointer ceil.
+				u64.LeBytes(5),        // length of code.
+				[]byte{1, 2, 3, 4, 5}, // code.
+			),
+		},
+		{
+			in: []*code{{withEnsureTermination: true, stackPointerCeil: 12345, codeSegment: []byte{1, 2, 3, 4, 5}}},
+			exp: concat(
+				[]byte(wazeroMagic),
+				[]byte{byte(len(testVersion))},
+				[]byte(testVersion),
+				[]byte{1},             // ensure termination.
 				u32.LeBytes(1),        // number of functions.
 				u64.LeBytes(12345),    // stack pointer ceil.
 				u64.LeBytes(5),        // length of code.
@@ -52,6 +66,7 @@ func TestSerializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},      // ensure termination.
 				u32.LeBytes(2), // number of functions.
 				// Function index = 0.
 				u64.LeBytes(12345),    // stack pointer ceil.
@@ -111,6 +126,7 @@ func TestDeserializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},             // ensure termination.
 				u32.LeBytes(1),        // number of functions.
 				u64.LeBytes(12345),    // stack pointer ceil.
 				u64.LeBytes(5),        // length of code.
@@ -122,12 +138,31 @@ func TestDeserializeCodes(t *testing.T) {
 			expStaleCache: false,
 			expErr:        "",
 		},
+		{
+			name: "one function with ensure termination",
+			in: concat(
+				[]byte(wazeroMagic),
+				[]byte{byte(len(testVersion))},
+				[]byte(testVersion),
+				[]byte{1},             // ensure termination.
+				u32.LeBytes(1),        // number of functions.
+				u64.LeBytes(12345),    // stack pointer ceil.
+				u64.LeBytes(5),        // length of code.
+				[]byte{1, 2, 3, 4, 5}, // code.
+			),
+			expCodes: []*code{
+				{stackPointerCeil: 12345, codeSegment: []byte{1, 2, 3, 4, 5}, withEnsureTermination: true},
+			},
+			expStaleCache: false,
+			expErr:        "",
+		},
 		{
 			name: "two functions",
 			in: concat(
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},      // ensure termination.
 				u32.LeBytes(2), // number of functions.
 				// Function index = 0.
 				u64.LeBytes(12345),    // stack pointer ceil.
@@ -151,6 +186,7 @@ func TestDeserializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},      // ensure termination.
 				u32.LeBytes(2), // number of functions.
 				// Function index = 0.
 				u64.LeBytes(12345),    // stack pointer ceil.
@@ -166,6 +202,7 @@ func TestDeserializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},      // ensure termination.
 				u32.LeBytes(2), // number of functions.
 				// Function index = 0.
 				u64.LeBytes(12345),    // stack pointer ceil.
@@ -182,6 +219,7 @@ func TestDeserializeCodes(t *testing.T) {
 				[]byte(wazeroMagic),
 				[]byte{byte(len(testVersion))},
 				[]byte(testVersion),
+				[]byte{0},      // ensure termination.
 				u32.LeBytes(2), // number of functions.
 				// Function index = 0.
 				u64.LeBytes(12345),    // stack pointer ceil.
@@ -217,6 +255,7 @@ func TestEngine_getCodesFromCache(t *testing.T) {
 		[]byte(wazeroMagic),
 		[]byte{byte(len(testVersion))},
 		[]byte(testVersion),
+		[]byte{0},      // ensure termination.
 		u32.LeBytes(2), // number of functions.
 		// Function index = 0.
 		u64.LeBytes(12345),    // stack pointer ceil.
@@ -353,6 +392,7 @@ func TestEngine_addCodesToCache(t *testing.T) {
 			[]byte(wazeroMagic),
 			[]byte{byte(len(testVersion))},
 			[]byte(testVersion),
+			[]byte{0},
 			u32.LeBytes(1),   // number of functions.
 			u64.LeBytes(123), // stack pointer ceil.
 			u64.LeBytes(3),   // length of code.
diff --git a/internal/integration_test/README.md b/internal/integration_test/README.md
deleted file mode 100644
index 03e719d2..00000000
--- a/internal/integration_test/README.md
+++ /dev/null
@@ -1,11 +0,0 @@
-This directory contains tests which use multiple packages. For example:
-
-* `bench` contains benchmark tests.
-* `engine` contains variety of end-to-end tests, mainly to ensure the consistency in the behavior between engines.
-* `fuzzcases` contains variety of test cases found by the [fuzz](./fuzz) testing.
-* `post1_0` contains end-to-end tests for features [finished](https://github.com/WebAssembly/proposals/blob/main/finished-proposals.md) after WebAssembly 1.0 (20191205).
-* `spectest` contains end-to-end tests with the [WebAssembly specification tests](https://github.com/WebAssembly/spec/tree/wg-1.0/test/core).
-* `vs` tests and benchmarks VS other WebAssembly runtimes.
-
-*Note*: This doesn't contain WASI tests, as there's not yet an [official testsuite](https://github.com/WebAssembly/WASI/issues/9).
-Meanwhile, WASI functions are unit tested including via Text Format imports [here](../../wasi/wasi_test.go)
diff --git a/internal/integration_test/filecache/filecache_test.go b/internal/integration_test/filecache/filecache_test.go
index 1fb61fd1..bdd57d6d 100644
--- a/internal/integration_test/filecache/filecache_test.go
+++ b/internal/integration_test/filecache/filecache_test.go
@@ -8,13 +8,19 @@ import (
 	"os/exec"
 	"strings"
 	"testing"
+	"time"
 
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/experimental/logging"
 	"github.com/tetratelabs/wazero/internal/engine/compiler"
 	"github.com/tetratelabs/wazero/internal/filecache"
 	"github.com/tetratelabs/wazero/internal/integration_test/spectest"
 	v1 "github.com/tetratelabs/wazero/internal/integration_test/spectest/v1"
 	"github.com/tetratelabs/wazero/internal/platform"
+	"github.com/tetratelabs/wazero/internal/testing/binaryencoding"
 	"github.com/tetratelabs/wazero/internal/testing/require"
+	"github.com/tetratelabs/wazero/internal/wasm"
 )
 
 func TestSpecTestCompilerCache(t *testing.T) {
@@ -64,3 +70,220 @@ func TestSpecTestCompilerCache(t *testing.T) {
 		spectest.Run(t, v1.Testcases, context.Background(), fc, compiler.NewEngine, v1.EnabledFeatures)
 	}
 }
+
+// TestListeners ensures that compilation cache works as expected on and off with respect to listeners.
+func TestListeners(t *testing.T) {
+	if !platform.CompilerSupported() {
+		t.Skip()
+	}
+
+	var (
+		zero    uint32 = 0
+		wasmBin        = binaryencoding.EncodeModule(&wasm.Module{
+			TypeSection:     []wasm.FunctionType{{}},
+			FunctionSection: []wasm.Index{0, 0, 0, 0},
+			CodeSection: []wasm.Code{
+				{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}},
+				{Body: []byte{wasm.OpcodeCall, 2, wasm.OpcodeEnd}},
+				{Body: []byte{wasm.OpcodeCall, 3, wasm.OpcodeEnd}},
+				{Body: []byte{wasm.OpcodeEnd}},
+			},
+			StartSection: &zero,
+			NameSection: &wasm.NameSection{
+				FunctionNames: wasm.NameMap{{Index: 0, Name: "1"}, {Index: 1, Name: "2"}, {Index: 2, Name: "3"}, {Index: 3, Name: "4"}},
+				ModuleName:    "test",
+			},
+		})
+	)
+
+	t.Run("always on", func(t *testing.T) {
+		dir := t.TempDir()
+
+		out := bytes.NewBuffer(nil)
+		ctxWithListener := context.WithValue(context.Background(),
+			experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+
+		{
+			cc, err := wazero.NewCompilationCacheWithDir(dir)
+			require.NoError(t, err)
+			rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+
+			r := wazero.NewRuntimeWithConfig(ctxWithListener, rc)
+			_, err = r.CompileModule(ctxWithListener, wasmBin)
+			require.NoError(t, err)
+			err = r.Close(ctxWithListener)
+			require.NoError(t, err)
+		}
+
+		cc, err := wazero.NewCompilationCacheWithDir(dir)
+		require.NoError(t, err)
+		rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+		r := wazero.NewRuntimeWithConfig(ctxWithListener, rc)
+		_, err = r.Instantiate(ctxWithListener, wasmBin)
+		require.NoError(t, err)
+		err = r.Close(ctxWithListener)
+		require.NoError(t, err)
+
+		// Ensures that compilation cache works with listeners.
+		require.Equal(t, `--> test.1()
+	--> test.2()
+		--> test.3()
+			--> test.4()
+			<--
+		<--
+	<--
+<--
+`, out.String())
+	})
+
+	t.Run("with->without", func(t *testing.T) {
+		dir := t.TempDir()
+
+		// Compile with listeners.
+		{
+			cc, err := wazero.NewCompilationCacheWithDir(dir)
+			require.NoError(t, err)
+			rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+
+			out := bytes.NewBuffer(nil)
+			ctxWithListener := context.WithValue(context.Background(),
+				experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+			r := wazero.NewRuntimeWithConfig(ctxWithListener, rc)
+			_, err = r.CompileModule(ctxWithListener, wasmBin)
+			require.NoError(t, err)
+			err = r.Close(ctxWithListener)
+			require.NoError(t, err)
+		}
+
+		// Then compile without listeners -> run it.
+		cc, err := wazero.NewCompilationCacheWithDir(dir)
+		require.NoError(t, err)
+		rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+		r := wazero.NewRuntimeWithConfig(context.Background(), rc)
+		_, err = r.Instantiate(context.Background(), wasmBin)
+		require.NoError(t, err)
+		err = r.Close(context.Background())
+		require.NoError(t, err)
+	})
+
+	t.Run("without->with", func(t *testing.T) {
+		dir := t.TempDir()
+
+		// Compile without listeners.
+		{
+			cc, err := wazero.NewCompilationCacheWithDir(dir)
+			require.NoError(t, err)
+			rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+			r := wazero.NewRuntimeWithConfig(context.Background(), rc)
+			_, err = r.CompileModule(context.Background(), wasmBin)
+			require.NoError(t, err)
+			err = r.Close(context.Background())
+			require.NoError(t, err)
+		}
+
+		// Then compile with listeners -> run it.
+		out := bytes.NewBuffer(nil)
+		ctxWithListener := context.WithValue(context.Background(),
+			experimental.FunctionListenerFactoryKey{}, logging.NewLoggingListenerFactory(out))
+
+		cc, err := wazero.NewCompilationCacheWithDir(dir)
+		require.NoError(t, err)
+		rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc)
+		r := wazero.NewRuntimeWithConfig(ctxWithListener, rc)
+		_, err = r.Instantiate(ctxWithListener, wasmBin)
+		require.NoError(t, err)
+		err = r.Close(ctxWithListener)
+		require.NoError(t, err)
+
+		// Ensures that compilation cache works with listeners.
+		require.Equal(t, `--> test.1()
+	--> test.2()
+		--> test.3()
+			--> test.4()
+			<--
+		<--
+	<--
+<--
+`, out.String())
+	})
+}
+
+// TestWithCloseOnContextDone ensures that compilation cache works as expected on and off with respect to WithCloseOnContextDone config.
+func TestWithCloseOnContextDone(t *testing.T) {
+	if !platform.CompilerSupported() {
+		t.Skip()
+	}
+
+	var (
+		zero    uint32 = 0
+		wasmBin        = binaryencoding.EncodeModule(&wasm.Module{
+			TypeSection:     []wasm.FunctionType{{}},
+			FunctionSection: []wasm.Index{0},
+			CodeSection: []wasm.Code{
+				{Body: []byte{
+					wasm.OpcodeLoop, 0,
+					wasm.OpcodeBr, 0,
+					wasm.OpcodeEnd,
+					wasm.OpcodeEnd,
+				}},
+			},
+			StartSection: &zero,
+		})
+	)
+
+	t.Run("always on", func(t *testing.T) {
+		dir := t.TempDir()
+		ctx := context.Background()
+		{
+			cc, err := wazero.NewCompilationCacheWithDir(dir)
+			require.NoError(t, err)
+			rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc).WithCloseOnContextDone(true)
+
+			r := wazero.NewRuntimeWithConfig(ctx, rc)
+			_, err = r.CompileModule(ctx, wasmBin)
+			require.NoError(t, err)
+			err = r.Close(ctx)
+			require.NoError(t, err)
+		}
+
+		cc, err := wazero.NewCompilationCacheWithDir(dir)
+		require.NoError(t, err)
+		rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc).WithCloseOnContextDone(true)
+		r := wazero.NewRuntimeWithConfig(ctx, rc)
+
+		timeoutCtx, done := context.WithTimeout(ctx, time.Second)
+		defer done()
+		_, err = r.Instantiate(timeoutCtx, wasmBin)
+		require.EqualError(t, err, "module closed with context deadline exceeded")
+		err = r.Close(ctx)
+		require.NoError(t, err)
+	})
+
+	t.Run("off->on", func(t *testing.T) {
+		dir := t.TempDir()
+		ctx := context.Background()
+		{
+			cc, err := wazero.NewCompilationCacheWithDir(dir)
+			require.NoError(t, err)
+			rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc).WithCloseOnContextDone(false)
+
+			r := wazero.NewRuntimeWithConfig(ctx, rc)
+			_, err = r.CompileModule(ctx, wasmBin)
+			require.NoError(t, err)
+			err = r.Close(ctx)
+			require.NoError(t, err)
+		}
+
+		cc, err := wazero.NewCompilationCacheWithDir(dir)
+		require.NoError(t, err)
+		rc := wazero.NewRuntimeConfigCompiler().WithCompilationCache(cc).WithCloseOnContextDone(true)
+		r := wazero.NewRuntimeWithConfig(ctx, rc)
+
+		timeoutCtx, done := context.WithTimeout(ctx, time.Second)
+		defer done()
+		_, err = r.Instantiate(timeoutCtx, wasmBin)
+		require.EqualError(t, err, "module closed with context deadline exceeded")
+		err = r.Close(ctx)
+		require.NoError(t, err)
+	})
+}
diff --git a/internal/integration_test/spectest/spectest.go b/internal/integration_test/spectest/spectest.go
index 4415e59a..144a44a1 100644
--- a/internal/integration_test/spectest/spectest.go
+++ b/internal/integration_test/spectest/spectest.go
@@ -393,7 +393,7 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 						mod, err := binaryformat.DecodeModule(buf, enabledFeatures, wasm.MemoryLimitPages, false, false, false)
 						require.NoError(t, err, msg)
 						require.NoError(t, mod.Validate(enabledFeatures))
-						mod.AssignModuleID(buf)
+						mod.AssignModuleID(buf, false, false)
 
 						moduleName := c.Name
 						if moduleName == "" {
@@ -541,7 +541,7 @@ func Run(t *testing.T, testDataFS embed.FS, ctx context.Context, fc filecache.Ca
 							err = mod.Validate(s.EnabledFeatures)
 							require.NoError(t, err, msg)
 
-							mod.AssignModuleID(buf)
+							mod.AssignModuleID(buf, false, false)
 
 							maybeSetMemoryCap(mod)
 							mod.BuildFunctionDefinitions()
@@ -577,7 +577,7 @@ func requireInstantiationError(t *testing.T, ctx context.Context, s *wasm.Store,
 		return
 	}
 
-	mod.AssignModuleID(buf)
+	mod.AssignModuleID(buf, false, false)
 
 	maybeSetMemoryCap(mod)
 	mod.BuildMemoryDefinitions()
diff --git a/internal/wasm/module_test.go b/internal/wasm/module_test.go
index aaecb2d5..dadb2bda 100644
--- a/internal/wasm/module_test.go
+++ b/internal/wasm/module_test.go
@@ -1007,3 +1007,32 @@ func TestModule_declaredFunctionIndexes(t *testing.T) {
 		})
 	}
 }
+
+func TestModule_AssignModuleID(t *testing.T) {
+	getID := func(bin []byte, withListener, withEnsureTermination bool) ModuleID {
+		m := Module{}
+		m.AssignModuleID(bin, withListener, withEnsureTermination)
+		return m.ID
+	}
+
+	// Ensures that different args always produce the different IDs.
+	exists := map[ModuleID]struct{}{}
+	for _, tc := range []struct {
+		bin                                 []byte
+		withListener, withEnsureTermination bool
+	}{
+		{bin: []byte{1, 2, 3}, withListener: false, withEnsureTermination: false},
+		{bin: []byte{1, 2, 3}, withListener: false, withEnsureTermination: true},
+		{bin: []byte{1, 2, 3}, withListener: true, withEnsureTermination: false},
+		{bin: []byte{1, 2, 3}, withListener: true, withEnsureTermination: true},
+		{bin: []byte{1, 2, 3, 4}, withListener: false, withEnsureTermination: false},
+		{bin: []byte{1, 2, 3, 4}, withListener: false, withEnsureTermination: true},
+		{bin: []byte{1, 2, 3, 4}, withListener: true, withEnsureTermination: false},
+		{bin: []byte{1, 2, 3, 4}, withListener: true, withEnsureTermination: true},
+	} {
+		id := getID(tc.bin, tc.withListener, tc.withEnsureTermination)
+		_, exist := exists[id]
+		require.False(t, exist)
+		exists[id] = struct{}{}
+	}
+}
diff --git a/runtime_test.go b/runtime_test.go
index 6ee0af77..b7fce389 100644
--- a/runtime_test.go
+++ b/runtime_test.go
@@ -159,7 +159,7 @@ func TestRuntime_CompileModule_Errors(t *testing.T) {
 	}{
 		{
 			name:        "nil",
-			expectedErr: "binary == nil",
+			expectedErr: "invalid magic number",
 		},
 		{
 			name:        "invalid binary",