Repository: tetratelabs/wazero
Pull request #1285 (2023-03-24 16:30:49)

Shrink the store's nameToModule map
Putting this out there. Go does not shrink maps when elements are deleted so a long lived map with lots of elements added and deleted can consume a lot of memory.

This change would track the largest size the nameToNode map has grown to and if that size is more than twice the size it needs to be it gets replaced with a new map. Since we wouldn't want to be constantly resizing small maps I've set a floor and set the initial capacity of the map to that floor. I some what arbitrarily chose 10,000 as the floor.

The scenario this helps with is if a long running application has bursty periods of time when it instantiates a lot of named modules but most of the time has a very low volume of named modules.

I think it might be nice to make an option to disable or adjust the floor for applications running in memory constrained environments who know they will never need to instantiate a lot of named modules.

If there is interest in merging this I'll update the tests and add that option.

ShivanshVij: It might also make sense to add a knob disabling this - we want the most throughput possible on the system and a massive preallocated map works for that. 

Not the most efficient use of memory, sure, but you can always buy more ram. 
codefromthecrypt: FWIW I prefer a better, affordable default value vs a million config switches for something that is unlikely to be set by most people. Config is hugely distracting and I recall some places I worked at you'd have a bunch of weedy config and it was hard to feel the project was relevant as they didn't config important things and fuzzed over minute knobs.

my 2p
ckaznocha: Updated to shrink the `nameToModule` to adjust to the changes in #1359. Added tests and reduced the initial cap to 100. No settings added for now. Easer to add later than it is remove premature settings!
ncruces: Using `initialNameToModuleListSize` seems absolutely necessary to use as a minimum when reclaiming space (to avoid degenerate allocation patterns). But I wonder if it really is important to set the initial capacity.

Since anonymous modules many runtimes will have… less than a handful names in the map? The fixed cost of 100 is probably low, and if it's beneficial to other use cases, I guess it's fine. But I just wonder if it makes sense.
ckaznocha: @ncruces fair point. Removed the initial map size so shrinking shouldn't have much impact until the map grows organically.

@codefromthecrypt benchmarks
```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
                                       │    old.txt    │               new.txt                │
                                       │    sec/op     │    sec/op     vs base                │
Initialization/interpreter-10             11.31µ ± ∞ ¹   12.33µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/interpreter-multiple-10   12.064µ ± ∞ ¹   9.467µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-10                10.76µ ± ∞ ¹   10.18µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-multiple-10       10.48µ ± ∞ ¹   10.50µ ± ∞ ¹       ~ (p=1.000 n=1) ²
geomean                                   11.14µ         10.57µ        -5.07%
¹ need >= 6 samples for confidence interval at level 0.95
² need >= 4 samples to detect a difference at alpha level 0.05

                                       │    old.txt    │                new.txt                │
                                       │     B/op      │     B/op       vs base                │
Initialization/interpreter-10            132.3Ki ± ∞ ¹   132.3Ki ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/interpreter-multiple-10   132.3Ki ± ∞ ¹   132.3Ki ± ∞ ¹       ~ (p=1.000 n=1) ³
Initialization/compiler-10               132.3Ki ± ∞ ¹   132.3Ki ± ∞ ¹       ~ (p=1.000 n=1) ³
Initialization/compiler-multiple-10      132.3Ki ± ∞ ¹   132.3Ki ± ∞ ¹       ~ (p=1.000 n=1) ³
geomean                                  132.3Ki         132.3Ki        -0.00%
¹ need >= 6 samples for confidence interval at level 0.95
² all samples are equal
³ need >= 4 samples to detect a difference at alpha level 0.05

                                       │   old.txt   │               new.txt               │
                                       │  allocs/op  │  allocs/op   vs base                │
Initialization/interpreter-10            23.00 ± ∞ ¹   23.00 ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/interpreter-multiple-10   23.00 ± ∞ ¹   23.00 ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-10               23.00 ± ∞ ¹   23.00 ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-multiple-10      23.00 ± ∞ ¹   23.00 ± ∞ ¹       ~ (p=1.000 n=1) ²
geomean                                  23.00         23.00        +0.00%
¹ need >= 6 samples for confidence interval at level 0.95
² all samples are equal
```

Patch:
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 35549ff6..83bfe5af 100644
--- a/internal/wasm/store.go
+++ b/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
 }
diff --git a/internal/wasm/store_module_list.go b/internal/wasm/store_module_list.go
index 8c61e51a..02e2842c 100644
--- a/internal/wasm/store_module_list.go
+++ b/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
 }
 


Test patch:
diff --git a/internal/wasm/store_module_list_test.go b/internal/wasm/store_module_list_test.go
index ff5b763a..75bdbf71 100644
--- a/internal/wasm/store_module_list_test.go
+++ b/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)
+		}
 	})
 }