Replace module name check linear scan with map lookup (#844)

Signed-off-by: Clifton Kaznocha <ckaznocha@users.noreply.github.com>

internal/integration_test/bench/bench_test.go

@@ -6,6 +6,8 @@ import (
 	_ "embed"
 	"fmt"
 	"runtime"
+	"strconv"
+	"sync"
 	"testing"
 
 	"github.com/tetratelabs/wazero"
@@ -44,11 +46,21 @@ func BenchmarkInitialization(b *testing.B) {
 		runInitializationBench(b, r)
 	})
 
+	b.Run("interpreter-multiple", func(b *testing.B) {
+		r := createRuntime(b, wazero.NewRuntimeConfigInterpreter())
+		runInitializationConcurrentBench(b, r)
+	})
+
 	if runtime.GOARCH == "amd64" || runtime.GOARCH == "arm64" {
 		b.Run("compiler", func(b *testing.B) {
 			r := createRuntime(b, wazero.NewRuntimeConfigCompiler())
 			runInitializationBench(b, r)
 		})
+
+		b.Run("compiler-multiple", func(b *testing.B) {
+			r := createRuntime(b, wazero.NewRuntimeConfigCompiler())
+			runInitializationConcurrentBench(b, r)
+		})
 	}
 }
 
@@ -102,6 +114,28 @@ func runInitializationBench(b *testing.B, r wazero.Runtime) {
 	}
 }
 
