Repository: tetratelabs/wazero
Pull request #360 (2022-03-11 05:53:14)

Adds Module.Close API based on runtime.Finalizer
This commit adds Module.Close API, which can be used to indicate that
the module will no longer be used and setting finalizers for the allocated
resources. Notably, this will set finalizers on compiledFunction to deallocate
mapped memory regions as well as removing them from the store-wide
storage. As a result, users can re-instantiate a module for the same name
while having the safety -- Calling Module.Close is safe even when there are
outstanding function calls as we never explicitly deallocate but indirectly
do that via Goruntine's finalizers.


partially resolves: #293 

Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


mathetake: OK almost ready, but I would like to at least representative test case in conccurent_test.go.. other than that I think this is ready
mathetake: will do the next week!

Patch:
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index 667e5f24..a5dc584d 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -26,6 +26,6 @@ type ModuleEngine interface {
 	// The ctx's context.Context will be the outer-most ancestor of the argument to wasm.Function.
 	Call(ctx *ModuleContext, f *FunctionInstance, params ...uint64) (results []uint64, err error)
 
-	// Release releases the resources allocated by functions in this ModuleEngine.
-	Release() error
+	// Close releases the resources allocated by functions in this ModuleEngine.
+	Close()
 }
diff --git a/internal/wasm/interpreter/interpreter.go b/internal/wasm/interpreter/interpreter.go
index 5288840b..2ffcbcde 100644
--- a/internal/wasm/interpreter/interpreter.go
+++ b/internal/wasm/interpreter/interpreter.go
@@ -148,6 +148,14 @@ type interpreterOp struct {
 	rs     []*wazeroir.InclusiveRange
 }
 
