Repository: tetratelabs/wazero
Pull request #1365 (2023-04-17 00:21:53)

Cleanup aliased modules when the main module is deleted
When a module is deleted which has been aliased, the aliases are not cleaned up meaning the original module will be closed and unusable but never get GC'd. I understand that aliases are currently only used in spect tests however `*store.AliasModule` is an exported method covered by tests so it feels right to fix the bug. This fix should have almost no meaningful impact to normal uses as `Module.aliases` will remain `nil` for all modules if no aliases are ever created.

```
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            10.94µ ± ∞ ¹   10.84µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/interpreter-multiple-10   9.698µ ± ∞ ¹   9.356µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-10               10.52µ ± ∞ ¹   10.42µ ± ∞ ¹       ~ (p=1.000 n=1) ²
Initialization/compiler-multiple-10      10.59µ ± ∞ ¹   11.07µ ± ∞ ¹       ~ (p=1.000 n=1) ²
geomean                                  10.42µ         10.40µ        -0.26%
¹ 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.02%
¹ 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               │
                                       │  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 83bfe5af..bc6bc627 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -119,6 +119,11 @@ type (
 		s *Store
 		// prev and next hold the nodes in the linked list of ModuleInstance held by Store.
 		prev, next *ModuleInstance
+		// aliases holds the module names that are aliases of this module registered in the store.
+		// Access to this field must be guarded by s.mux.
+		//
+		// Note: This is currently only used for spectests and will be nil in most cases.
+		aliases []string
 		// Definitions is derived from *Module, and is constructed during compilation phrase.
 		Definitions []FunctionDefinition
 	}
diff --git a/internal/wasm/store_module_list.go b/internal/wasm/store_module_list.go
index 02e2842c..bdfbbdec 100644
--- a/internal/wasm/store_module_list.go
+++ b/internal/wasm/store_module_list.go
@@ -30,6 +30,12 @@ func (s *Store) deleteModule(m *ModuleInstance) error {
 	if m.ModuleName != "" {
 		delete(s.nameToModule, m.ModuleName)
 
+		// Under normal circumstances, m.aliases will be nil this loop will not
+		// be entered unless aliases have been created. See `*store.AliasModule`
+		for _, alias := range m.aliases {
+			delete(s.nameToModule, alias)
+		}
+
 		// Shrink the map if it's allocated more than twice the size of the list
 		newCap := len(s.nameToModule)
 		if newCap < nameToModuleShrinkThreshold {
@@ -93,7 +99,15 @@ func (s *Store) registerModule(m *ModuleInstance) error {
 func (s *Store) AliasModule(src, dst string) error {
 	s.mux.Lock()
 	defer s.mux.Unlock()
-	s.nameToModule[dst] = s.nameToModule[src]
+	if _, ok := s.nameToModule[dst]; ok {
+		return nil
+	}
+	m, ok := s.nameToModule[src]
+	if !ok {
+		return nil
+	}
+	m.aliases = append(m.aliases, dst)
+	s.nameToModule[dst] = m
 	if len(s.nameToModule) > s.nameToModuleCap {
 		s.nameToModuleCap = len(s.nameToModule)
 	}


Test patch:
diff --git a/internal/wasm/store_module_list_test.go b/internal/wasm/store_module_list_test.go
index 75bdbf71..462c6880 100644
--- a/internal/wasm/store_module_list_test.go
+++ b/internal/wasm/store_module_list_test.go
@@ -108,17 +108,28 @@ func TestStore_module(t *testing.T) {
 }
 
 func TestStore_AliasModule(t *testing.T) {
-	s := newStore()
-	m1 := &ModuleInstance{ModuleName: "m1"}
-	s.nameToModule[m1.ModuleName] = m1
-
 	t.Run("alias module", func(t *testing.T) {
+		s := newStore()
+		m1 := &ModuleInstance{ModuleName: "m1"}
+		s.nameToModule[m1.ModuleName] = m1
+
 		require.NoError(t, s.AliasModule("m1", "m2"))
 		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)
 	})
+
+	t.Run("delete aliased module", func(t *testing.T) {
+		s := newStore()
+		m1 := &ModuleInstance{ModuleName: "m1"}
+		s.nameToModule[m1.ModuleName] = m1
+
+		require.NoError(t, s.AliasModule("m1", "m2"))
+		require.NoError(t, s.deleteModule(m1))
+		_, ok := s.nameToModule["m2"]
+		require.False(t, ok)
+	})
 }
 
 func TestStore_nameToModuleCap(t *testing.T) {