Disable comp cache for host modules (#949)

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

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
 	}
 

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(
+				{}: valid,
-					[]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.
-				),
 			}},
 			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}

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

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)
 		})
 	}
 }

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