+// Close implements wasm.Engine Close
+func (me *moduleEngine) Close() {
+	// Release all the function instances declared in this module.
+	for _, cf := range me.compiledFunctions[me.importedFunctionCounts:] {
+		me.parentEngine.deleteCompiledFunction(cf.funcInstance)
+	}
+}
+
 // NewModuleEngine implements internalwasm.Engine NewModuleEngine
 func (e *engine) NewModuleEngine(importedFunctions, moduleFunctions []*wasm.FunctionInstance) (wasm.ModuleEngine, error) {
 	me := &moduleEngine{
@@ -169,14 +177,14 @@ func (e *engine) NewModuleEngine(importedFunctions, moduleFunctions []*wasm.Func
 		if f.Kind == wasm.FunctionKindWasm {
 			ir, err := wazeroir.Compile(f)
 			if err != nil {
-				_ = me.Release()
+				me.Close()
 				// TODO(Adrian): extract Module.funcDesc so that errors here have more context
 				return nil, fmt.Errorf("function[%d/%d] failed to lower to wazeroir: %w", i, len(moduleFunctions)-1, err)
 			}
 
 			compiled, err = e.lowerIROps(f, ir.Operations)
 			if err != nil {
-				_ = me.Release()
+				me.Close()
 				return nil, fmt.Errorf("function[%d/%d] failed to convert wazeroir operations: %w", i, len(moduleFunctions)-1, err)
 			}
 		} else {
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index 69b51a6f..a2fa58dc 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -8,6 +8,7 @@ import (
 	"runtime/debug"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"unsafe"
 
 	wasm "github.com/tetratelabs/wazero/internal/wasm"
@@ -33,6 +34,11 @@ type (
 		// parentEngine holds *engine from which this module engine is created from.
 		parentEngine           *engine
 		importedFunctionCounts int
+
+		// Closed indicates whether Close was called for this moduleEngine, and used to prevent
+		// double-free of compiledFunctions.
+		// Note: intentionally use uint32 instead of bool for atomic operation.
+		closed uint32
 	}
 
 	// callEngine holds context per moduleEngine.Call, and shared across all the
@@ -348,7 +354,7 @@ func (e *engine) NewModuleEngine(importedFunctions, moduleFunctions []*wasm.Func
 			compiled, err = compileHostFunction(f)
 		}
 		if err != nil {
-			_ = me.Release()
+			me.Close()
 			return nil, fmt.Errorf("function[%d/%d] %w", i, len(moduleFunctions)-1, err)
 		}
 		me.compiledFunctions = append(me.compiledFunctions, compiled)
@@ -366,16 +372,20 @@ func (me *moduleEngine) FunctionAddress(index wasm.Index) uintptr {
 	return uintptr(unsafe.Pointer(me.compiledFunctions[index]))
 }
 
-// Release implements internalwasm.ModuleEngine Release
-func (me *moduleEngine) Release() error {
-	// Release all the function instances declared in this module.
-	for _, cf := range me.compiledFunctions[me.importedFunctionCounts:] {
-		if err := munmapCodeSegment(cf.codeSegment); err != nil {
-			return err
+// Close implements internalwasm.ModuleEngine Close
+func (me *moduleEngine) Close() {
+	// Setting finalizer multiple times result in panic, so we guard with CAS.
+	if atomic.CompareAndSwapUint32(&me.closed, 0, 1) {
+		// Release all the function instances declared in this module.
+		for _, cf := range me.compiledFunctions[me.importedFunctionCounts:] {
+			runtime.SetFinalizer(cf, func(compiledFn *compiledFunction) {
+				if err := munmapCodeSegment(compiledFn.codeSegment); err != nil {
+					panic(err) // mumap failure cannot recover.
+				}
+			})
+			me.parentEngine.deleteCompiledFunction(cf.source)
 		}
-		me.parentEngine.deleteCompiledFunction(cf.source)
 	}
-	return nil
 }
 
 func (e *engine) deleteCompiledFunction(f *wasm.FunctionInstance) {
diff --git a/internal/wasm/module_context.go b/internal/wasm/module_context.go
index 9adec56d..3dc36c46 100644
--- a/internal/wasm/module_context.go
+++ b/internal/wasm/module_context.go
@@ -10,11 +10,12 @@ import (
 // compile time check to ensure ModuleContext implements wasm.Module
 var _ publicwasm.Module = &ModuleContext{}
 
-func NewModuleContext(ctx context.Context, engine Engine, instance *ModuleInstance) *ModuleContext {
+func NewModuleContext(ctx context.Context, store *Store, instance *ModuleInstance) *ModuleContext {
 	return &ModuleContext{
 		ctx:    ctx,
 		memory: instance.Memory,
 		module: instance,
+		store:  store,
 	}
 }
 
@@ -25,6 +26,7 @@ type ModuleContext struct {
 	module *ModuleInstance
 	// memory is returned by Memory and overridden WithMemory
 	memory publicwasm.Memory
+	store  *Store
 }
 
 // WithMemory allows overriding memory without re-allocation when the result would be the same.
@@ -55,6 +57,10 @@ func (m *ModuleContext) WithContext(ctx context.Context) publicwasm.Module {
 	return m
 }
 
+func (m *ModuleContext) Close() {
+	m.store.CloseModule(m.module.Name)
+}
+
 // Memory implements wasm.Module Memory
 func (m *ModuleContext) Memory() publicwasm.Memory {
 	return m.module.Memory
@@ -90,10 +96,11 @@ func (f *FunctionInstance) ResultTypes() []publicwasm.ValueType {
 
 // Call implements wasm.Function Call
 func (f *FunctionInstance) Call(ctx publicwasm.Module, params ...uint64) ([]uint64, error) {
+	mod := f.Module
 	if modCtx, ok := ctx.(*ModuleContext); !ok { // allow nil to substitute for the defining module
-		return f.Module.Engine.Call(f.Module.Ctx, f, params...)
+		return mod.Engine.Call(mod.Ctx, f, params...)
 	} else { // TODO: check if the importing context is correct
-		return f.Module.Engine.Call(modCtx, f, params...)
+		return mod.Engine.Call(modCtx, f, params...)
 	}
 }
 
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index 8b7c681b..e023950b 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -74,17 +74,6 @@ type (
 
 		// Engine implements function calls for this module.
 		Engine ModuleEngine
-
-		// mux is used to guard the fields from concurrent access.
-		mux sync.Mutex
-
-		// dependentCount is the current number of modules which import this module. On Store.ReleaseModule, this number
-		// must be zero otherwise it fails.
-		dependentCount int // guarded by mux
-
-		// dependencies holds imported modules. This is used when releasing this module instance, or decrementing the
-		// dependentCount of the imported modules.
-		dependencies map[*ModuleInstance]struct{}
 	}
 
 	// ExportInstance represents an exported instance in a Store.
@@ -170,10 +159,9 @@ const (
 // addSections adds section elements to the ModuleInstance
 func (m *ModuleInstance) addSections(module *Module, importedFunctions, functions []*FunctionInstance,
 	importedGlobals, globals []*GlobalInstance, importedTable, table *TableInstance,
-	memory, importedMemory *MemoryInstance, typeInstances []*TypeInstance, moduleImports map[*ModuleInstance]struct{}) {
+	memory, importedMemory *MemoryInstance, typeInstances []*TypeInstance) {
 
 	m.Types = typeInstances
-	m.dependencies = moduleImports
 
 	m.Functions = append(m.Functions, importedFunctions...)
 	for i, f := range functions {
@@ -290,15 +278,7 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 		return nil, err
 	}
 
-	moduleImportsFinalized := false
-	importedFunctions, importedGlobals, importedTable, importedMemory, moduleImports, err := s.resolveImports(module)
-	defer func() {
-		if !moduleImportsFinalized {
-			for moduleImport := range moduleImports {
-				moduleImport.decDependentCount()
-			}
-		}
-	}()
+	importedFunctions, importedGlobals, importedTable, importedMemory, err := s.resolveImports(module)
 	if err != nil {
 		s.deleteModule(name)
 		return nil, err
@@ -320,11 +300,12 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 	// Now we have all instances from imports and local ones, so ready to create a new ModuleInstance.
 	m := &ModuleInstance{Name: name}
 	m.addSections(module, importedFunctions, functions, importedGlobals,
-		globals, importedTable, table, importedMemory, memory, types, moduleImports)
+		globals, importedTable, table, importedMemory, memory, types)
 
 	// Plus we are ready to compile functions.
 	m.Engine, err = s.engine.NewModuleEngine(importedFunctions, functions)
 	if err != nil {
+		s.deleteModule(name)
 		return nil, fmt.Errorf("compilation failed: %w", err)
 	}
 
@@ -343,10 +324,7 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 	m.applyData(module.DataSection)
 
 	// Build the default context for calls to this module.
-	m.Ctx = NewModuleContext(s.ctx, s.engine, m)
-
-	// Plus, we can finalize the module import reference count.
-	moduleImportsFinalized = true
+	m.Ctx = NewModuleContext(s.ctx, s, m)
 
 	// Execute the start function.
 	if module.StartSection != nil {
@@ -363,45 +341,13 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 	return m.Ctx, nil
 }
 
-// ReleaseModule deallocates resources if a module with the given name exists.
-func (s *Store) ReleaseModule(moduleName string) error {
+// CloseModule deallocates resources if a module with the given name exists.
+func (s *Store) CloseModule(moduleName string) {
 	m := s.module(moduleName)
-	if m == nil {
-		return nil // already released
-	}
-
-	m.mux.Lock()
-	defer m.mux.Unlock()
-
-	if m.dependentCount > 0 {
-		// This case other modules are importing this module instance and still alive.
-		return fmt.Errorf("%d modules import this and need to be closed first", m.dependentCount)
-	}
-
-	// TODO: check outstanding calls and wait until they exit.
-
-	for mod := range m.dependencies {
-		mod.decDependentCount()
-	}
-
-	if err := m.Engine.Release(); err != nil {
-		return fmt.Errorf("unable to release function instance: %w", err)
+	if m != nil {
+		m.Engine.Close()
+		s.deleteModule(moduleName)
 	}
-
-	s.deleteModule(moduleName)
-	return nil
-}
-
-func (m *ModuleInstance) decDependentCount() {
-	m.mux.Lock()
-	defer m.mux.Unlock()
-	m.dependentCount--
-}
-
-func (m *ModuleInstance) incDependentCount() {
-	m.mux.Lock()
-	defer m.mux.Unlock()
-	m.dependentCount++
 }
 
 // deleteModule makes the moduleName available for instantiation again.
@@ -449,13 +395,11 @@ func (s *Store) module(moduleName string) *ModuleInstance {
 func (s *Store) resolveImports(module *Module) (
 	importedFunctions []*FunctionInstance, importedGlobals []*GlobalInstance,
 	importedTable *TableInstance, importedMemory *MemoryInstance,
-	moduleImports map[*ModuleInstance]struct{},
 	err error,
 ) {
 	s.mux.RLock()
 	defer s.mux.RUnlock()
 
-	moduleImports = map[*ModuleInstance]struct{}{}
 	for idx, i := range module.ImportSection {
 		m, ok := s.modules[i.Module]
 		if !ok {
@@ -463,11 +407,6 @@ func (s *Store) resolveImports(module *Module) (
 			return
 		}
 
-		if _, ok := moduleImports[m]; !ok {
-			m.incDependentCount() // TODO: check if the module is already released. See #293
-			moduleImports[m] = struct{}{}
-		}
-
 		var imported *ExportInstance
 		imported, err = m.getExport(i.Name, i.Type)
 		if err != nil {
diff --git a/wasm/wasm.go b/wasm/wasm.go
index 2763feb0..5a8e196b 100644
--- a/wasm/wasm.go
+++ b/wasm/wasm.go
@@ -87,6 +87,10 @@ type Module interface {
 
 	// ExportedGlobal a global exported from this module or nil if it wasn't.
 	ExportedGlobal(name string) Global
+
+	// Close releases resources allocated for this Module. Using this while having outstanding function calls is
+	// safe. After calling this function, re-instantiating a module for the same name is allowed.
+	Close()
 }
 
 // Function is a WebAssembly 1.0 (20191205) function exported from an instantiated module (wazero.Runtime InstantiateModule).


Test patch:
diff --git a/internal/wasm/interpreter/interpreter_test.go b/internal/wasm/interpreter/interpreter_test.go
index 320f6469..ee41163c 100644
--- a/internal/wasm/interpreter/interpreter_test.go
+++ b/internal/wasm/interpreter/interpreter_test.go
@@ -91,7 +91,7 @@ func TestEngine_Call_HostFn(t *testing.T) {
 
 	e := NewEngine()
 	module := &wasm.ModuleInstance{Memory: memory}
-	modCtx := wasm.NewModuleContext(context.Background(), e, module)
+	modCtx := wasm.NewModuleContext(context.Background(), wasm.NewStore(context.Background(), e, wasm.Features20191205), module)
 	f := &wasm.FunctionInstance{
 		GoFunc: &hostFn,
 		Kind:   wasm.FunctionKindGoModule,
@@ -281,7 +281,7 @@ func TestEngineCompile_Errors(t *testing.T) {
 	})
 }
 
-func TestRelease(t *testing.T) {
+func TestClose(t *testing.T) {
 	newFunctionInstance := func(id int) *wasm.FunctionInstance {
 		return &wasm.FunctionInstance{
 			Name: strconv.Itoa(id), Type: &wasm.FunctionType{}, Body: []byte{wasm.OpcodeEnd}, Module: &wasm.ModuleInstance{}}
@@ -327,8 +327,7 @@ func TestRelease(t *testing.T) {
 				require.Contains(t, e.compiledFunctions, f)
 			}
 
-			err = me.Release()
-			require.NoError(t, err)
+			me.Close()
 
 			require.Len(t, e.compiledFunctions, len(tc.importedFunctions))
 			for _, f := range tc.importedFunctions {
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index c63dd278..a9c96368 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -6,6 +6,7 @@ import (
 	"reflect"
 	"runtime"
 	"strconv"
+	"sync"
 	"testing"
 	"unsafe"
 
@@ -146,7 +147,7 @@ func TestEngine_Call_HostFn(t *testing.T) {
 
 	e := NewEngine()
 	module := &wasm.ModuleInstance{Memory: memory}
-	modCtx := wasm.NewModuleContext(context.Background(), e, module)
+	modCtx := wasm.NewModuleContext(context.Background(), wasm.NewStore(context.Background(), e, wasm.Features20191205), module)
 	f := &wasm.FunctionInstance{
 		GoFunc: &hostFn,
 		Kind:   wasm.FunctionKindGoModule,
@@ -252,10 +253,12 @@ func TestRelease(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			e := newEngine()
+			var importedModuleEngine *moduleEngine
 			if len(tc.importedFunctions) > 0 {
 				modEngine, err := e.NewModuleEngine(nil, tc.importedFunctions)
 				require.NoError(t, err)
-				require.Len(t, modEngine.(*moduleEngine).compiledFunctions, len(tc.importedFunctions))
+				importedModuleEngine = modEngine.(*moduleEngine)
+				require.Len(t, importedModuleEngine.compiledFunctions, len(tc.importedFunctions))
 			}
 
 			modEngine, err := e.NewModuleEngine(tc.importedFunctions, tc.moduleFunctions)
@@ -263,16 +266,30 @@ func TestRelease(t *testing.T) {
 			require.Len(t, modEngine.(*moduleEngine).compiledFunctions, len(tc.importedFunctions)+len(tc.moduleFunctions))
 
 			require.Len(t, e.compiledFunctions, len(tc.importedFunctions)+len(tc.moduleFunctions))
+
+			var importedMappedRegions [][]byte
 			for _, f := range tc.importedFunctions {
 				require.Contains(t, e.compiledFunctions, f)
+				importedMappedRegions = append(importedMappedRegions, e.compiledFunctions[f].codeSegment)
 			}
+			var mappedRegions [][]byte
 			for _, f := range tc.moduleFunctions {
 				require.Contains(t, e.compiledFunctions, f)
+				mappedRegions = append(mappedRegions, e.compiledFunctions[f].codeSegment)
 			}
 
-			err = modEngine.Release()
-			require.NoError(t, err)
+			const goroutines = 100
+			var wg sync.WaitGroup
+			wg.Add(goroutines)
+			for i := 0; i < goroutines; i++ {
+				go func() {
+					defer wg.Done()
+					modEngine.Close() // concurrent mutliple execution of Close must be guarded by atomic.
+				}()
+			}
+			wg.Wait()
 
+			require.True(t, modEngine.(*moduleEngine).closed == 1)
 			require.Len(t, e.compiledFunctions, len(tc.importedFunctions))
 			for _, f := range tc.importedFunctions {
 				require.Contains(t, e.compiledFunctions, f)
@@ -280,6 +297,29 @@ func TestRelease(t *testing.T) {
 			for i, f := range tc.moduleFunctions {
 				require.NotContains(t, e.compiledFunctions, f, i)
 			}
+
+			// Release the reference to ModuleEngine -> *compiledFunctions for non-imported compiledFunctions
+			// must be munmapped.
+			modEngine = nil
+			runtime.GC()
+
+			for _, mappedRegion := range mappedRegions {
+				// munmap twice should result in error.
+				err = munmapCodeSegment(mappedRegion)
+				require.Error(t, err)
+			}
+
+			if len(tc.importedFunctions) > 0 {
+				importedModuleEngine.Close()
+				importedModuleEngine = nil
+				runtime.GC()
+
+				for _, mappedRegion := range importedMappedRegions {
+					// munmap twice should result in error.
+					err = munmapCodeSegment(mappedRegion)
+					require.Error(t, err)
+				}
+			}
 		})
 	}
 }
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 48281ff6..b1395074 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -97,7 +97,7 @@ func TestPublicModule_String(t *testing.T) {
 	require.Equal(t, "Module[module]", s.Module(m.module.Name).String())
 }
 
-func TestStore_ReleaseModule(t *testing.T) {
+func TestStore_CloseModule(t *testing.T) {
 	const importedModuleName = "imported"
 	const importingModuleName = "test"
 
@@ -141,19 +141,21 @@ func TestStore_ReleaseModule(t *testing.T) {
 			}, importingModuleName)
 			require.NoError(t, err)
 
-			// We shouldn't be able to release the imported module as it is in use!
-			require.Error(t, s.ReleaseModule(importedModuleName))
+			_, ok := s.modules[importedModuleName]
+			require.True(t, ok)
 
-			// Can release the importing module
-			require.NoError(t, s.ReleaseModule(importingModuleName))
-			require.Nil(t, s.modules[importingModuleName])
+			_, ok = s.modules[importingModuleName]
+			require.True(t, ok)
+
+			// Release the importing module
+			s.CloseModule(importingModuleName)
 			require.NotContains(t, s.modules, importingModuleName)
 
 			// Can re-release the importing module
-			require.NoError(t, s.ReleaseModule(importingModuleName))
+			s.CloseModule(importingModuleName)
 
-			// Now we should be able to release the imported module.
-			require.NoError(t, s.ReleaseModule(importedModuleName))
+			// Now we release the imported module.
+			s.CloseModule(importedModuleName)
 			require.Nil(t, s.modules[importedModuleName])
 			require.NotContains(t, s.modules, importedModuleName)
 		})
@@ -195,36 +197,22 @@ func TestStore_concurrent(t *testing.T) {
 			defer wg.Done()
 			_, err := s.Instantiate(importingModule, strconv.Itoa(i))
 			require.NoError(t, err)
-
-			if i == goroutines/2 {
-				// Trying to release the imported module concurrently, but should fail as it's in use.
-				err := s.ReleaseModule(importedModuleName)
-				require.Error(t, err)
-			}
 		}(i)
-
-		// TODO: in addition to the normal instantiation, we should try to instantiate host module in conjunction.
-		// after making store.NewHostModule goroutine-safe.
 	}
 	wg.Wait()
 
-	// At this point 1000 modules import host modules.
-	require.Equal(t, goroutines, hm.dependentCount)
-
 	// Concurrent release.
 	wg.Add(goroutines)
 	for i := 0; i < goroutines; i++ {
 		go func(i int) {
 			defer wg.Done()
-			err := s.ReleaseModule(strconv.Itoa(i))
+			s.CloseModule(strconv.Itoa(i))
 			require.NoError(t, err)
 		}(i)
 	}
 	wg.Wait()
 
-	// Now all the importing instances were released, the imported module can be freed.
-	require.Zero(t, hm.dependentCount)
-	require.NoError(t, s.ReleaseModule(hm.Name))
+	s.CloseModule(hm.Name)
 
 	// All instances are freed.
 	require.Len(t, s.modules, 0)
@@ -265,9 +253,6 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 			},
 		}, importingModuleName)
 		require.EqualError(t, err, "module[non-exist] not instantiated")
-
-		// hm.dependentCount must be intact as the instantiation failed.
-		require.Zero(t, hm.dependentCount)
 	})
 
 	t.Run("compilation failed", func(t *testing.T) {
@@ -294,9 +279,6 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 			},
 		}, importingModuleName)
 		require.EqualError(t, err, "compilation failed: some compilation error")
-
-		// hm.dependentCount must be intact as the instantiation failed.
-		require.Zero(t, hm.dependentCount)
 	})
 
 	t.Run("start func failed", func(t *testing.T) {
@@ -321,9 +303,6 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 			},
 		}, importingModuleName)
 		require.EqualError(t, err, "start function[1] failed: call failed")
-
-		// hm.dependentCount must stay incremented as the instantiation itself has already succeeded.
-		require.Equal(t, 1, hm.dependentCount)
 	})
 }
 
@@ -458,9 +437,7 @@ func (e *mockModuleEngine) Call(ctx *ModuleContext, f *FunctionInstance, _ ...ui
 	return
 }
 
-func (e *mockModuleEngine) Release() error {
-	return nil
-}
+func (e *mockModuleEngine) Close() {}
 
 func TestStore_getTypeInstance(t *testing.T) {
 	t.Run("too many functions", func(t *testing.T) {
@@ -588,13 +565,13 @@ func TestStore_resolveImports(t *testing.T) {
 
 	t.Run("module not instantiated", func(t *testing.T) {
 		s := newStore()
-		_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: "unknown", Name: "unknown"}}})
+		_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: "unknown", Name: "unknown"}}})
 		require.EqualError(t, err, "module[unknown] not instantiated")
 	})
 	t.Run("export instance not found", func(t *testing.T) {
 		s := newStore()
 		s.modules[moduleName] = &ModuleInstance{Exports: map[string]*ExportInstance{}, Name: moduleName}
-		_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: "unknown"}}})
+		_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: "unknown"}}})
 		require.EqualError(t, err, "\"unknown\" is not exported in module \"test\"")
 	})
 	t.Run("func", func(t *testing.T) {
@@ -616,18 +593,15 @@ func TestStore_resolveImports(t *testing.T) {
 					{Module: moduleName, Name: "", Type: ExternTypeFunc, DescFunc: 1},
 				},
 			}
-			functions, _, _, _, moduleImports, err := s.resolveImports(m)
+			functions, _, _, _, err := s.resolveImports(m)
 			require.NoError(t, err)
-			require.Contains(t, moduleImports, s.modules[moduleName])
 			require.Contains(t, functions, f)
 			require.Contains(t, functions, g)
-			// DependentCount must be one even when importing two instances.
-			require.Equal(t, 1, s.modules[moduleName].dependentCount)
 		})
 		t.Run("type out of range", func(t *testing.T) {
 			s := newStore()
 			s.modules[moduleName] = &ModuleInstance{Exports: map[string]*ExportInstance{name: {}}, Name: moduleName}
-			_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeFunc, DescFunc: 100}}})
+			_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeFunc, DescFunc: 100}}})
 			require.EqualError(t, err, "import[0] func[test.target]: function type out of range")
 		})
 		t.Run("signature mismatch", func(t *testing.T) {
@@ -639,7 +613,7 @@ func TestStore_resolveImports(t *testing.T) {
 				TypeSection:   []*FunctionType{{Results: []ValueType{ValueTypeF32}}},
 				ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeFunc, DescFunc: 0}},
 			}
