Repository: tetratelabs/wazero
Pull request #952 (2022-12-21 05:08:05)

internal/gojs: reuse single wazero.Runtime to fix flake in tests
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/internal/engine/compiler/engine_cache.go b/internal/engine/compiler/engine_cache.go
index 65642bda..0e9ab719 100644
--- a/internal/engine/compiler/engine_cache.go
+++ b/internal/engine/compiler/engine_cache.go
@@ -53,7 +53,7 @@ func (e *engine) getCodesFromMemory(module *wasm.Module) (codes []*code, ok bool
 }
 
 func (e *engine) addCodesToCache(module *wasm.Module, codes []*code) (err error) {
-	if e.Cache == nil {
+	if e.Cache == nil || module.IsHostModule {
 		return
 	}
 	err = e.Cache.Add(module.ID, serializeCodes(e.wazeroVersion, codes))
@@ -61,7 +61,7 @@ func (e *engine) addCodesToCache(module *wasm.Module, codes []*code) (err error)
 }
 
 func (e *engine) getCodesFromCache(module *wasm.Module) (codes []*code, hit bool, err error) {
-	if e.Cache == nil {
+	if e.Cache == nil || module.IsHostModule {
 		return
 	}
 
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index 64cd2cf8..9736c735 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -104,7 +104,7 @@ func NewHostModule(
 		}
 	}
 
-	// Assigns the ModuleID by calculating sha256 on inputs as host modules do not have `wasm` to hash.
+	m.IsHostModule = true
 	m.AssignModuleID([]byte(fmt.Sprintf("%s:%v:%v", moduleName, nameToGoFunc, enabledFeatures)))
 	m.BuildFunctionDefinitions()
 	return
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 8183ba02..6d20d2ff 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -180,6 +180,9 @@ type Module struct {
 	// ID is the sha256 value of the source wasm and is used for caching.
 	ID ModuleID
 
+	// IsHostModule true if this is the host module, false otherwise.
+	IsHostModule bool
+
 	// FunctionDefinitionSection is a wazero-specific section built on Validate.
 	FunctionDefinitionSection []*FunctionDefinition
 


Test patch:
diff --git a/internal/engine/compiler/engine_cache_test.go b/internal/engine/compiler/engine_cache_test.go
index e13e2a31..8ef85dd6 100644
--- a/internal/engine/compiler/engine_cache_test.go
+++ b/internal/engine/compiler/engine_cache_test.go
@@ -2,6 +2,7 @@ package compiler
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -212,10 +213,26 @@ func TestDeserializeCodes(t *testing.T) {
 }
 
 func TestEngine_getCodesFromCache(t *testing.T) {
+	valid := concat(
+		[]byte(wazeroMagic),
+		[]byte{byte(len(testVersion))},
+		[]byte(testVersion),
+		u32.LeBytes(2), // number of functions.
+		// Function index = 0.
+		u64.LeBytes(12345),    // stack pointer ceil.
+		u64.LeBytes(5),        // length of code.
+		[]byte{1, 2, 3, 4, 5}, // code.
+		// Function index = 1.
+		u64.LeBytes(0xffffffff), // stack pointer ceil.
+		u64.LeBytes(3),          // length of code.
+		[]byte{1, 2, 3},         // code.
+	)
+
 	tests := []struct {
 		name       string
 		ext        *testCache
 		key        wasm.ModuleID
+		isHostMod  bool
 		expCodes   []*code
 		expHit     bool
 		expErr     string
@@ -226,6 +243,11 @@ func TestEngine_getCodesFromCache(t *testing.T) {
 			name: "not hit",
 			ext:  &testCache{caches: map[wasm.ModuleID][]byte{}},
 		},
+		{
+			name:      "host module",
+			ext:       &testCache{caches: map[wasm.ModuleID][]byte{{}: valid}},
+			isHostMod: true,
+		},
 		{
 			name:   "error in Cache.Get",
 			ext:    &testCache{caches: map[wasm.ModuleID][]byte{{}: {}}},
@@ -249,20 +271,7 @@ func TestEngine_getCodesFromCache(t *testing.T) {
 		{
 			name: "hit",
 			ext: &testCache{caches: map[wasm.ModuleID][]byte{
-				{}: concat(
-					[]byte(wazeroMagic),
-					[]byte{byte(len(testVersion))},
-					[]byte(testVersion),
-					u32.LeBytes(2), // number of functions.
-					// Function index = 0.
-					u64.LeBytes(12345),    // stack pointer ceil.
-					u64.LeBytes(5),        // length of code.
-					[]byte{1, 2, 3, 4, 5}, // code.
-					// Function index = 1.
-					u64.LeBytes(0xffffffff), // stack pointer ceil.
-					u64.LeBytes(3),          // length of code.
-					[]byte{1, 2, 3},         // code.
-				),
+				{}: valid,
 			}},
 			expHit: true,
 			expCodes: []*code{
@@ -275,7 +284,7 @@ func TestEngine_getCodesFromCache(t *testing.T) {
 	for _, tc := range tests {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			m := &wasm.Module{ID: tc.key}
+			m := &wasm.Module{ID: tc.key, IsHostModule: tc.isHostMod}
 			for _, expC := range tc.expCodes {
 				expC.sourceModule = m
 			}
@@ -308,6 +317,18 @@ func TestEngine_addCodesToCache(t *testing.T) {
 		err := e.addCodesToCache(nil, nil)
 		require.NoError(t, err)
 	})
+	t.Run("host module", func(t *testing.T) {
+		tc := &testCache{caches: map[wasm.ModuleID][]byte{}}
+		e := engine{Cache: tc}
+		codes := []*code{{stackPointerCeil: 123, codeSegment: []byte{1, 2, 3}}}
+		m := &wasm.Module{ID: sha256.Sum256(nil), IsHostModule: true} // Host module!
+		err := e.addCodesToCache(m, codes)
+		require.NoError(t, err)
+		// Check the host module not cached.
+		content, ok := tc.caches[m.ID]
+		require.False(t, ok)
+		require.Nil(t, content)
+	})
 	t.Run("add", func(t *testing.T) {
 		ext := &testCache{caches: map[wasm.ModuleID][]byte{}}
 		e := engine{Cache: ext}
diff --git a/internal/gojs/compiler_test.go b/internal/gojs/compiler_test.go
index 39e00668..ffaaf43b 100644
--- a/internal/gojs/compiler_test.go
+++ b/internal/gojs/compiler_test.go
@@ -24,18 +24,21 @@ import (
 func compileAndRun(ctx context.Context, arg string, config wazero.ModuleConfig) (stdout, stderr string, err error) {
 	var stdoutBuf, stderrBuf bytes.Buffer
 
-	r := wazero.NewRuntimeWithConfig(testCtx, wazero.NewRuntimeConfig())
-	defer r.Close(ctx)
-
-	gojs.MustInstantiate(ctx, r)
-
-	compiled, compileErr := r.CompileModule(ctx, testBin)
-	if compileErr != nil {
-		err = compileErr
+	ns := rt.NewNamespace(ctx)
+	builder := rt.NewHostModuleBuilder("go")
+	gojs.NewFunctionExporter().
+		ExportFunctions(builder)
+	if _, err = builder.Instantiate(ctx, ns); err != nil {
 		return
 	}
 
-	err = gojs.Run(ctx, r, compiled, config.WithStdout(&stdoutBuf).WithStderr(&stderrBuf).
+	// Note: this hits the file cache.
+	compiled, err := rt.CompileModule(testCtx, testBin)
+	if err != nil {
+		log.Panicln(err)
+	}
+
+	err = gojs.Run(ctx, ns, compiled, config.WithStdout(&stdoutBuf).WithStderr(&stderrBuf).
 		WithArgs("test", arg))
 	stdout = stdoutBuf.String()
 	stderr = stderrBuf.String()
@@ -54,6 +57,7 @@ var (
 		"sub":          {Mode: fs.ModeDir},
 		"sub/test.txt": {Data: []byte("greet sub dir\n")},
 	}
+	rt wazero.Runtime
 )
 
 func TestMain(m *testing.M) {
@@ -87,13 +91,17 @@ func TestMain(m *testing.M) {
 
 	// Seed wazero's compilation cache to see any error up-front and to prevent
 	// one test from a cache-miss performance penalty.
-	rt := wazero.NewRuntimeWithConfig(testCtx, wazero.NewRuntimeConfig())
-	defer rt.Close(testCtx)
+	rt = wazero.NewRuntimeWithConfig(testCtx, wazero.NewRuntimeConfig())
 	_, err = rt.CompileModule(testCtx, testBin)
 	if err != nil {
 		log.Panicln(err)
 	}
-	rt.Close(testCtx)
+
+	var exit int
+	defer func() {
+		rt.Close(testCtx)
+		os.Exit(exit)
+	}()
 
 	// Configure fs test data
 	if d, err := fs.Sub(testFS, "sub"); err != nil {
@@ -101,7 +109,7 @@ func TestMain(m *testing.M) {
 	} else if err = fstest.TestFS(d, "test.txt"); err != nil {
 		log.Panicln(err)
 	}
-	os.Exit(m.Run())
+	exit = m.Run()
 }
 
 // compileJsWasm allows us to generate a binary with runtime.GOOS=js and
diff --git a/internal/wasm/host_test.go b/internal/wasm/host_test.go
index 3ecf97e7..4d33e3b5 100644
--- a/internal/wasm/host_test.go
+++ b/internal/wasm/host_test.go
@@ -119,6 +119,7 @@ func TestNewHostModule(t *testing.T) {
 			m, e := NewHostModule(tc.moduleName, tc.nameToGoFunc, tc.funcToNames, api.CoreFeaturesV2)
 			require.NoError(t, e)
 			requireHostModuleEquals(t, tc.expected, m)
+			require.True(t, m.IsHostModule)
 		})
 	}
 }