Reduce lock contention during module close (#900)

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

internal/integration_test/bench/bench_test.go

@@ -125,10 +125,11 @@ func runInitializationConcurrentBench(b *testing.B, r wazero.Runtime) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		go func(name string) {
-			_, err := r.InstantiateModule(testCtx, compiled, config.WithName(name))
+			m, err := r.InstantiateModule(testCtx, compiled, config.WithName(name))
 			if err != nil {
 				b.Error(err)
 			}
+			m.Close(testCtx)
 			wg.Done()
 		}(strconv.Itoa(i))
 	}

internal/wasm/namespace.go

@@ -8,13 +8,19 @@ import (
 	"github.com/tetratelabs/wazero/api"
 )
 
+// nameListNode is a node in a doubly linked list of names.
+type nameListNode struct {
+	name       string
+	next, prev *nameListNode
+}
+
 // Namespace is a collection of instantiated modules which cannot conflict on name.
 type Namespace struct {
 	// moduleNamesList ensures modules are closed in reverse initialization order.
-	moduleNamesList []string // guarded by mux
+	moduleNamesList *nameListNode // guarded by mux
 
 	// moduleNamesSet ensures no race conditions instantiating two modules of the same name
-	moduleNamesSet map[string]struct{} // guarded by mux
+	moduleNamesSet map[string]*nameListNode // guarded by mux
 
 	// modules holds the instantiated Wasm modules by module name from Instantiate.
 	modules map[string]*ModuleInstance // guarded by mux
@@ -27,7 +33,7 @@ type Namespace struct {
 func newNamespace() *Namespace {
 	return &Namespace{
 		moduleNamesList: nil,
-		moduleNamesSet:  map[string]struct{}{},
+		moduleNamesSet:  map[string]*nameListNode{},
 		modules:         map[string]*ModuleInstance{},
 	}
 }
@@ -43,15 +49,23 @@ func (ns *Namespace) addModule(m *ModuleInstance) {
 func (ns *Namespace) deleteModule(moduleName string) {
 	ns.mux.Lock()
 	defer ns.mux.Unlock()
+	node, ok := ns.moduleNamesSet[moduleName]
+	if !ok {
+		return
+	}
+
+	// remove this module name
+	if node.prev == nil {
+		ns.moduleNamesList = node.next
+	} else {
+		node.prev.next = node.next
+	}
+	if node.next != nil {
+		node.next.prev = node.prev
+	}
+
 	delete(ns.modules, moduleName)
 	delete(ns.moduleNamesSet, moduleName)
-	// remove this module name
-	for i, n := range ns.moduleNamesList {
-		if n == moduleName {
-			ns.moduleNamesList = append(ns.moduleNamesList[:i], ns.moduleNamesList[i+1:]...)
-			break
-		}
-	}
 }
 
 // module returns the module of the given name or nil if not in this namespace
@@ -86,8 +100,17 @@ func (ns *Namespace) requireModuleName(moduleName string) error {
 	if _, ok := ns.moduleNamesSet[moduleName]; ok {
 		return fmt.Errorf("module[%s] has already been instantiated", moduleName)
 	}
-	ns.moduleNamesSet[moduleName] = struct{}{}
-	ns.moduleNamesList = append(ns.moduleNamesList, moduleName)
+
+	// add the newest node to the moduleNamesList as the head.
+	node := &nameListNode{
+		name: moduleName,
+		next: ns.moduleNamesList,
+	}
+	if node.next != nil {
+		node.next.prev = node
+	}
+	ns.moduleNamesList = node
+	ns.moduleNamesSet[moduleName] = node
 	return nil
 }
 
@@ -103,16 +126,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.moduleNamesList) - 1; i >= 0; i-- {
+	for node := ns.moduleNamesList; node != nil; node = node.next {
 		// If closing this module errs, proceed anyway to close the others.
-		if m, ok := ns.modules[ns.moduleNamesList[i]]; ok {
+		if m, ok := ns.modules[node.name]; ok {
 			if _, e := m.CallCtx.close(ctx, exitCode); e != nil && err == nil {
 				err = e // first error
 			}
 		}
 	}
 	ns.moduleNamesList = nil
-	ns.moduleNamesSet = map[string]struct{}{}
+	ns.moduleNamesSet = map[string]*nameListNode{}
 	ns.modules = map[string]*ModuleInstance{}
 	return
 }

internal/wasm/namespace_test.go

@@ -47,8 +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, map[string]struct{}{m1.Name: {}}, ns.moduleNamesSet)
-		require.Equal(t, []string{m1.Name}, ns.moduleNamesList)
+		require.Equal(t, map[string]*nameListNode{m1.Name: {name: m1.Name}}, ns.moduleNamesSet)
+		require.Equal(t, &nameListNode{name: m1.Name}, ns.moduleNamesList)
 	})
 
 	t.Run("ok if missing", func(t *testing.T) {
@@ -60,7 +60,7 @@ func TestNamespace_deleteModule(t *testing.T) {
 
 		require.Zero(t, len(ns.modules))
 		require.Zero(t, len(ns.moduleNamesSet))
-		require.Zero(t, len(ns.moduleNamesList))
+		require.Nil(t, ns.moduleNamesList)
 	})
 }
 