-			_, _, _, _, _, err := s.resolveImports(m)
+			_, _, _, _, err := s.resolveImports(m)
 			require.EqualError(t, err, "import[0] func[test.target]: signature mismatch: v_f32 != v_v")
 		})
 	})
@@ -648,10 +622,9 @@ func TestStore_resolveImports(t *testing.T) {
 			s := newStore()
 			inst := &GlobalInstance{Type: &GlobalType{ValType: ValueTypeI32}}
 			s.modules[moduleName] = &ModuleInstance{Exports: map[string]*ExportInstance{name: {Type: ExternTypeGlobal, Global: inst}}, Name: moduleName}
-			_, globals, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: inst.Type}}})
+			_, globals, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: inst.Type}}})
 			require.NoError(t, err)
 			require.Contains(t, globals, inst)
-			require.Equal(t, 1, s.modules[moduleName].dependentCount)
 		})
 		t.Run("mutability mismatch", func(t *testing.T) {
 			s := newStore()
@@ -659,7 +632,7 @@ func TestStore_resolveImports(t *testing.T) {
 				Type:   ExternTypeGlobal,
 				Global: &GlobalInstance{Type: &GlobalType{Mutable: false}},
 			}}, Name: moduleName}
-			_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: &GlobalType{Mutable: true}}}})
+			_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: &GlobalType{Mutable: true}}}})
 			require.EqualError(t, err, "import[0] global[test.target]: mutability mismatch: true != false")
 		})
 		t.Run("type mismatch", func(t *testing.T) {
@@ -668,7 +641,7 @@ func TestStore_resolveImports(t *testing.T) {
 				Type:   ExternTypeGlobal,
 				Global: &GlobalInstance{Type: &GlobalType{ValType: ValueTypeI32}},
 			}}, Name: moduleName}
