Shrink the store's nameToModule map (#1285)

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

internal/wasm/store.go

@@ -12,6 +12,11 @@ import (
 	"github.com/tetratelabs/wazero/sys"
 )
 
+// nameToModuleShrinkThreshold is the size the nameToModule map can grow to
+// before it starts to be monitored for shrinking.
+// The capacity will never be smaller than this once the threshold is met.
+const nameToModuleShrinkThreshold = 100
+
 type (
 	// Store is the runtime representation of "instantiated" Wasm module and objects.
 	// Multiple modules can be instantiated within a single store, and each instance,
@@ -32,6 +37,10 @@ type (
 		// It ensures no race conditions instantiating two modules of the same name.
 		nameToModule map[string]*ModuleInstance // guarded by mux
 
+		// nameToModuleCap tracks the growth of the nameToModule map in order to
+		// track when to shrink it.
+		nameToModuleCap int // guarded by mux
+
 		// EnabledFeatures are read-only to allow optimizations.
 		EnabledFeatures api.CoreFeatures
 
@@ -259,6 +268,7 @@ func NewStore(enabledFeatures api.CoreFeatures, engine Engine) *Store {
 	}
 	return &Store{
 		nameToModule:     map[string]*ModuleInstance{},
+		nameToModuleCap:  nameToModuleShrinkThreshold,
 		EnabledFeatures:  enabledFeatures,
 		Engine:           engine,
 		typeIDs:          typeIDs,
@@ -613,6 +623,7 @@ func (s *Store) CloseWithExitCode(ctx context.Context, exitCode uint32) (err err
 	}
 	s.moduleList = nil
 	s.nameToModule = nil
+	s.nameToModuleCap = 0
 	s.typeIDs = nil
 	return
 }

internal/wasm/store_module_list.go

@@ -29,6 +29,20 @@ func (s *Store) deleteModule(m *ModuleInstance) error {
 
 	if m.ModuleName != "" {
 		delete(s.nameToModule, m.ModuleName)
+
+		// Shrink the map if it's allocated more than twice the size of the list
+		newCap := len(s.nameToModule)
+		if newCap < nameToModuleShrinkThreshold {
+			newCap = nameToModuleShrinkThreshold
+		}
+		if newCap*2 <= s.nameToModuleCap {
+			nameToModule := make(map[string]*ModuleInstance, newCap)
+			for k, v := range s.nameToModule {
+				nameToModule[k] = v
+			}
+			s.nameToModule = nameToModule
+			s.nameToModuleCap = newCap
+		}
 	}
 	return nil
 }
@@ -59,6 +73,9 @@ func (s *Store) registerModule(m *ModuleInstance) error {
 			return fmt.Errorf("module[%s] has already been instantiated", m.ModuleName)
 		}
 		s.nameToModule[m.ModuleName] = m
+		if len(s.nameToModule) > s.nameToModuleCap {
+			s.nameToModuleCap = len(s.nameToModule)
+		}
 	}
 
 	// Add the newest node to the moduleNamesList as the head.
@@ -77,6 +94,9 @@ func (s *Store) AliasModule(src, dst string) error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
 	s.nameToModule[dst] = s.nameToModule[src]
+	if len(s.nameToModule) > s.nameToModuleCap {
+		s.nameToModuleCap = len(s.nameToModule)
+	}
 	return nil
 }
 

internal/wasm/store_module_list_test.go

@@ -2,6 +2,7 @@ package wasm
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/tetratelabs/wazero/internal/testing/require"
@@ -15,6 +16,7 @@ func TestStore_registerModule(t *testing.T) {
 		require.NoError(t, s.registerModule(m1))
 		require.Equal(t, map[string]*ModuleInstance{m1.ModuleName: m1}, s.nameToModule)
 		require.Equal(t, m1, s.moduleList)
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
 	})
 
 	t.Run("adds second module", func(t *testing.T) {
@@ -22,6 +24,7 @@ func TestStore_registerModule(t *testing.T) {
 		require.NoError(t, s.registerModule(m2))
 		require.Equal(t, map[string]*ModuleInstance{m1.ModuleName: m1, m2.ModuleName: m2}, s.nameToModule)
 		require.Equal(t, m2, s.moduleList)
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
 	})
 
 	t.Run("error on duplicated non anonymous", func(t *testing.T) {
@@ -44,6 +47,7 @@ func TestStore_deleteModule(t *testing.T) {
 		// Leaves the other module alone.
 		require.Equal(t, map[string]*ModuleInstance{m1.ModuleName: m1}, s.nameToModule)
 		require.Equal(t, m1, s.moduleList)
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
 	})
 
 	t.Run("ok if missing", func(t *testing.T) {
@@ -55,6 +59,7 @@ func TestStore_deleteModule(t *testing.T) {
 
 		require.Zero(t, len(s.nameToModule))
 		require.Nil(t, s.moduleList)
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
 	})
 
 	t.Run("delete middle", func(t *testing.T) {
@@ -112,6 +117,56 @@ func TestStore_AliasModule(t *testing.T) {
 		require.Equal(t, map[string]*ModuleInstance{"m1": m1, "m2": m1}, s.nameToModule)
 		// Doesn't affect module names
 		require.Nil(t, s.moduleList)
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
+	})
+}
+
+func TestStore_nameToModuleCap(t *testing.T) {
+	t.Run("nameToModuleCap grows beyond initial cap", func(t *testing.T) {
+		s := newStore()
+		for i := 0; i < 300; i++ {
+			require.NoError(t, s.registerModule(&ModuleInstance{ModuleName: fmt.Sprintf("m%d", i)}))
+		}
+
+		require.Equal(t, 300, s.nameToModuleCap)
+	})
+
+	t.Run("nameToModuleCap shrinks by half the cap", func(t *testing.T) {
+		s := newStore()
+		for i := 0; i < 400; i++ {
+			require.NoError(t, s.registerModule(&ModuleInstance{ModuleName: fmt.Sprintf("m%d", i)}))
+		}
+
+		for i := 0; i < 250; i++ {
+			require.NoError(t, s.deleteModule(s.nameToModule[fmt.Sprintf("m%d", i)]))
+		}
+
+		require.Equal(t, 200, s.nameToModuleCap)
+	})
+
+	t.Run("nameToModuleCap does not shrink below initial size", func(t *testing.T) {
+		s := newStore()
+		for i := 0; i < 400; i++ {
+			require.NoError(t, s.registerModule(&ModuleInstance{ModuleName: fmt.Sprintf("m%d", i)}))
+		}
+
+		for i := 0; i < 350; i++ {
+			require.NoError(t, s.deleteModule(s.nameToModule[fmt.Sprintf("m%d", i)]))
+		}
+
+		require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
+	})
+
+	t.Run("nameToModuleCap does not grow when if nameToModule does not grow", func(t *testing.T) {
+		s := newStore()
+		for i := 0; i < 99; i++ {
+			require.NoError(t, s.registerModule(&ModuleInstance{ModuleName: fmt.Sprintf("m%d", i)}))
+		}
+		for i := 0; i < 400; i++ {
+			require.NoError(t, s.registerModule(&ModuleInstance{ModuleName: fmt.Sprintf("m%d", i+99)}))
+			require.NoError(t, s.deleteModule(s.nameToModule[fmt.Sprintf("m%d", i+99)]))
+			require.Equal(t, nameToModuleShrinkThreshold, s.nameToModuleCap)
+		}
 	})
 }