Repository: tetratelabs/wazero
Pull request #844 (2022-11-07 18:44:21)

Replace module name check linear scan with map lookup
The check in module initialization to ensure a uniqueness of a module name in a a namespace does a linear scan of a string slice. Because the namespace needs to be fully locked during this check it means module initializations slows down as you have more active modules. This change replaces the linear scan with a map lookup. I've added a benchmark to track the performance of having multiple simiultainus module instances.

Before:
```
❯ go test -benchmem -memprofile memprofile.out -cpuprofile profile.out -blockprofile blockprofile.out -run=^$ -bench ^BenchmarkInitialization github.com/tetratelabs/wazero/internal/integration_test/bench
goos: darwin
goarch: amd64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
cpu: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
BenchmarkInitialization/interpreter-12             29836             40640 ns/op          142743 B/op        144 allocs/op
BenchmarkInitialization/interpreter-multiple-12    25544             83117 ns/op          143663 B/op        150 allocs/op
BenchmarkInitialization/compiler-12                31640             38680 ns/op          145935 B/op        128 allocs/op
BenchmarkInitialization/compiler-multiple-12       28410             72372 ns/op          146676 B/op        134 allocs/op
PASS
ok      github.com/tetratelabs/wazero/internal/integration_test/bench   8.945s
```

After:
```
❯ go test -benchmem -memprofile memprofile.out -cpuprofile profile.out -blockprofile blockprofile.out -run=^$ -bench ^BenchmarkInitialization github.com/tetratelabs/wazero/internal/integration_test/bench
goos: darwin
goarch: amd64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
cpu: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
BenchmarkInitialization/interpreter-12             29815             40490 ns/op          142742 B/op        144 allocs/op
BenchmarkInitialization/interpreter-multiple-12    49501             27282 ns/op          143808 B/op        150 allocs/op
BenchmarkInitialization/compiler-12                25664             39947 ns/op          145934 B/op        128 allocs/op
BenchmarkInitialization/compiler-multiple-12       48870             39809 ns/op          146652 B/op        134 allocs/op
PASS
ok      github.com/tetratelabs/wazero/internal/integration_test/bench   7.726s
```

I've batched in a few other micro optimizations but I don't expect any of them to make much of a meaningful difference.


Patch:
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 998668a3..967e5a30 100644
--- a/internal/wasm/module.go
+++ b/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)
 	}
diff --git a/internal/wasm/namespace.go b/internal/wasm/namespace.go
index 7c412d47..78d44992 100644
--- a/internal/wasm/namespace.go
+++ 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
 }
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index c8b936b3..ee3c74f7 100644
--- a/internal/wasm/store.go
+++ b/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


Test patch:
diff --git a/internal/integration_test/bench/bench_test.go b/internal/integration_test/bench/bench_test.go
index 3cb91917..58db9c76 100644
--- a/internal/integration_test/bench/bench_test.go
+++ b/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)
diff --git a/internal/wasm/namespace_test.go b/internal/wasm/namespace_test.go
index 2d8f1097..87ffb46f 100644
--- a/internal/wasm/namespace_test.go
+++ b/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
 }