-			_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: &GlobalType{ValType: ValueTypeF64}}}})
+			_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeGlobal, DescGlobal: &GlobalType{ValType: ValueTypeF64}}}})
 			require.EqualError(t, err, "import[0] global[test.target]: value type mismatch: f64 != i32")
 		})
 	})
@@ -681,10 +654,9 @@ func TestStore_resolveImports(t *testing.T) {
 				Type:   ExternTypeMemory,
 				Memory: memoryInst,
 			}}, Name: moduleName}
-			_, _, _, memory, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: &Memory{Max: &max}}}})
+			_, _, _, memory, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: &Memory{Max: &max}}}})
 			require.NoError(t, err)
 			require.Equal(t, memory, memoryInst)
-			require.Equal(t, 1, s.modules[moduleName].dependentCount)
 		})
 		t.Run("minimum size mismatch", func(t *testing.T) {
 			s := newStore()
@@ -693,7 +665,7 @@ func TestStore_resolveImports(t *testing.T) {
 				Type:   ExternTypeMemory,
 				Memory: &MemoryInstance{Min: importMemoryType.Min - 1},
 			}}, Name: moduleName}
-			_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: importMemoryType}}})
+			_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: importMemoryType}}})
 			require.EqualError(t, err, "import[0] memory[test.target]: minimum size mismatch: 2 > 1")
 		})
 		t.Run("maximum size mismatch", func(t *testing.T) {
@@ -704,7 +676,7 @@ func TestStore_resolveImports(t *testing.T) {
 				Type:   ExternTypeMemory,
 				Memory: &MemoryInstance{},
 			}}, Name: moduleName}