+func runInitializationConcurrentBench(b *testing.B, r wazero.Runtime) {
+	compiled := runCompilation(b, r)
+	defer compiled.Close(testCtx)
+	// Configure with real sources to avoid performance hit initializing fake ones. These sources are not used
+	// in the benchmark.
+	config := wazero.NewModuleConfig().WithSysNanotime().WithSysWalltime().WithRandSource(rand.Reader)
+	wg := &sync.WaitGroup{}
+	wg.Add(b.N)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		go func(name string) {
+			_, err := r.InstantiateModule(testCtx, compiled, config.WithName(name))
+			if err != nil {
+				b.Error(err)
+			}
+			wg.Done()
+		}(strconv.Itoa(i))
+	}
+	wg.Wait()
+	b.StopTimer()
+}
+
 func runAllInvocationBenches(b *testing.B, m api.Module) {
 	runBase64Benches(b, m)
 	runFibBenches(b, m)

internal/wasm/module.go

@@ -559,7 +559,8 @@ func (m *Module) validateDataCountSection() (err error) {
 }
 
 func (m *Module) buildGlobals(importedGlobals []*GlobalInstance) (globals []*GlobalInstance) {
-	for _, gs := range m.GlobalSection {
+	globals = make([]*GlobalInstance, len(m.GlobalSection))
+	for i, gs := range m.GlobalSection {
 		g := &GlobalInstance{Type: gs.Type}
 		switch v := executeConstExpression(importedGlobals, gs.Init).(type) {
 		case uint32:
@@ -577,7 +578,7 @@ func (m *Module) buildGlobals(importedGlobals []*GlobalInstance) (globals []*Glo
 		default:
 			panic(fmt.Errorf("BUG: invalid conversion %d", v))
 		}
-		globals = append(globals, g)
+		globals[i] = g
 	}
 	return
 }
@@ -707,9 +708,10 @@ func (f *FunctionType) key() string {
 		ret += ValueTypeName(b)
 	}
 	if len(f.Params) == 0 {
-		ret += "v"
+		ret += "v_"
+	} else {
+		ret += "_"
 	}
-	ret += "_"
 	for _, b := range f.Results {
 		ret += ValueTypeName(b)
 	}

internal/wasm/namespace.go

@@ -10,8 +10,11 @@ import (
 
 // Namespace is a collection of instantiated modules which cannot conflict on name.
 type Namespace struct {
-	// moduleNames ensures no race conditions instantiating two modules of the same name
-	moduleNames []string // guarded by mux
+	// moduleNamesList ensures modules are closed in reverse initialization order.
+	moduleNamesList []string // guarded by mux
+
+	// moduleNamesSet ensures no race conditions instantiating two modules of the same name
+	moduleNamesSet map[string]struct{} // guarded by mux
 
 	// modules holds the instantiated Wasm modules by module name from Instantiate.
 	modules map[string]*ModuleInstance // guarded by mux
@@ -23,8 +26,9 @@ type Namespace struct {
 // newNamespace returns an empty namespace.
 func newNamespace() *Namespace {
 	return &Namespace{
-		moduleNames: nil,
-		modules:     map[string]*ModuleInstance{},
+		moduleNamesList: nil,
+		moduleNamesSet:  map[string]struct{}{},
+		modules:         map[string]*ModuleInstance{},
 	}
 }
 
@@ -40,10 +44,11 @@ func (ns *Namespace) deleteModule(moduleName string) {
 	ns.mux.Lock()
 	defer ns.mux.Unlock()
 	delete(ns.modules, moduleName)
+	delete(ns.moduleNamesSet, moduleName)
 	// remove this module name
-	for i, n := range ns.moduleNames {
+	for i, n := range ns.moduleNamesList {
 		if n == moduleName {
-			ns.moduleNames = append(ns.moduleNames[:i], ns.moduleNames[i+1:]...)
+			ns.moduleNamesList = append(ns.moduleNamesList[:i], ns.moduleNamesList[i+1:]...)
 			break
 		}
 	}
@@ -58,10 +63,11 @@ func (ns *Namespace) module(moduleName string) *ModuleInstance {
 
 // requireModules returns all instantiated modules whose names equal the keys in the input, or errs if any are missing.
 func (ns *Namespace) requireModules(moduleNames map[string]struct{}) (map[string]*ModuleInstance, error) {
+	ret := make(map[string]*ModuleInstance, len(moduleNames))
+
 	ns.mux.RLock()
 	defer ns.mux.RUnlock()
 
-	ret := make(map[string]*ModuleInstance, len(moduleNames))
 	for n := range moduleNames {
 		m, ok := ns.modules[n]
 		if !ok {
@@ -77,12 +83,11 @@ func (ns *Namespace) requireModules(moduleNames map[string]struct{}) (map[string
 func (ns *Namespace) requireModuleName(moduleName string) error {
 	ns.mux.Lock()
 	defer ns.mux.Unlock()
-	for _, n := range ns.moduleNames {
-		if n == moduleName {
-			return fmt.Errorf("module[%s] has already been instantiated", moduleName)
-		}
+	if _, ok := ns.moduleNamesSet[moduleName]; ok {
+		return fmt.Errorf("module[%s] has already been instantiated", moduleName)
 	}
-	ns.moduleNames = append(ns.moduleNames, moduleName)
+	ns.moduleNamesSet[moduleName] = struct{}{}
+	ns.moduleNamesList = append(ns.moduleNamesList, moduleName)
 	return nil
 }
 
@@ -98,15 +103,16 @@ func (ns *Namespace) CloseWithExitCode(ctx context.Context, exitCode uint32) (er
 	ns.mux.Lock()
 	defer ns.mux.Unlock()
 	// Close modules in reverse initialization order.
-	for i := len(ns.moduleNames) - 1; i >= 0; i-- {
+	for i := len(ns.moduleNamesList) - 1; i >= 0; i-- {
 		// If closing this module errs, proceed anyway to close the others.
-		if m, ok := ns.modules[ns.moduleNames[i]]; ok {
+		if m, ok := ns.modules[ns.moduleNamesList[i]]; ok {
 			if _, e := m.CallCtx.close(ctx, exitCode); e != nil && err == nil {
 				err = e // first error
 			}
 		}
 	}
-	ns.moduleNames = nil
+	ns.moduleNamesList = nil
+	ns.moduleNamesSet = map[string]struct{}{}
 	ns.modules = map[string]*ModuleInstance{}
 	return
 }

internal/wasm/namespace_test.go

@@ -23,7 +23,8 @@ func TestNamespace_addModule(t *testing.T) {
 
 		require.Equal(t, map[string]*ModuleInstance{m1.Name: m1}, ns.modules)
 		// Doesn't affect module names
-		require.Nil(t, ns.moduleNames)
+		require.Zero(t, len(ns.moduleNamesSet))
+		require.Nil(t, ns.moduleNamesList)
 	})
 
 	t.Run("redundant ok", func(t *testing.T) {
@@ -46,7 +47,8 @@ func TestNamespace_deleteModule(t *testing.T) {
 
 		// Leaves the other module alone
 		require.Equal(t, map[string]*ModuleInstance{m1.Name: m1}, ns.modules)
-		require.Equal(t, []string{m1.Name}, ns.moduleNames)
+		require.Equal(t, map[string]struct{}{m1.Name: {}}, ns.moduleNamesSet)
+		require.Equal(t, []string{m1.Name}, ns.moduleNamesList)
 	})
 
 	t.Run("ok if missing", func(t *testing.T) {
@@ -57,7 +59,8 @@ func TestNamespace_deleteModule(t *testing.T) {
 		ns.deleteModule(m1.Name)
 
 		require.Zero(t, len(ns.modules))
-		require.Zero(t, len(ns.moduleNames))
+		require.Zero(t, len(ns.moduleNamesSet))
+		require.Zero(t, len(ns.moduleNamesList))
 	})
 }
 
@@ -90,14 +93,15 @@ func TestNamespace_requireModules(t *testing.T) {
 }
 
 func TestNamespace_requireModuleName(t *testing.T) {
-	ns := &Namespace{}
+	ns := &Namespace{moduleNamesSet: map[string]struct{}{}}
 
 	t.Run("first", func(t *testing.T) {
 		err := ns.requireModuleName("m1")
 		require.NoError(t, err)
 
 		// Ensure it adds the module name, and doesn't impact the module list.
-		require.Equal(t, []string{"m1"}, ns.moduleNames)
+		require.Equal(t, []string{"m1"}, ns.moduleNamesList)
+		require.Equal(t, map[string]struct{}{"m1": {}}, ns.moduleNamesSet)
 		require.Zero(t, len(ns.modules))
 	})
 	t.Run("second", func(t *testing.T) {
@@ -105,7 +109,8 @@ func TestNamespace_requireModuleName(t *testing.T) {
 		require.NoError(t, err)
 
 		// Appends in order.
-		require.Equal(t, []string{"m1", "m2"}, ns.moduleNames)
+		require.Equal(t, []string{"m1", "m2"}, ns.moduleNamesList)
+		require.Equal(t, map[string]struct{}{"m1": {}, "m2": {}}, ns.moduleNamesSet)
 	})
 	t.Run("existing", func(t *testing.T) {
 		err := ns.requireModuleName("m2")
@@ -121,7 +126,8 @@ func TestNamespace_AliasModule(t *testing.T) {
 	ns.AliasModule("m1", "m2")
 	require.Equal(t, map[string]*ModuleInstance{"m1": m1, "m2": m1}, ns.modules)
 	// Doesn't affect module names
-	require.Nil(t, ns.moduleNames)
+	require.Zero(t, len(ns.moduleNamesSet))
+	require.Nil(t, ns.moduleNamesList)
 }
 
 func TestNamespace_CloseWithExitCode(t *testing.T) {
@@ -158,7 +164,8 @@ func TestNamespace_CloseWithExitCode(t *testing.T) {
 
 			// Namespace state zeroed
 			require.Zero(t, len(ns.modules))
-			require.Zero(t, len(ns.moduleNames))
+			require.Zero(t, len(ns.moduleNamesSet))
+			require.Zero(t, len(ns.moduleNamesList))
 		})
 	}
 
@@ -183,7 +190,8 @@ func TestNamespace_CloseWithExitCode(t *testing.T) {
 
 		// Namespace state zeroed
 		require.Zero(t, len(ns.modules))
-		require.Zero(t, len(ns.moduleNames))
+		require.Zero(t, len(ns.moduleNamesSet))
+		require.Zero(t, len(ns.moduleNamesList))
 	})
 }
 
@@ -209,6 +217,7 @@ func newTestNamespace() (*Namespace, *ModuleInstance, *ModuleInstance) {
 	m2.CallCtx = NewCallContext(ns, m2, nil)
 
 	ns.modules = map[string]*ModuleInstance{m1.Name: m1, m2.Name: m2}
-	ns.moduleNames = []string{m1.Name, m2.Name}
+	ns.moduleNamesSet = map[string]struct{}{m1.Name: {}, m2.Name: {}}
+	ns.moduleNamesList = []string{m1.Name, m2.Name}
 	return ns, m1, m2
 }

internal/wasm/store.go

@@ -594,7 +594,7 @@ func (s *Store) getFunctionTypeIDs(ts []*FunctionType) ([]FunctionTypeID, error)
 }
 
 func (s *Store) getFunctionTypeID(t *FunctionType) (FunctionTypeID, error) {
-	key := t.String()
+	key := t.key()
 	s.mux.RLock()
 	id, ok := s.typeIDs[key]
 	s.mux.RUnlock()
@@ -605,11 +605,11 @@ func (s *Store) getFunctionTypeID(t *FunctionType) (FunctionTypeID, error) {
 		if id, ok = s.typeIDs[key]; ok {
 			return id, nil
 		}
-		l := uint32(len(s.typeIDs))
-		if l >= s.functionMaxTypes {
+		l := len(s.typeIDs)
+		if uint32(l) >= s.functionMaxTypes {
 			return 0, fmt.Errorf("too many function types in a store")
 		}
-		id = FunctionTypeID(len(s.typeIDs))
+		id = FunctionTypeID(l)
 		s.typeIDs[key] = id
 	}
 	return id, nil