@@ -93,24 +93,27 @@ func TestNamespace_requireModules(t *testing.T) {
 }
 
 func TestNamespace_requireModuleName(t *testing.T) {
-	ns := &Namespace{moduleNamesSet: map[string]struct{}{}}
+	ns := &Namespace{moduleNamesSet: map[string]*nameListNode{}}
 
 	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.moduleNamesList)
-		require.Equal(t, map[string]struct{}{"m1": {}}, ns.moduleNamesSet)
+		require.Equal(t, &nameListNode{name: "m1"}, ns.moduleNamesList)
+		require.Equal(t, map[string]*nameListNode{"m1": {name: "m1"}}, ns.moduleNamesSet)
 		require.Zero(t, len(ns.modules))
 	})
 	t.Run("second", func(t *testing.T) {
 		err := ns.requireModuleName("m2")
 		require.NoError(t, err)
+		m2Node := &nameListNode{name: "m2"}
+		m1Node := &nameListNode{name: "m1", prev: m2Node}
+		m2Node.next = m1Node
 
 		// Appends in order.
-		require.Equal(t, []string{"m1", "m2"}, ns.moduleNamesList)
-		require.Equal(t, map[string]struct{}{"m1": {}, "m2": {}}, ns.moduleNamesSet)
+		require.Equal(t, m2Node, ns.moduleNamesList)
+		require.Equal(t, map[string]*nameListNode{"m1": m1Node, "m2": m2Node}, ns.moduleNamesSet)
 	})
 	t.Run("existing", func(t *testing.T) {
 		err := ns.requireModuleName("m2")
@@ -165,7 +168,7 @@ func TestNamespace_CloseWithExitCode(t *testing.T) {
 			// Namespace state zeroed
 			require.Zero(t, len(ns.modules))
 			require.Zero(t, len(ns.moduleNamesSet))
-			require.Zero(t, len(ns.moduleNamesList))
+			require.Nil(t, ns.moduleNamesList)
 		})
 	}
 
@@ -191,7 +194,7 @@ func TestNamespace_CloseWithExitCode(t *testing.T) {
 		// Namespace state zeroed
 		require.Zero(t, len(ns.modules))
 		require.Zero(t, len(ns.moduleNamesSet))
-		require.Zero(t, len(ns.moduleNamesList))
+		require.Nil(t, ns.moduleNamesList)
 	})
 }
 
@@ -217,7 +220,9 @@ func newTestNamespace() (*Namespace, *ModuleInstance, *ModuleInstance) {
 	m2.CallCtx = NewCallContext(ns, m2, nil)
 
 	ns.modules = map[string]*ModuleInstance{m1.Name: m1, m2.Name: m2}
-	ns.moduleNamesSet = map[string]struct{}{m1.Name: {}, m2.Name: {}}
-	ns.moduleNamesList = []string{m1.Name, m2.Name}
+	node1 := &nameListNode{name: m1.Name}
+	node2 := &nameListNode{name: m2.Name, next: node1}
+	ns.moduleNamesSet = map[string]*nameListNode{m1.Name: node1, m2.Name: node2}
+	ns.moduleNamesList = node2
 	return ns, m1, m2
 }