-			_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: importMemoryType}}})
+			_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeMemory, DescMem: importMemoryType}}})
 			require.EqualError(t, err, "import[0] memory[test.target]: maximum size mismatch: 10, but actual has no max")
 		})
 	})
@@ -759,34 +731,3 @@ func TestModuleInstance_applyData(t *testing.T) {
 	})
 	require.Equal(t, []byte{0xa, 0xf, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x5}, m.Memory.Buffer)
 }
-
-func TestModuleInstance_decDependentCount(t *testing.T) {
-	count := 100
-	m := ModuleInstance{dependentCount: count}
-
-	wg := sync.WaitGroup{}
-	wg.Add(count)
-	for i := 0; i < count; i++ {
-		go func() {
-			defer wg.Done()
-			m.decDependentCount()
-		}()
-	}
-	wg.Wait()
-	require.Zero(t, m.dependentCount)
-}
-
-func TestModuleInstance_incDependentCount(t *testing.T) {
-	count := 100
-	m := ModuleInstance{}
-	wg := sync.WaitGroup{}
-	wg.Add(count)
-	for i := 0; i < count; i++ {
-		go func() {
-			defer wg.Done()
-			m.incDependentCount()
-		}()
-	}
-	wg.Wait()
-	require.Equal(t, count, m.dependentCount)
-}
diff --git a/internal/wasm/table_test.go b/internal/wasm/table_test.go
index 1d68f31b..aae5b0a0 100644
--- a/internal/wasm/table_test.go
+++ b/internal/wasm/table_test.go
@@ -18,10 +18,9 @@ func TestStore_resolveImports_table(t *testing.T) {
 			Type:  ExternTypeTable,
 			Table: tableInst,
 		}}, Name: moduleName}
-		_, _, table, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: &Table{Max: &max}}}})
+		_, _, table, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: &Table{Max: &max}}}})
 		require.NoError(t, err)
 		require.Equal(t, table, tableInst)
-		require.Equal(t, 1, s.modules[moduleName].dependentCount)
 	})
 	t.Run("minimum size mismatch", func(t *testing.T) {
 		s := newStore()
@@ -30,7 +29,7 @@ func TestStore_resolveImports_table(t *testing.T) {
 			Type:  ExternTypeTable,
 			Table: &TableInstance{Min: importTableType.Min - 1},
 		}}, Name: moduleName}
-		_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: importTableType}}})
+		_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: importTableType}}})
 		require.EqualError(t, err, "import[0] table[test.target]: minimum size mismatch: 2 > 1")
 	})
 	t.Run("maximum size mismatch", func(t *testing.T) {
@@ -41,7 +40,7 @@ func TestStore_resolveImports_table(t *testing.T) {
 			Type:  ExternTypeTable,
 			Table: &TableInstance{Min: importTableType.Min - 1},
 		}}, Name: moduleName}
-		_, _, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: importTableType}}})
+		_, _, _, _, err := s.resolveImports(&Module{ImportSection: []*Import{{Module: moduleName, Name: name, Type: ExternTypeTable, DescTable: importTableType}}})
 		require.EqualError(t, err, "import[0] table[test.target]: maximum size mismatch: 10, but actual has no max")
 	})
 }
diff --git a/tests/engine/adhoc_test.go b/tests/engine/adhoc_test.go
index 19226f27..b87e4121 100644
--- a/tests/engine/adhoc_test.go
+++ b/tests/engine/adhoc_test.go
@@ -50,6 +50,9 @@ func runAdhocTests(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConfig)
 	t.Run("host function with float type", func(t *testing.T) {
 		testHostFunctions(t, newRuntimeConfig)
 	})
+	t.Run("close with outstanding calls", func(t *testing.T) {
+		testAdhocCloseWhileExecution(t, newRuntimeConfig)
+	})
 }
 
 func testHugeStack(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConfig) {
@@ -228,3 +231,62 @@ func testHostFunctions(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConf
 		})
 	}
 }
+
+// testAdhocCloseWhileExecution ensures that calling Module.Close with outstanding calls is safe.
+func testAdhocCloseWhileExecution(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConfig) {
+	t.Run("singleton", func(t *testing.T) {
+		r := wazero.NewRuntimeWithConfig(newRuntimeConfig())
+		var moduleCloser func()
+		_, err := r.NewModuleBuilder("host").ExportFunctions(map[string]interface{}{
+			"close_module": func() { moduleCloser() }, // Closing while executing itself.
+		}).Instantiate()
+		require.NoError(t, err)
+
+		m, err := r.InstantiateModuleFromSource([]byte(`(module $test
+	(import "host" "close_module" (func $close_module ))
+
+	(func $close_while_execution
+		call $close_module
+	)
+	(export "close_while_execution" (func $close_while_execution))
+)`))
+		require.NoError(t, err)
+
+		moduleCloser = m.Close
+
+		_, err = m.ExportedFunction("close_while_execution").Call(nil)
+		require.NoError(t, err)
+	})
+	t.Run("close imported module", func(t *testing.T) {
+		r := wazero.NewRuntimeWithConfig(newRuntimeConfig())
+		importedModule, err := r.NewModuleBuilder("host").ExportFunctions(map[string]interface{}{
+			"already_closed": func() {},
+		}).Instantiate()
+		require.NoError(t, err)
+
+		m, err := r.InstantiateModuleFromSource([]byte(`(module $test
+		(import "host" "already_closed" (func $already_closed ))
+
+		(func $close_parent_before_execution
+			call $already_closed
+		)
+		(export "close_parent_before_execution" (func $close_parent_before_execution))
+	)`))
+		require.NoError(t, err)
+
+		// Closing the imported module before making call should also safe.
+		importedModule.Close()
+
+		// Even we can re-enstantiate the module for the same name.
+		importedModuleNew, err := r.NewModuleBuilder("host").ExportFunctions(map[string]interface{}{
+			"already_closed": func() {
+				panic("unreachable") // The new module's function must not be called.
+			},
+		}).Instantiate()
+		require.NoError(t, err)
+		defer importedModuleNew.Close()
+
+		_, err = m.ExportedFunction("close_parent_before_execution").Call(nil)
+		require.NoError(t, err)
+	})
+}
diff --git a/tests/engine/concurrency_test.go b/tests/engine/concurrency_test.go
index 35f95828..f0b58d9c 100644
--- a/tests/engine/concurrency_test.go
+++ b/tests/engine/concurrency_test.go
@@ -3,6 +3,9 @@ package adhoc
 import (
 	"sync"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
 
 	"github.com/tetratelabs/wazero"
 )
@@ -12,13 +15,12 @@ func TestJITConcurrency(t *testing.T) {
 		t.Skip()
 	}
 	runAdhocTestsUnderHighConcurrency(t, wazero.NewRuntimeConfigJIT)
-	// TODO: Add concurrent instantiation, invocation and release on a single store test case in https://github.com/tetratelabs/wazero/issues/293
-
+	singleModuleHighConcurrency(t, wazero.NewRuntimeConfigJIT)
 }
 
 func TestInterpreterConcurrency(t *testing.T) {
 	runAdhocTestsUnderHighConcurrency(t, wazero.NewRuntimeConfigInterpreter)
-	// TODO: Add concurrent instantiation, invocation and release on a single store test case in https://github.com/tetratelabs/wazero/issues/293
+	singleModuleHighConcurrency(t, wazero.NewRuntimeConfigInterpreter)
 }
 
 func runAdhocTestsUnderHighConcurrency(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConfig) {
@@ -54,3 +56,103 @@ func runAdhocTestUnderHighConcurrency(t *testing.T, newRuntimeConfig func() *waz
 	}
 	wg.Wait()
 }
+
+func singleModuleHighConcurrency(t *testing.T, newRuntimeConfig func() *wazero.RuntimeConfig) {
+	args := []uint64{1, 123}
+	exp := args[0] + args[1]
+	const goroutines = 1000
+	const delay = 100 * time.Millisecond
+
+	t.Run("single module", func(t *testing.T) {
+		r := wazero.NewRuntimeWithConfig(newRuntimeConfig())
+		imported, err := r.NewModuleBuilder("host").ExportFunctions(map[string]interface{}{
+			"delay": func() { time.Sleep(delay) },
+		}).Instantiate()
+		require.NoError(t, err)
+
+		source := []byte(`(module
+			(import "host" "delay" (func $delay ))
+			(func $delay_add
+				(param $value_1 i32) (param $value_2 i32)
+				(result i32)
+				local.get 0
+				local.get 1
+				i32.add
+				call $delay
+			)
+			(export "delay_add" (func $delay_add))
+		)`)
+
+		module, err := r.InstantiateModuleFromSource(source)
+		require.NoError(t, err)
+
+		t.Run("close importing module while in use", func(t *testing.T) {
+			fn := module.ExportedFunction("delay_add")
+			require.NotNil(t, fn)
+
+			var wg sync.WaitGroup
+			wg.Add(goroutines)
+			for i := 0; i < goroutines; i++ {
+				if i == 200 {
+					// Close the importing module.
+					module.Close()
+				} else if i == 400 {
+					// Re-instantiate the importing module, and swap the function.
+					module, err = r.InstantiateModuleFromSource(source)
+					require.NoError(t, err)
+					fn = module.ExportedFunction("delay_add")
+					require.NotNil(t, fn)
+				}
+				go func() {
+					defer wg.Done()
+					res, err := fn.Call(nil, args...)
+					require.NoError(t, err)
+					require.Equal(t, exp, res[0])
+				}()
+			}
+			wg.Wait()
+		})
+		t.Run("close imported module while in use", func(t *testing.T) {
+			fn := module.ExportedFunction("delay_add")
+			require.NotNil(t, fn)
+
+			var newImportedModuleShouldPanic bool = true
+			var wg sync.WaitGroup
+			wg.Add(goroutines)
+			for i := 0; i < goroutines; i++ {
+				if i == 200 {
+					// Close the imported module.
+					imported.Close()
+					// Re-instantiate the imported module but at this point importing module should
+					// not use the new one.
+					imported, err = r.NewModuleBuilder("host").ExportFunctions(map[string]interface{}{
+						"delay": func() {
+							time.Sleep(delay)
+							if newImportedModuleShouldPanic {
+								panic("unreachable")
+							}
+						},
+					}).Instantiate()
+					require.NoError(t, err)
+				} else if i == 400 {
+					// Re-instantiate the importing module, and swap the function which will use
+					// the new imported module.
+					module.Close()
+					module, err = r.InstantiateModuleFromSource(source)
+					require.NoError(t, err)
+					fn = module.ExportedFunction("delay_add")
+					require.NotNil(t, fn)
+					// The new imported module is now in use.
+					newImportedModuleShouldPanic = false
+				}
+				go func() {
+					defer wg.Done()
+					res, err := fn.Call(nil, 1, 123)
+					require.NoError(t, err)
+					require.Equal(t, exp, res[0])
+				}()
+			}
+			wg.Wait()
+		})
+	})
+}