Repository: tetratelabs/wazero
Pull request #438 (2022-04-04 07:46:55)

Moves close responsibility to the module context
This moves the responsibility to close a module from the engine to the
context. The reason for this is that the engine is what defines the
function. When a module is closed, it is often from an imported host
function. If it is on the module engine, it is easy to accidentally
close an imported module.

This refactors the WASI tests also, to ensure they aren't cheating too
much. This allows us to know for example "proc_exit" works without too
much trouble.


Patch:
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index bff97921..ea1f0be6 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -26,7 +26,6 @@ type ModuleEngine interface {
 	// The ctx's context.Context will be the outer-most ancestor of the argument to api.Function.
 	Call(ctx *ModuleContext, f *FunctionInstance, params ...uint64) (results []uint64, err error)
 
-	// CloseWithExitCode releases the resources allocated by functions in this ModuleEngine and ensures new calls (Call)
-	// return a sys.ExitError with the given code. This returns false if already closed.
-	CloseWithExitCode(exitCode uint32) (bool, error)
+	// Close releases all the function instances declared in this module.
+	Close()
 }
diff --git a/internal/wasm/interpreter/interpreter.go b/internal/wasm/interpreter/interpreter.go
index e295ef91..8b25c7bb 100644
--- a/internal/wasm/interpreter/interpreter.go
+++ b/internal/wasm/interpreter/interpreter.go
@@ -9,14 +9,12 @@ import (
 	"runtime/debug"
 	"strings"
 	"sync"
-	"sync/atomic"
 
 	"github.com/tetratelabs/wazero/internal/buildoptions"
 	"github.com/tetratelabs/wazero/internal/moremath"
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wasmruntime"
 	"github.com/tetratelabs/wazero/internal/wazeroir"
-	"github.com/tetratelabs/wazero/sys"
 )
 
 var callStackCeiling = buildoptions.CallStackCeiling
@@ -64,14 +62,6 @@ type moduleEngine struct {
 	// parentEngine holds *engine from which this module engine is created from.
 	parentEngine          *engine
 	importedFunctionCount uint32
-
-	// closed is the pointer used both to guard moduleEngine.CloseWithExitCode and to store the exit code.
-	//
-	// The update value is 1 + exitCode << 32. This ensures an exit code of zero isn't mistaken for never closed.
-	//
-	// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
-	// See /RATIONALE.md
-	closed uint64
 }
 
 // callEngine holds context per moduleEngine.Call, and shared across all the
@@ -148,7 +138,7 @@ type callFrame struct {
 
 type compiledFunction struct {
 	moduleEngine *moduleEngine
-	funcInstance *wasm.FunctionInstance
+	source       *wasm.FunctionInstance
 	body         []*interpreterOp
 	hostFn       *reflect.Value
 }
@@ -184,18 +174,18 @@ func (e *engine) NewModuleEngine(name string, importedFunctions, moduleFunctions
 		if f.Kind == wasm.FunctionKindWasm {
 			ir, err := wazeroir.Compile(f)
 			if err != nil {
-				me.doClose() // safe because the reference to me was never leaked.
+				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.doClose() // safe because the reference to me was never leaked.
+				me.Close()
 				return nil, fmt.Errorf("function[%d/%d] failed to convert wazeroir operations: %w", i, len(moduleFunctions)-1, err)
 			}
 		} else {
-			compiled = &compiledFunction{hostFn: f.GoFunc, funcInstance: f}
+			compiled = &compiledFunction{hostFn: f.GoFunc, source: f}
 		}
 		compiled.moduleEngine = me
 		me.compiledFunctions = append(me.compiledFunctions, compiled)
@@ -215,7 +205,7 @@ func (e *engine) NewModuleEngine(name string, importedFunctions, moduleFunctions
 func (me *moduleEngine) Release() error {
 	// Release all the function instances declared in this module.
 	for _, cf := range me.compiledFunctions[me.importedFunctionCount:] {
-		me.parentEngine.deleteCompiledFunction(cf.funcInstance)
+		me.parentEngine.deleteCompiledFunction(cf.source)
 	}
 	return nil
 }
@@ -223,7 +213,7 @@ func (me *moduleEngine) Release() error {
 // lowerIROps lowers the wazeroir operations to engine friendly struct.
 func (e *engine) lowerIROps(f *wasm.FunctionInstance,
 	ops []wazeroir.Operation) (*compiledFunction, error) {
-	ret := &compiledFunction{funcInstance: f}
+	ret := &compiledFunction{source: f}
 	labelAddress := map[string]uint64{}
 	onLabelAddressResolved := map[string][]func(addr uint64){}
 	for _, original := range ops {
@@ -500,12 +490,12 @@ func (me *moduleEngine) Name() string {
 }
 
 // Call implements the same method as documented on wasm.ModuleEngine.
-func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance, params ...uint64) (results []uint64, err error) {
+func (me *moduleEngine) Call(m *wasm.ModuleContext, f *wasm.FunctionInstance, params ...uint64) (results []uint64, err error) {
 	// Note: The input parameters are pre-validated, so a compiled function is only absent on close. Updates to
 	// compiledFunctions on close aren't locked, neither is this read.
 	compiled := me.compiledFunctions[f.Index]
 	if compiled == nil { // Lazy check the cause as it could be because the module was already closed.
-		if err = failIfClosed(me); err == nil {
+		if err = m.FailIfClosed(); err == nil {
 			panic(fmt.Errorf("BUG: %s.compiledFunctions[%d] was nil before close", me.name, f.Index))
 		}
 		return
@@ -521,7 +511,7 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 	defer func() {
 		// If the module closed during the call, and the call didn't err for another reason, set an ExitError.
 		if err == nil {
-			err = failIfClosed(me)
+			err = m.FailIfClosed()
 		}
 		// TODO: ^^ Will not fail if the function was imported from a closed module.
 
@@ -533,7 +523,7 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 			traces := make([]string, len(ce.frames))
 			for i := 0; i < len(traces); i++ {
 				frame := ce.popFrame()
-				name := frame.f.funcInstance.DebugName
+				name := frame.f.source.DebugName
 				// TODO: include DWARF symbols. See #58
 				traces[i] = fmt.Sprintf("\t%d: %s", i, name)
 			}
@@ -555,9 +545,9 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 		ce.push(param)
 	}
 	if f.Kind == wasm.FunctionKindWasm {
-		ce.callNativeFunc(ctx, compiled)
+		ce.callNativeFunc(m, compiled)
 	} else {
-		ce.callHostFunc(ctx, compiled)
+		ce.callHostFunc(m, compiled)
 	}
 	results = make([]uint64, len(f.Type.Results))
 	for i := range results {
@@ -566,20 +556,12 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 	return
 }
 
-// failIfClosed returns a sys.ExitError if moduleEngine.CloseWithExitCode was called.
-func failIfClosed(me *moduleEngine) error {
-	if closed := atomic.LoadUint64(&me.closed); closed != 0 {
-		return sys.NewExitError(me.name, uint32(closed>>32)) // Unpack the high order bits as the exit code.
-	}
-	return nil
-}
-
 func (ce *callEngine) callHostFunc(ctx *wasm.ModuleContext, f *compiledFunction) {
 	tp := f.hostFn.Type()
 	in := make([]reflect.Value, tp.NumIn())
 
 	wasmParamOffset := 0
-	if f.funcInstance.Kind != wasm.FunctionKindGoNoContext {
+	if f.source.Kind != wasm.FunctionKindGoNoContext {
 		wasmParamOffset = 1
 	}
 	for i := len(in) - 1; i >= wasmParamOffset; i-- {
@@ -601,11 +583,11 @@ func (ce *callEngine) callHostFunc(ctx *wasm.ModuleContext, f *compiledFunction)
 
 	// A host function is invoked with the calling frame's memory, which may be different if in another module.
 	if len(ce.frames) > 0 {
-		ctx = ctx.WithMemory(ce.frames[len(ce.frames)-1].f.funcInstance.Module.Memory)
+		ctx = ctx.WithMemory(ce.frames[len(ce.frames)-1].f.source.Module.Memory)
 	}
 
 	// Handle any special parameter zero
-	if val := wasm.GetHostFunctionCallContextValue(f.funcInstance.Kind, ctx); val != nil {
+	if val := wasm.GetHostFunctionCallContextValue(f.source.Kind, ctx); val != nil {
 		in[0] = *val
 	}
 
@@ -630,7 +612,7 @@ func (ce *callEngine) callHostFunc(ctx *wasm.ModuleContext, f *compiledFunction)
 
 func (ce *callEngine) callNativeFunc(ctx *wasm.ModuleContext, f *compiledFunction) {
 	frame := &callFrame{f: f}
-	moduleInst := f.funcInstance.Module
+	moduleInst := f.source.Module
 	memoryInst := moduleInst.Memory
 	globals := moduleInst.Globals
 	table := moduleInst.Table
@@ -689,7 +671,7 @@ func (ce *callEngine) callNativeFunc(ctx *wasm.ModuleContext, f *compiledFunctio
 				targetCompiledFunction, ok := table.Table[offset].(*compiledFunction)
 				if !ok {
 					panic(wasmruntime.ErrRuntimeInvalidTableAccess)
-				} else if uint64(targetCompiledFunction.funcInstance.TypeID) != op.us[1] {
+				} else if uint64(targetCompiledFunction.source.TypeID) != op.us[1] {
 					panic(wasmruntime.ErrRuntimeIndirectCallTypeMismatch)
 				}
 
@@ -1608,19 +1590,9 @@ func (ce *callEngine) callNativeFunc(ctx *wasm.ModuleContext, f *compiledFunctio
 	ce.popFrame()
 }
 
-// CloseWithExitCode implements the same method as documented on wasm.ModuleEngine.
-func (me *moduleEngine) CloseWithExitCode(exitCode uint32) (bool, error) {
-	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
-	if !atomic.CompareAndSwapUint64(&me.closed, 0, closed) {
-		return false, nil
-	}
-	me.doClose()
-	return true, nil
-}
-
-// doClose releases all the function instances declared in this module.
-func (me *moduleEngine) doClose() {
+// Close releases all the function instances declared in this module.
+func (me *moduleEngine) Close() {
 	for _, cf := range me.compiledFunctions[me.importedFunctionCount:] {
-		me.parentEngine.deleteCompiledFunction(cf.funcInstance)
+		me.parentEngine.deleteCompiledFunction(cf.source)
 	}
 }
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index a3f0f5d7..603779e5 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -8,14 +8,12 @@ import (
 	"runtime/debug"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"unsafe"
 
 	"github.com/tetratelabs/wazero/internal/buildoptions"
 	"github.com/tetratelabs/wazero/internal/wasm"
 	"github.com/tetratelabs/wazero/internal/wasmruntime"
 	"github.com/tetratelabs/wazero/internal/wazeroir"
-	"github.com/tetratelabs/wazero/sys"
 )
 
 type (
@@ -41,14 +39,6 @@ type (
 		// parentEngine holds *engine from which this module engine is created from.
 		parentEngine          *engine
 		importedFunctionCount uint32
-
-		// closed is the pointer used both to guard moduleEngine.CloseWithExitCode and to store the exit code.
-		//
-		// The update value is 1 + exitCode << 32. This ensures an exit code of zero isn't mistaken for never closed.
-		//
-		// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
-		// See /RATIONALE.md
-		closed uint64
 	}
 
 	// callEngine holds context per moduleEngine.Call, and shared across all the
@@ -282,7 +272,7 @@ const (
 	jitCallStatusCodeCallBuiltInFunction
 	// jitCallStatusCodeUnreachable means the function invocation reaches "unreachable" instruction.
 	jitCallStatusCodeUnreachable
-	// jitCallStatusCodeInvalidFloatToIntConversion means a invalid conversion of integer to floats happened.
+	// jitCallStatusCodeInvalidFloatToIntConversion means an invalid conversion of integer to floats happened.
 	jitCallStatusCodeInvalidFloatToIntConversion
 	// jitCallStatusCodeMemoryOutOfBounds means an out of bounds memory access happened.
 	jitCallStatusCodeMemoryOutOfBounds
@@ -383,12 +373,12 @@ func (e *engine) NewModuleEngine(name string, importedFunctions, moduleFunctions
 			compiled, err = compileHostFunction(f)
 		}
 		if err != nil {
-			me.doClose() // safe because the reference to me was never leaked.
+			me.Close() // safe because the reference to me was never leaked.
 			return nil, fmt.Errorf("function[%d/%d] %w", i, len(moduleFunctions)-1, err)
 		}
 
 		// As this uses mmap, we need a finalizer in case moduleEngine.Close was never called. Regardless, we need a
-		// finalizer due to how moduleEngine.doClose is implemented.
+		// finalizer due to how moduleEngine.Close is implemented.
 		e.setFinalizer(compiled, releaseCompiledFunction)
 
 		me.compiledFunctions = append(me.compiledFunctions, compiled)
@@ -404,7 +394,7 @@ func (e *engine) NewModuleEngine(name string, importedFunctions, moduleFunctions
 	return me, nil
 }
 
-// doClose is guarded by the caller with CAS, which means it happens only once. However, there is a race-condition
+// Close is guarded by the caller with CAS, which means it happens only once. However, there is a race-condition
 // inside the critical section: functions are removed from the parent engine, but there's no guard to prevent this
 // moduleInstance from making new calls. This means at inside the critical section there could be in-flight calls,
 // and even after it new calls can be made, given a reference to this moduleEngine.
@@ -427,7 +417,7 @@ func (e *engine) NewModuleEngine(name string, importedFunctions, moduleFunctions
 //    * Knowing which in-flight calls from other modules, that can use this module may be tricky
 //    * Pure wasm functions can be left to complete.
 //    * Host functions are the only unknowns (ex can do I/O) so they may need to be tracked.
-func (me *moduleEngine) doClose() {
+func (me *moduleEngine) Close() {
 	// Release all the function instances declared in this module.
 	for _, cf := range me.compiledFunctions[me.importedFunctionCount:] {
 		// NOTE: we still rely on the finalizer of cf until the notes on this function are addressed.
@@ -460,12 +450,12 @@ func (me *moduleEngine) Name() string {
 }
 
 // Call implements the same method as documented on wasm.ModuleEngine.
-func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance, params ...uint64) (results []uint64, err error) {
+func (me *moduleEngine) Call(m *wasm.ModuleContext, f *wasm.FunctionInstance, params ...uint64) (results []uint64, err error) {
 	// Note: The input parameters are pre-validated, so a compiled function is only absent on close. Updates to
 	// compiledFunctions on close aren't locked, neither is this read.
 	compiled := me.compiledFunctions[f.Index]
 	if compiled == nil { // Lazy check the cause as it could be because the module was already closed.
-		if err = failIfClosed(me); err == nil {
+		if err = m.FailIfClosed(); err == nil {
 			panic(fmt.Errorf("BUG: %s.compiledFunctions[%d] was nil before close", me.name, f.Index))
 		}
 		return
@@ -486,7 +476,7 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 	defer func() {
 		// If the module closed during the call, and the call didn't err for another reason, set an ExitError.
 		if err == nil {
-			err = failIfClosed(me)
+			err = m.FailIfClosed()
 		}
 		// TODO: ^^ Will not fail if the function was imported from a closed module.
 
@@ -525,9 +515,9 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 	}
 
 	if f.Kind == wasm.FunctionKindWasm {
-		ce.execWasmFunction(ctx, compiled)
+		ce.execWasmFunction(m, compiled)
 	} else {
-		ce.execHostFunction(f.Kind, compiled.source.GoFunc, ctx)
+		ce.execHostFunction(f.Kind, compiled.source.GoFunc, m)
 	}
 
 	// Note the top value is the tail of the results,
@@ -539,24 +529,6 @@ func (me *moduleEngine) Call(ctx *wasm.ModuleContext, f *wasm.FunctionInstance,
 	return
 }
 
-// failIfClosed returns a sys.ExitError if moduleEngine.CloseWithExitCode was called.
-func failIfClosed(me *moduleEngine) error {
-	if closed := atomic.LoadUint64(&me.closed); closed != 0 {
-		return sys.NewExitError(me.name, uint32(closed>>32)) // Unpack the high order bits as the exit code.
-	}
-	return nil
-}
-
-// CloseWithExitCode implements the same method as documented on wasm.ModuleEngine.
-func (me *moduleEngine) CloseWithExitCode(exitCode uint32) (bool, error) {
-	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
-	if !atomic.CompareAndSwapUint64(&me.closed, 0, closed) {
-		return false, nil
-	}
-	me.doClose()
-	return true, nil
-}
-
 func NewEngine() wasm.Engine {
 	return newEngine()
 }
diff --git a/internal/wasm/module_context.go b/internal/wasm/module_context.go
index c511e707..ebd93529 100644
--- a/internal/wasm/module_context.go
+++ b/internal/wasm/module_context.go
@@ -3,15 +3,18 @@ package wasm
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 
 	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/sys"
 )
 
 // compile time check to ensure ModuleContext implements api.Module
 var _ api.Module = &ModuleContext{}
 
-func NewModuleContext(ctx context.Context, store *Store, instance *ModuleInstance, sys *SysContext) *ModuleContext {
-	return &ModuleContext{ctx: ctx, memory: instance.Memory, module: instance, store: store, sys: sys}
+func NewModuleContext(ctx context.Context, store *Store, instance *ModuleInstance, Sys *SysContext) *ModuleContext {
+	zero := uint64(0)
+	return &ModuleContext{ctx: ctx, memory: instance.Memory, module: instance, store: store, Sys: Sys, closed: &zero}
 }
 
 // ModuleContext implements api.Module
@@ -23,9 +26,25 @@ type ModuleContext struct {
 	memory api.Memory
 	store  *Store
 
-	// sys is not exposed publicly. This is currently only used by wasi.
 	// Note: This is a part of ModuleContext so that scope is correct and Close is coherent.
-	sys *SysContext
+	// Sys is exposed only for WASI
+	Sys *SysContext
+
+	// closed is the pointer used both to guard moduleEngine.CloseWithExitCode and to store the exit code.
+	//
+	// The update value is 1 + exitCode << 32. This ensures an exit code of zero isn't mistaken for never closed.
+	//
+	// Note: Exclusively reading and updating this with atomics guarantees cross-goroutine observations.
+	// See /RATIONALE.md
+	closed *uint64
+}
+
+// FailIfClosed returns a sys.ExitError if CloseWithExitCode was called.
+func (m *ModuleContext) FailIfClosed() error {
+	if closed := atomic.LoadUint64(m.closed); closed != 0 {
+		return sys.NewExitError(m.module.Name, uint32(closed>>32)) // Unpack the high order bits as the exit code.
+	}
+	return nil
 }
 
 // Name implements the same method as documented on api.Module
@@ -36,7 +55,7 @@ func (m *ModuleContext) Name() string {
 // WithMemory allows overriding memory without re-allocation when the result would be the same.
 func (m *ModuleContext) WithMemory(memory *MemoryInstance) *ModuleContext {
 	if memory != nil && memory != m.memory { // only re-allocate if it will change the effective memory
-		return &ModuleContext{module: m.module, memory: memory, ctx: m.ctx, sys: m.sys}
+		return &ModuleContext{module: m.module, memory: memory, ctx: m.ctx, Sys: m.Sys, closed: m.closed}
 	}
 	return m
 }
@@ -51,15 +70,10 @@ func (m *ModuleContext) Context() context.Context {
 	return m.ctx
 }
 
-// Sys is exposed only for WASI.
-func (m *ModuleContext) Sys() *SysContext {
-	return m.sys
-}
-
 // WithContext implements the same method as documented on api.Module
 func (m *ModuleContext) WithContext(ctx context.Context) api.Module {
 	if ctx != nil && ctx != m.ctx { // only re-allocate if it will change the effective context
-		return &ModuleContext{module: m.module, memory: m.memory, ctx: ctx, sys: m.sys}
+		return &ModuleContext{module: m.module, memory: m.memory, ctx: ctx, Sys: m.Sys, closed: m.closed}
 	}
 	return m
 }
@@ -71,9 +85,13 @@ func (m *ModuleContext) Close() (err error) {
 
 // CloseWithExitCode implements the same method as documented on api.Module.
 func (m *ModuleContext) CloseWithExitCode(exitCode uint32) (err error) {
-	if err = m.store.CloseModuleWithExitCode(m.module.Name, exitCode); err != nil {
-		return err
-	} else if sys := m.sys; sys != nil { // ex nil if from ModuleBuilder
+	closed := uint64(1) + uint64(exitCode)<<32 // Store exitCode as high-order bits.
+	if !atomic.CompareAndSwapUint64(m.closed, 0, closed) {
+		return nil
+	}
+	m.module.Engine.Close()
+	m.store.deleteModule(m.Name())
+	if sys := m.Sys; sys != nil { // ex nil if from ModuleBuilder
 		return sys.Close()
 	}
 	return
@@ -99,20 +117,48 @@ func (m *ModuleContext) ExportedFunction(name string) api.Function {
 	if err != nil {
 		return nil
 	}
-	return exp.Function
+	if exp.Function.Module == m.module {
+		return exp.Function
+	} else {
+		return &importedFn{importingModule: m, importedFn: exp.Function}
+	}
+}
+
+// importedFn implements api.Function and ensures the call context of an imported function is the importing module.
+type importedFn struct {
+	importingModule *ModuleContext
+	importedFn      *FunctionInstance
+}
+
+// ParamTypes implements the same method as documented on api.Function
+func (f *importedFn) ParamTypes() []api.ValueType {
+	return f.importedFn.ParamTypes()
+}
+
+// ResultTypes implements the same method as documented on api.Function
+func (f *importedFn) ResultTypes() []api.ValueType {
+	return f.importedFn.ResultTypes()
+}
+
+// Call implements the same method as documented on api.Function
+func (f *importedFn) Call(m api.Module, params ...uint64) (ret []uint64, err error) {
+	if m == nil {
+		return f.importedFn.Call(f.importingModule, params...)
+	}
+	return f.importedFn.Call(m, params...)
 }
 
-// ParamTypes implements api.Function ParamTypes
+// ParamTypes implements the same method as documented on api.Function
 func (f *FunctionInstance) ParamTypes() []api.ValueType {
 	return f.Type.Params
 }
 
-// ResultTypes implements api.Function ResultTypes
+// ResultTypes implements the same method as documented on api.Function
 func (f *FunctionInstance) ResultTypes() []api.ValueType {
 	return f.Type.Results
 }
 
-// Call implements api.Function Call
+// Call implements the same method as documented on api.Function
 func (f *FunctionInstance) Call(m api.Module, params ...uint64) (ret []uint64, err error) {
 	mod := f.Module
 	modCtx, ok := m.(*ModuleContext)
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index d34af01a..53fc3050 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -183,16 +183,6 @@ func (m *ModuleInstance) buildExports(exports map[string]*Export) {
 		switch exp.Type {
 		case ExternTypeFunc:
 			ei = &ExportInstance{Type: exp.Type, Function: m.Functions[index]}
-			// The module instance of the host function is a fake that only includes the function and its types.
-			// We need to assign the ModuleInstance when re-exporting so that any memory defined in the target is
-			// available to the wasm.ModuleContext Memory.
-			//
-			// TODO(adrian): shouldn't this be race if multiple modules re-exporting?
-			// maybe we should limit this to only tests.
-			if ei.Function.GoFunc != nil {
-				ei.Function.Module = m
-				ei.Function.Index = index
-			}
 		case ExternTypeGlobal:
 			ei = &ExportInstance{Type: exp.Type, Global: m.Globals[index]}
 		case ExternTypeMemory:
@@ -327,18 +317,6 @@ func (s *Store) Instantiate(ctx context.Context, module *Module, name string, sy
 	return m.Ctx, nil
 }
 
-// CloseModuleWithExitCode deallocates resources if a module with the given name exists.
-func (s *Store) CloseModuleWithExitCode(moduleName string, exitCode uint32) error {
-	if m := s.module(moduleName); m == nil {
-		return nil
-	} else if ok, err := m.Engine.CloseWithExitCode(exitCode); err != nil {
-		return err
-	} else if ok {
-		s.deleteModule(moduleName)
-	}
-	return nil
-}
-
 // deleteModule makes the moduleName available for instantiation again.
 func (s *Store) deleteModule(moduleName string) {
 	s.mux.Lock()
diff --git a/wasi/wasi.go b/wasi/wasi.go
index 8d6b8833..24d9988c 100644
--- a/wasi/wasi.go
+++ b/wasi/wasi.go
@@ -1498,7 +1498,7 @@ func sysCtx(m api.Module) *wasm.SysContext {
 	if internal, ok := m.(*wasm.ModuleContext); !ok {
 		panic(fmt.Errorf("unsupported wasm.Module implementation: %v", m))
 	} else {
-		return internal.Sys()
+		return internal.Sys
 	}
 }
 


Test patch:
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index a90c2fdd..e22e46df 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -40,7 +40,7 @@ func RunTestEngine_NewModuleEngine(t *testing.T, et EngineTester) {
 	t.Run("sets module name", func(t *testing.T) {
 		me, err := e.NewModuleEngine(t.Name(), nil, nil, nil, nil)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, me, 0)
+		defer me.Close()
 		require.Equal(t, t.Name(), me.Name())
 	})
 }
@@ -63,7 +63,7 @@ func RunTestModuleEngine_Call(t *testing.T, et EngineTester) {
 	// Compile the module
 	me, err := e.NewModuleEngine(module.Name, nil, module.Functions, nil, nil)
 	require.NoError(t, err)
-	defer closeModuleEngineWithExitCode(t, me, 0)
+	defer me.Close()
 	linkModuleToEngine(module, me)
 
 	// Ensure the base case doesn't fail: A single parameter should work as that matches the function signature.
@@ -94,7 +94,7 @@ func RunTestEngine_NewModuleEngine_InitTable(t *testing.T, et EngineTester) {
 		// Instantiate the module, which has nothing but an empty table.
 		me, err := e.NewModuleEngine(t.Name(), importedFunctions, moduleFunctions, table, tableInit)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, me, 0)
+		defer me.Close()
 
 		// Since there are no elements to initialize, we expect the table to be nil.
 		require.Equal(t, table.Table, make([]interface{}, 2))
@@ -114,7 +114,7 @@ func RunTestEngine_NewModuleEngine_InitTable(t *testing.T, et EngineTester) {
 		// Instantiate the module whose table points to its own functions.
 		me, err := e.NewModuleEngine(t.Name(), importedFunctions, moduleFunctions, table, tableInit)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, me, 0)
+		defer me.Close()
 
 		// The functions mapped to the table are defined in the same moduleEngine
 		require.Equal(t, table.Table, et.InitTable(me, table.Min, tableInit))
@@ -134,12 +134,12 @@ func RunTestEngine_NewModuleEngine_InitTable(t *testing.T, et EngineTester) {
 		// Imported functions are compiled before the importing module is instantiated.
 		imported, err := e.NewModuleEngine(t.Name(), nil, importedFunctions, nil, nil)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, imported, 0)
+		defer imported.Close()
 
 		// Instantiate the importing module, which is whose table is initialized.
 		importing, err := e.NewModuleEngine(t.Name(), importedFunctions, moduleFunctions, table, tableInit)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, importing, 0)
+		defer importing.Close()
 
 		// A moduleEngine's compiled function slice includes its imports, so the offsets is absolute.
 		require.Equal(t, table.Table, et.InitTable(importing, table.Min, tableInit))
@@ -164,12 +164,12 @@ func RunTestEngine_NewModuleEngine_InitTable(t *testing.T, et EngineTester) {
 		// Imported functions are compiled before the importing module is instantiated.
 		imported, err := e.NewModuleEngine(t.Name(), nil, importedFunctions, nil, nil)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, imported, 0)
+		defer imported.Close()
 
 		// Instantiate the importing module, which is whose table is initialized.
 		importing, err := e.NewModuleEngine(t.Name(), importedFunctions, moduleFunctions, table, tableInit)
 		require.NoError(t, err)
-		defer closeModuleEngineWithExitCode(t, importing, 0)
+		defer importing.Close()
 
 		// A moduleEngine's compiled function slice includes its imports, so the offsets are absolute.
 		require.Equal(t, table.Table, et.InitTable(importing, table.Min, tableInit))
@@ -203,7 +203,7 @@ func runTestModuleEngine_Call_HostFn_ModuleContext(t *testing.T, et EngineTester
 
 	me, err := e.NewModuleEngine(t.Name(), nil, module.Functions, nil, nil)
 	require.NoError(t, err)
-	defer closeModuleEngineWithExitCode(t, me, 0)
+	defer me.Close()
 
 	t.Run("defaults to module memory when call stack empty", func(t *testing.T) {
 		// When calling a host func directly, there may be no stack. This ensures the module's memory is used.
@@ -219,38 +219,44 @@ func RunTestModuleEngine_Call_HostFn(t *testing.T, et EngineTester) {
 
 	e := et.NewEngine()
 
-	imported, importedMe := setupCallTests(t, e)
-	defer closeModuleEngineWithExitCode(t, importedMe, 0)
+	imported, importedMe, importing, importingMe := setupCallTests(t, e)
+	defer importingMe.Close()
+	defer importedMe.Close()
 
 	// Ensure the base case doesn't fail: A single parameter should work as that matches the function signature.
 	tests := []struct {
-		name string
-		me   wasm.ModuleEngine
-		fn   *wasm.FunctionInstance
+		name   string
+		module *wasm.ModuleContext
+		fn     *wasm.FunctionInstance
 	}{
 		{
-			name: wasmFnName,
-			me:   importedMe,
-			fn:   imported.Exports[wasmFnName].Function,
+			name:   wasmFnName,
+			module: imported.Ctx,
+			fn:     imported.Exports[wasmFnName].Function,
 		},
 		{
-			name: hostFnName,
-			me:   importedMe,
-			fn:   imported.Exports[hostFnName].Function,
+			name:   hostFnName,
+			module: imported.Ctx,
+			fn:     imported.Exports[hostFnName].Function,
 		},
 		{
-			name: callHostFnName,
-			me:   importedMe,
-			fn:   imported.Exports[callHostFnName].Function,
+			name:   callHostFnName,
+			module: imported.Ctx,
+			fn:     imported.Exports[callHostFnName].Function,
+		},
+		{
+			name:   callImportCallHostFnName,
+			module: importing.Ctx,
+			fn:     importing.Exports[callImportCallHostFnName].Function,
 		},
 	}
 	for _, tt := range tests {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
+			m := tc.module
 			f := tc.fn
-			m := f.Module
-			results, err := m.Engine.Call(m.Ctx, f, 1)
+			results, err := f.Module.Engine.Call(m, f, 1)
 			require.NoError(t, err)
 			require.Equal(t, uint64(1), results[0])
 		})
@@ -260,12 +266,13 @@ func RunTestModuleEngine_Call_HostFn(t *testing.T, et EngineTester) {
 func RunTestModuleEngine_Call_Errors(t *testing.T, et EngineTester) {
 	e := et.NewEngine()
 
-	imported, importedMe := setupCallTests(t, e)
-	defer closeModuleEngineWithExitCode(t, importedMe, 0)
+	imported, importedMe, importing, importingMe := setupCallTests(t, e)
+	defer importingMe.Close()
+	defer importedMe.Close()
 
 	tests := []struct {
 		name        string
-		me          wasm.ModuleEngine
+		module      *wasm.ModuleContext
 		fn          *wasm.FunctionInstance
 		input       []uint64
 		expectedErr string
@@ -273,99 +280,123 @@ func RunTestModuleEngine_Call_Errors(t *testing.T, et EngineTester) {
 		{
 			name:        "host function not enough parameters",
 			input:       []uint64{},
-			me:          importedMe,
+			module:      imported.Ctx,
 			fn:          imported.Exports[hostFnName].Function,
 			expectedErr: `expected 1 params, but passed 0`,
 		},
 		{
 			name:        "host function too many parameters",
 			input:       []uint64{1, 2},
-			me:          importedMe,
+			module:      imported.Ctx,
 			fn:          imported.Exports[hostFnName].Function,
 			expectedErr: `expected 1 params, but passed 2`,
 		},
 		{
 			name:        "wasm function not enough parameters",
 			input:       []uint64{},
-			me:          importedMe,
+			module:      imported.Ctx,
 			fn:          imported.Exports[wasmFnName].Function,
 			expectedErr: `expected 1 params, but passed 0`,
 		},
 		{
 			name:        "wasm function too many parameters",
 			input:       []uint64{1, 2},
-			me:          importedMe,
+			module:      imported.Ctx,
 			fn:          imported.Exports[wasmFnName].Function,
 			expectedErr: `expected 1 params, but passed 2`,
 		},
 		{
-			name:  "wasm function panics with wasmruntime.Error",
-			input: []uint64{0},
-			me:    importedMe,
-			fn:    imported.Exports[wasmFnName].Function,
+			name:   "wasm function panics with wasmruntime.Error",
+			input:  []uint64{0},
+			module: imported.Ctx,
+			fn:     imported.Exports[wasmFnName].Function,
 			expectedErr: `wasm runtime error: integer divide by zero
 wasm backtrace:
 	0: wasm_div_by`,
 		},
 		{
-			name:  "host function that panics",
-			input: []uint64{math.MaxUint32},
-			me:    importedMe,
-			fn:    imported.Exports[hostFnName].Function,
+			name:   "host function that panics",
+			input:  []uint64{math.MaxUint32},
+			module: imported.Ctx,
+			fn:     imported.Exports[hostFnName].Function,
 			expectedErr: `wasm runtime error: host-function panic
 wasm backtrace:
 	0: host_div_by`,
 		},
 		{
-			name:  "host function panics with runtime.Error",
-			input: []uint64{0},
-			me:    importedMe,
-			fn:    imported.Exports[hostFnName].Function, // TODO: This should be a normal runtime error
+			name:   "host function panics with runtime.Error",
+			input:  []uint64{0},
+			module: imported.Ctx,
+			fn:     imported.Exports[hostFnName].Function, // TODO: This should be a normal runtime error
 			expectedErr: `wasm runtime error: runtime error: integer divide by zero
 wasm backtrace:
 	0: host_div_by`,
 		},
 		{
-			name:  "wasm calls host function that panics",
-			input: []uint64{math.MaxUint32},
-			me:    importedMe,
-			fn:    imported.Exports[callHostFnName].Function,
+			name:   "wasm calls host function that panics",
+			input:  []uint64{math.MaxUint32},
+			module: imported.Ctx,
+			fn:     imported.Exports[callHostFnName].Function,
 			expectedErr: `wasm runtime error: host-function panic
 wasm backtrace:
 	0: host_div_by
 	1: call->host_div_by`,
 		},
 		{
-			name:  "wasm calls imported wasm that calls host function panics with runtime.Error",
-			input: []uint64{0},
-			me:    importedMe,
-			fn:    imported.Exports[callHostFnName].Function,
+			name:   "wasm calls imported wasm that calls host function panics with runtime.Error",
+			input:  []uint64{0},
+			module: importing.Ctx,
+			fn:     importing.Exports[callImportCallHostFnName].Function,
 			expectedErr: `wasm runtime error: runtime error: integer divide by zero
 wasm backtrace:
 	0: host_div_by
-	1: call->host_div_by`,
+	1: call->host_div_by
+	2: call_import->call->host_div_by`,
+		},
+		{
+			name:   "wasm calls imported wasm that calls host function that panics",
+			input:  []uint64{math.MaxUint32},
+			module: importing.Ctx,
+			fn:     importing.Exports[callImportCallHostFnName].Function,
+			expectedErr: `wasm runtime error: host-function panic
+wasm backtrace:
+	0: host_div_by
+	1: call->host_div_by
+	2: call_import->call->host_div_by`,
+		},
+		{
+			name:   "wasm calls imported wasm calls host function panics with runtime.Error",
+			input:  []uint64{0},
+			module: importing.Ctx,
+			fn:     importing.Exports[callImportCallHostFnName].Function,
+			expectedErr: `wasm runtime error: runtime error: integer divide by zero
+wasm backtrace:
+	0: host_div_by
+	1: call->host_div_by
+	2: call_import->call->host_div_by`,
 		},
 	}
 	for _, tt := range tests {
 		tc := tt
 		t.Run(tc.name, func(t *testing.T) {
+			m := tc.module
 			f := tc.fn
-			m := f.Module
-			_, err := m.Engine.Call(m.Ctx, f, tc.input...)
+			_, err := f.Module.Engine.Call(m, f, tc.input...)
 			require.EqualError(t, err, tc.expectedErr)
 
 			// Ensure the module still works
-			ret, err := m.Engine.Call(m.Ctx, f, 1)
+			results, err := f.Module.Engine.Call(m, f, 1)
 			require.NoError(t, err)
-			require.Equal(t, uint64(1), ret[0])
+			require.Equal(t, uint64(1), results[0])
 		})
 	}
 }
 
 const (
-	wasmFnName     = "wasm_div_by"
-	hostFnName     = "host_div_by"
-	callHostFnName = "call->" + hostFnName
+	wasmFnName               = "wasm_div_by"
+	hostFnName               = "host_div_by"
+	callHostFnName           = "call->" + hostFnName
+	callImportCallHostFnName = "call_import->" + callHostFnName
 )
 
 // (func (export "wasm_div_by") (param i32) (result i32) (i32.div_u (i32.const 1) (local.get 0)))
@@ -378,7 +409,7 @@ func divBy(d uint32) uint32 {
 	return 1 / d // go panics if d == 0
 }
 
-func setupCallTests(t *testing.T, e wasm.Engine) (*wasm.ModuleInstance, wasm.ModuleEngine) {
+func setupCallTests(t *testing.T, e wasm.Engine) (*wasm.ModuleInstance, wasm.ModuleEngine, *wasm.ModuleInstance, wasm.ModuleEngine) {
 	i32 := wasm.ValueTypeI32
 	ft := &wasm.FunctionType{Params: []wasm.ValueType{i32}, Results: []wasm.ValueType{i32}}
 	wasmFn := &wasm.FunctionInstance{
@@ -412,7 +443,30 @@ func setupCallTests(t *testing.T, e wasm.Engine) (*wasm.ModuleInstance, wasm.Mod
 	require.NoError(t, err)
 	linkModuleToEngine(imported, importedMe)
 
-	return imported, importedMe
+	// To test stack traces, call the same function from another module
+	importing := &wasm.ModuleInstance{Name: t.Name() + "-importing"}
+
+	// Don't add imported functions yet as NewModuleEngine requires them split.
+	importedFunctions := []*wasm.FunctionInstance{callHostFn}
+
+	// Add the exported function.
+	callImportedHostFn := &wasm.FunctionInstance{
+		Kind:  wasm.FunctionKindWasm,
+		Type:  ft,
+		Body:  []byte{wasm.OpcodeLocalGet, 0, wasm.OpcodeCall, 0 /* only one imported function */, wasm.OpcodeEnd},
+		Index: 1, // after import
+	}
+	addFunction(importing, callImportCallHostFnName, callImportedHostFn)
+
+	// Compile the importing module
+	importingMe, err := e.NewModuleEngine(importing.Name, importedFunctions, importing.Functions, nil, nil)
+	require.NoError(t, err)
+	linkModuleToEngine(importing, importingMe)
+
+	// Add the imported functions back to the importing module.
+	importing.Functions = append(importedFunctions, importing.Functions...)
+
+	return imported, importedMe, importing, importingMe
 }
 
 // linkModuleToEngine assigns fields that wasm.Store would on instantiation. These includes fields both interpreter and
@@ -439,10 +493,3 @@ func addFunction(module *wasm.ModuleInstance, funcName string, fn *wasm.Function
 	// This link is essential for all engines. For example, functions call other functions defined in the same module.
 	fn.Module = module
 }
-
-// closeModuleEngineWithExitCode allows unit tests to check `CloseWithExitCode` didn't err.
-func closeModuleEngineWithExitCode(t *testing.T, me wasm.ModuleEngine, exitCode uint32) bool {
-	ok, err := me.CloseWithExitCode(exitCode)
-	require.NoError(t, err)
-	return ok
-}
diff --git a/internal/wasm/interpreter/interpreter_test.go b/internal/wasm/interpreter/interpreter_test.go
index 4e6f2f55..a139a2f4 100644
--- a/internal/wasm/interpreter/interpreter_test.go
+++ b/internal/wasm/interpreter/interpreter_test.go
@@ -129,7 +129,7 @@ func TestInterpreter_CallEngine_callNativeFunc_signExtend(t *testing.T) {
 				ce := &callEngine{}
 				f := &compiledFunction{
 					moduleEngine: &moduleEngine{},
-					funcInstance: &wasm.FunctionInstance{Module: &wasm.ModuleInstance{}},
+					source:       &wasm.FunctionInstance{Module: &wasm.ModuleInstance{}},
 					body: []*interpreterOp{
 						{kind: wazeroir.OperationKindConstI32, us: []uint64{uint64(uint32(tc.in))}},
 						{kind: translateToIROperationKind(tc.opcode)},
@@ -182,7 +182,7 @@ func TestInterpreter_CallEngine_callNativeFunc_signExtend(t *testing.T) {
 				ce := &callEngine{}
 				f := &compiledFunction{
 					moduleEngine: &moduleEngine{},
-					funcInstance: &wasm.FunctionInstance{Module: &wasm.ModuleInstance{}},
+					source:       &wasm.FunctionInstance{Module: &wasm.ModuleInstance{}},
 					body: []*interpreterOp{
 						{kind: wazeroir.OperationKindConstI64, us: []uint64{uint64(tc.in)}},
 						{kind: translateToIROperationKind(tc.opcode)},
@@ -289,9 +289,7 @@ func TestInterpreter_Close(t *testing.T) {
 				require.Contains(t, e.compiledFunctions, f)
 			}
 
-			closed, err := me.CloseWithExitCode(0)
-			require.True(t, closed)
-			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 923adb27..8d1224d1 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -224,7 +224,7 @@ func TestJIT_NewModuleEngine_CompiledFunctions(t *testing.T) {
 	}
 	modE, err := e.NewModuleEngine(t.Name(), nil, importedFunctions, nil, nil)
 	require.NoError(t, err)
-	defer modE.CloseWithExitCode(0) //nolint
+	defer modE.Close()
 	imported := modE.(*moduleEngine)
 
 	importingFinalizer := fakeFinalizer{}
@@ -238,7 +238,7 @@ func TestJIT_NewModuleEngine_CompiledFunctions(t *testing.T) {
 
 	modE, err = e.NewModuleEngine(t.Name(), importedFunctions, moduleFunctions, nil, nil)
 	require.NoError(t, err)
-	defer modE.CloseWithExitCode(0) //nolint
+	defer modE.Close()
 	importing := modE.(*moduleEngine)
 
 	// Ensure the importing module didn't try to finalize the imported functions.
@@ -360,12 +360,7 @@ func TestJIT_ModuleEngine_Close(t *testing.T) {
 				require.Contains(t, e.compiledFunctions, f)
 			}
 
-			closed, err := importing.CloseWithExitCode(0)
-			require.True(t, closed)
-			require.NoError(t, err)
-
-			// Closing should flip the status bit, so that it cannot be closed again.
-			require.Equal(t, uint64(1), importing.(*moduleEngine).closed)
+			importing.Close()
 
 			// Closing the importing module shouldn't delete the imported functions from the engine.
 			require.Len(t, e.compiledFunctions, len(tc.importedFunctions))
@@ -379,9 +374,7 @@ func TestJIT_ModuleEngine_Close(t *testing.T) {
 			}
 
 			if len(tc.importedFunctions) > 0 {
-				closed, err = imported.CloseWithExitCode(0)
-				require.True(t, closed)
-				require.NoError(t, err)
+				imported.Close()
 			}
 
 			// When all modules are closed, the engine should be empty.
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 01644c5f..d80d68c4 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -105,7 +105,7 @@ func TestStore_Instantiate(t *testing.T) {
 		require.Equal(t, s.modules[""], mod.module)
 		require.Equal(t, s.modules[""].Memory, mod.memory)
 		require.Equal(t, s, mod.store)
-		require.Equal(t, sys, mod.sys)
+		require.Equal(t, sys, mod.Sys)
 	})
 }
 
@@ -153,21 +153,21 @@ func TestStore_CloseModule(t *testing.T) {
 			}, importingModuleName, nil)
 			require.NoError(t, err)
 
-			_, ok := s.modules[importedModuleName]
+			imported, ok := s.modules[importedModuleName]
 			require.True(t, ok)
 
-			_, ok = s.modules[importingModuleName]
+			importing, ok := s.modules[importingModuleName]
 			require.True(t, ok)
 
 			// Close the importing module
-			require.NoError(t, s.CloseModuleWithExitCode(importingModuleName, 0))
+			require.NoError(t, importing.Ctx.CloseWithExitCode(0))
 			require.NotContains(t, s.modules, importingModuleName)
 
 			// Can re-close the importing module
-			require.NoError(t, s.CloseModuleWithExitCode(importingModuleName, 0))
+			require.NoError(t, importing.Ctx.CloseWithExitCode(0))
 
 			// Now we close the imported module.
-			require.NoError(t, s.CloseModuleWithExitCode(importedModuleName, 0))
+			require.NoError(t, imported.Ctx.CloseWithExitCode(0))
 			require.Nil(t, s.modules[importedModuleName])
 			require.NotContains(t, s.modules, importedModuleName)
 		})
@@ -210,14 +210,14 @@ func TestStore_hammer(t *testing.T) {
 	hammer.NewHammer(t, P, N).Run(func(name string) {
 		mod, instantiateErr := s.Instantiate(context.Background(), importingModule, name, DefaultSysContext())
 		require.NoError(t, instantiateErr)
-		require.NoError(t, s.CloseModuleWithExitCode(mod.Name(), 0))
+		require.NoError(t, mod.CloseWithExitCode(0))
 	}, nil)
 	if t.Failed() {
 		return // At least one test failed, so return now.
 	}
 
 	// Close the imported module.
-	require.NoError(t, s.CloseModuleWithExitCode(imported.Name(), 0))
+	require.NoError(t, imported.CloseWithExitCode(0))
 
 	// All instances are freed.
 	require.Len(t, s.modules, 0)
@@ -311,102 +311,81 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 	})
 }
 
-func TestStore_ExportImportedHostFunction(t *testing.T) {
-	m, err := NewHostModule("host", map[string]interface{}{"host_fn": func(api.Module) {}})
+func TestModuleContext_ExportedFunction(t *testing.T) {
+	host, err := NewHostModule("host", map[string]interface{}{"host_fn": func(api.Module) {}})
 	require.NoError(t, err)
 
 	s := newStore()
 
 	// Add the host module
-	_, err = s.Instantiate(context.Background(), m, m.NameSection.ModuleName, nil)
+	imported, err := s.Instantiate(context.Background(), host, host.NameSection.ModuleName, nil)
 	require.NoError(t, err)
+	defer imported.Close()
 
-	t.Run("Module is the importing module", func(t *testing.T) {
-		_, err = s.Instantiate(context.Background(), &Module{
+	t.Run("imported function", func(t *testing.T) {
+		importing, err := s.Instantiate(context.Background(), &Module{
 			TypeSection:   []*FunctionType{{}},
 			ImportSection: []*Import{{Type: ExternTypeFunc, Module: "host", Name: "host_fn", DescFunc: 0}},
 			MemorySection: &Memory{Min: 1},
 			ExportSection: map[string]*Export{"host.fn": {Type: ExternTypeFunc, Name: "host.fn", Index: 0}},
 		}, "test", nil)
 		require.NoError(t, err)
+		defer importing.Close()
 
-		mod, ok := s.modules["test"]
-		require.True(t, ok)
+		fn := importing.ExportedFunction("host.fn")
+		require.NotNil(t, fn)
 
-		ei, err := mod.getExport("host.fn", ExternTypeFunc)
-		require.NoError(t, err)
-		// We expect the host function to be called in context of the importing module.
-		// Otherwise, it would be the pseudo-module of the host, which only includes types and function definitions.
-		// Notably, this ensures the host function call context has the correct memory (from the importing module).
-		require.Equal(t, s.modules["test"], ei.Function.Module)
+		require.Equal(t, fn.(*importedFn).importedFn, imported.ExportedFunction("host_fn"))
+		require.Equal(t, fn.(*importedFn).importingModule, importing)
 	})
 }
 
 func TestFunctionInstance_Call(t *testing.T) {
-	type key string
-	storeCtx := context.WithValue(context.Background(), key("wa"), "zero")
-	notStoreCtx := context.WithValue(context.Background(), key("wazer"), "o")
+	store := NewStore(&mockEngine{shouldCompileFail: false, callFailIndex: -1}, Features20191205)
 
 	// Add the host module
 	functionName := "fn"
+
+	// This is a fake engine, so we don't capture inside the function body.
 	m, err := NewHostModule("host",
 		map[string]interface{}{functionName: func(api.Module) {}},
 	)
 	require.NoError(t, err)
 
-	tests := []struct {
-		name     string
-		ctx      context.Context
-		expected context.Context
-	}{
-		{
-			name:     "nil defaults to store context",
-			ctx:      nil,
-			expected: storeCtx,
-		},
-		{
-			name:     "set overrides store context",
-			ctx:      notStoreCtx,
-			expected: notStoreCtx,
-		},
-	}
-
-	for _, tt := range tests {
-		tc := tt
-
-		t.Run(tc.name, func(t *testing.T) {
-			store := NewStore(&mockEngine{shouldCompileFail: false, callFailIndex: -1}, Features20191205)
-
-			// Add the host module
-			hm, err := store.Instantiate(storeCtx, m, "host", nil)
-			require.NoError(t, err)
-
-			// Make a module to import the function
-			mod, err := store.Instantiate(storeCtx, &Module{
-				TypeSection: []*FunctionType{{}},
-				ImportSection: []*Import{{
-					Type:     ExternTypeFunc,
-					Module:   hm.Name(),
-					Name:     functionName,
-					DescFunc: 0,
-				}},
-				MemorySection: &Memory{Min: 1},
-				ExportSection: map[string]*Export{functionName: {Type: ExternTypeFunc, Name: functionName, Index: 0}},
-			}, "test", nil)
-			require.NoError(t, err)
-
-			// This fails if the function wasn't invoked, or had an unexpected context.
-			_, err = mod.ExportedFunction(functionName).Call(mod.WithContext(tc.ctx))
-			require.NoError(t, err)
+	// Add the host module
+	imported, err := store.Instantiate(context.Background(), m, "host", nil)
+	require.NoError(t, err)
+	defer imported.Close()
+
+	// Make a module to import the function
+	importing, err := store.Instantiate(context.Background(), &Module{
+		TypeSection: []*FunctionType{{}},
+		ImportSection: []*Import{{
+			Type:     ExternTypeFunc,
+			Module:   imported.Name(),
+			Name:     functionName,
+			DescFunc: 0,
+		}},
+		MemorySection: &Memory{Min: 1},
+		ExportSection: map[string]*Export{functionName: {Type: ExternTypeFunc, Name: functionName, Index: 0}},
+	}, "test", nil)
+	require.NoError(t, err)
+	defer imported.Close()
 
-			modEngine := store.modules["test"].Engine.(*mockModuleEngine)
-			if tc.expected == nil {
-				require.Nil(t, modEngine.ctx)
-			} else {
-				require.Equal(t, tc.expected, modEngine.ctx.Context())
-			}
-		})
-	}
+	fn := importing.ExportedFunction(functionName)
+	// me.ctx will hold the last seen context.
+	me := imported.module.Engine.(*mockModuleEngine)
+	t.Run("nil defaults to current module", func(t *testing.T) {
+		_, err := fn.Call(nil)
+		require.NoError(t, err)
+		require.Equal(t, importing, me.ctx)
+	})
+	t.Run("override current module context", func(t *testing.T) {
+		ctx := importing.WithContext(context.TODO())
+		_, err := fn.Call(ctx)
+		require.NoError(t, err)
+		require.Equal(t, ctx, me.ctx)
+	})
 }
 
 type mockEngine struct {
@@ -447,9 +426,8 @@ func (e *mockModuleEngine) Call(ctx *ModuleContext, f *FunctionInstance, _ ...ui
 	return
 }
 
-// CloseWithExitCode implements the same method as documented on wasm.ModuleEngine.
-func (e *mockModuleEngine) CloseWithExitCode(exitCode uint32) (bool, error) {
-	return true, nil
+// Close implements the same method as documented on wasm.ModuleEngine.
+func (e *mockModuleEngine) Close() {
 }
 
 func TestStore_getTypeInstance(t *testing.T) {
diff --git a/tests/engine/adhoc_test.go b/tests/engine/adhoc_test.go
index 56790df8..18a4d41a 100644
--- a/tests/engine/adhoc_test.go
+++ b/tests/engine/adhoc_test.go
@@ -111,21 +111,27 @@ func testRecursiveEntry(t *testing.T, r wazero.Runtime) {
 	require.NoError(t, err)
 }
 
+func TestImportedAndExportedFunc(t *testing.T) {
+	r := wazero.NewRuntime()
+	testImportedAndExportedFunc(t, r)
+}
+
 // testImportedAndExportedFunc fails if the engine cannot call an "imported-and-then-exported-back" function
 // Notably, this uses memory, which ensures api.Module is valid in both interpreter and JIT engines.
 func testImportedAndExportedFunc(t *testing.T, r wazero.Runtime) {
 	var memory *wasm.MemoryInstance
-	storeInt := func(nil api.Module, offset uint32, val uint64) uint32 {
-		if !nil.Memory().WriteUint64Le(offset, val) {
+	storeInt := func(m api.Module, offset uint32, val uint64) uint32 {
+		if !m.Memory().WriteUint64Le(offset, val) {
 			return 1
 		}
 		// sneak a reference to the memory, so we can check it later
-		memory = nil.Memory().(*wasm.MemoryInstance)
+		memory = m.Memory().(*wasm.MemoryInstance)
 		return 0
 	}
 
-	_, err := r.NewModuleBuilder("").ExportFunction("store_int", storeInt).Instantiate()
+	host, err := r.NewModuleBuilder("").ExportFunction("store_int", storeInt).Instantiate()
 	require.NoError(t, err)
+	defer host.Close()
 
 	module, err := r.InstantiateModuleFromCode([]byte(`(module $test
 		(import "" "store_int"
@@ -139,7 +145,8 @@ func testImportedAndExportedFunc(t *testing.T, r wazero.Runtime) {
 	defer module.Close()
 
 	// Call store_int and ensure it didn't return an error code.
-	results, err := module.ExportedFunction("store_int").Call(nil, 1, math.MaxUint64)
+	fn := module.ExportedFunction("store_int")
+	results, err := fn.Call(nil, 1, math.MaxUint64)
 	require.NoError(t, err)
 	require.Equal(t, uint64(0), results[0])
 
diff --git a/wasi/wasi_test.go b/wasi/wasi_test.go
index 228e0210..58876643 100644
--- a/wasi/wasi_test.go
+++ b/wasi/wasi_test.go
@@ -2,7 +2,6 @@ package wasi
 
 import (
 	"bytes"
-	"context"
 	_ "embed"
 	"errors"
 	"fmt"
@@ -17,17 +16,13 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/wasm"
-	"github.com/tetratelabs/wazero/internal/wasm/interpreter"
-	"github.com/tetratelabs/wazero/internal/wasm/text"
 	"github.com/tetratelabs/wazero/sys"
 )
 
-const moduleName = "test"
-
 func TestSnapshotPreview1_ArgsGet(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext([]string{"a", "bc"}, nil, nil)
 	require.NoError(t, err)
 
@@ -42,7 +37,7 @@ func TestSnapshotPreview1_ArgsGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	a, mod, fn := instantiateModule(t, ctx, functionArgsGet, importArgsGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionArgsGet, importArgsGet, sysCtx)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.ArgsGet", func(t *testing.T) {
@@ -72,11 +67,10 @@ func TestSnapshotPreview1_ArgsGet(t *testing.T) {
 }
 
 func TestSnapshotPreview1_ArgsGet_Errors(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext([]string{"a", "bc"}, nil, nil)
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionArgsGet, importArgsGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionArgsGet, importArgsGet, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -123,7 +117,6 @@ func TestSnapshotPreview1_ArgsGet_Errors(t *testing.T) {
 }
 
 func TestSnapshotPreview1_ArgsSizesGet(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext([]string{"a", "bc"}, nil, nil)
 	require.NoError(t, err)
 
@@ -137,7 +130,7 @@ func TestSnapshotPreview1_ArgsSizesGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	a, mod, fn := instantiateModule(t, ctx, functionArgsSizesGet, importArgsSizesGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionArgsSizesGet, importArgsSizesGet, sysCtx)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.ArgsSizesGet", func(t *testing.T) {
@@ -167,11 +160,10 @@ func TestSnapshotPreview1_ArgsSizesGet(t *testing.T) {
 }
 
 func TestSnapshotPreview1_ArgsSizesGet_Errors(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext([]string{"a", "bc"}, nil, nil)
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionArgsSizesGet, importArgsSizesGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionArgsSizesGet, importArgsSizesGet, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -215,7 +207,6 @@ func TestSnapshotPreview1_ArgsSizesGet_Errors(t *testing.T) {
 }
 
 func TestSnapshotPreview1_EnvironGet(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, []string{"a=b", "b=cd"}, nil)
 	require.NoError(t, err)
 
@@ -231,7 +222,7 @@ func TestSnapshotPreview1_EnvironGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	a, mod, fn := instantiateModule(t, ctx, functionEnvironGet, importEnvironGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionEnvironGet, importEnvironGet, sysCtx)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.EnvironGet", func(t *testing.T) {
@@ -261,11 +252,10 @@ func TestSnapshotPreview1_EnvironGet(t *testing.T) {
 }
 
 func TestSnapshotPreview1_EnvironGet_Errors(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, []string{"a=bc", "b=cd"}, nil)
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionEnvironGet, importEnvironGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionEnvironGet, importEnvironGet, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -311,7 +301,6 @@ func TestSnapshotPreview1_EnvironGet_Errors(t *testing.T) {
 }
 
 func TestSnapshotPreview1_EnvironSizesGet(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, []string{"a=b", "b=cd"}, nil)
 	require.NoError(t, err)
 
@@ -325,7 +314,7 @@ func TestSnapshotPreview1_EnvironSizesGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	a, mod, fn := instantiateModule(t, ctx, functionEnvironSizesGet, importEnvironSizesGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionEnvironSizesGet, importEnvironSizesGet, sysCtx)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.EnvironSizesGet", func(t *testing.T) {
@@ -355,11 +344,10 @@ func TestSnapshotPreview1_EnvironSizesGet(t *testing.T) {
 }
 
 func TestSnapshotPreview1_EnvironSizesGet_Errors(t *testing.T) {
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, []string{"a=b", "b=cd"}, nil)
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionEnvironSizesGet, importEnvironSizesGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionEnvironSizesGet, importEnvironSizesGet, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -404,11 +392,7 @@ func TestSnapshotPreview1_EnvironSizesGet_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_ClockResGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_ClockResGet(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionClockResGet, importClockResGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionClockResGet, importClockResGet, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.ClockResGet", func(t *testing.T) {
@@ -432,11 +416,7 @@ func TestSnapshotPreview1_ClockTimeGet(t *testing.T) {
 		'?', // stopped after encoding
 	}
 
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionClockTimeGet, importClockTimeGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionClockTimeGet, importClockTimeGet, nil)
 	defer mod.Close()
 
 	a.timeNowUnixNano = func() uint64 { return epochNanos }
@@ -470,11 +450,7 @@ func TestSnapshotPreview1_ClockTimeGet(t *testing.T) {
 func TestSnapshotPreview1_ClockTimeGet_Errors(t *testing.T) {
 	epochNanos := uint64(1640995200000000000) // midnight UTC 2022-01-01
 
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionClockTimeGet, importClockTimeGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionClockTimeGet, importClockTimeGet, nil)
 	defer mod.Close()
 
 	a.timeNowUnixNano = func() uint64 { return epochNanos }
@@ -511,11 +487,7 @@ func TestSnapshotPreview1_ClockTimeGet_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdAdvise only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdAdvise(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdAdvise, importFdAdvise, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdAdvise, importFdAdvise, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdAdvise", func(t *testing.T) {
@@ -533,11 +505,7 @@ func TestSnapshotPreview1_FdAdvise(t *testing.T) {
 
 // TestSnapshotPreview1_FdAllocate only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdAllocate(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdAllocate, importFdAllocate, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdAllocate, importFdAllocate, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdAllocate", func(t *testing.T) {
@@ -558,8 +526,6 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 	fdToKeep := uint32(4)  // another arbitrary fd
 
 	setupFD := func() (api.Module, api.Function, *wasiAPI) {
-		ctx := context.Background()
-
 		// fd_close needs to close an open file descriptor. Open two files so that we can tell which is closed.
 		path1, path2 := "a", "b"
 		testFs := fstest.MapFS{path1: {Data: make([]byte, 0)}, path2: {Data: make([]byte, 0)}}
@@ -574,7 +540,7 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		a, mod, fn := instantiateModule(t, ctx, functionFdClose, importFdClose, moduleName, sysCtx)
+		a, mod, fn := instantiateModule(t, functionFdClose, importFdClose, sysCtx)
 		return mod, fn, a
 	}
 
@@ -619,11 +585,7 @@ func TestSnapshotPreview1_FdClose(t *testing.T) {
 
 // TestSnapshotPreview1_FdDatasync only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdDatasync(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdDatasync, importFdDatasync, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdDatasync, importFdDatasync, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdDatasync", func(t *testing.T) {
@@ -647,11 +609,7 @@ func TestSnapshotPreview1_FdFdstatGet(t *testing.T) {
 
 // TestSnapshotPreview1_FdFdstatSetFlags only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFdstatSetFlags(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdFdstatSetFlags, importFdFdstatSetFlags, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdFdstatSetFlags, importFdFdstatSetFlags, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdFdstatSetFlags", func(t *testing.T) {
@@ -669,11 +627,7 @@ func TestSnapshotPreview1_FdFdstatSetFlags(t *testing.T) {
 
 // TestSnapshotPreview1_FdFdstatSetRights only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFdstatSetRights(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdFdstatSetRights, importFdFdstatSetRights, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdFdstatSetRights, importFdFdstatSetRights, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdFdstatSetRights", func(t *testing.T) {
@@ -691,11 +645,7 @@ func TestSnapshotPreview1_FdFdstatSetRights(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatGet(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdFilestatGet, importFdFilestatGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdFilestatGet, importFdFilestatGet, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdFilestatGet", func(t *testing.T) {
@@ -713,11 +663,7 @@ func TestSnapshotPreview1_FdFilestatGet(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatSetSize only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatSetSize(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdFilestatSetSize, importFdFilestatSetSize, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdFilestatSetSize, importFdFilestatSetSize, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdFilestatSetSize", func(t *testing.T) {
@@ -735,11 +681,7 @@ func TestSnapshotPreview1_FdFilestatSetSize(t *testing.T) {
 
 // TestSnapshotPreview1_FdFilestatSetTimes only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdFilestatSetTimes(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdFilestatSetTimes, importFdFilestatSetTimes, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdFilestatSetTimes, importFdFilestatSetTimes, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdFilestatSetTimes", func(t *testing.T) {
@@ -757,11 +699,7 @@ func TestSnapshotPreview1_FdFilestatSetTimes(t *testing.T) {
 
 // TestSnapshotPreview1_FdPread only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdPread(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdPread, importFdPread, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdPread, importFdPread, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdPread", func(t *testing.T) {
@@ -780,12 +718,11 @@ func TestSnapshotPreview1_FdPread(t *testing.T) {
 func TestSnapshotPreview1_FdPrestatGet(t *testing.T) {
 	fd := uint32(3) // arbitrary fd after 0, 1, and 2, that are stdin/out/err
 
-	ctx := context.Background()
 	pathName := "/tmp"
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{fd: {Path: pathName}})
 	require.NoError(t, err)
 
-	a, mod, fn := instantiateModule(t, ctx, functionFdPrestatGet, importFdPrestatGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdPrestatGet, importFdPrestatGet, sysCtx)
 	defer mod.Close()
 
 	resultPrestat := uint32(1) // arbitrary offset
@@ -827,11 +764,10 @@ func TestSnapshotPreview1_FdPrestatGet_Errors(t *testing.T) {
 	fd := uint32(3)           // fd 3 will be opened for the "/tmp" directory after 0, 1, and 2, that are stdin/out/err
 	validAddress := uint32(0) // Arbitrary valid address as arguments to fd_prestat_get. We chose 0 here.
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{fd: {Path: "/tmp"}})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionFdPrestatGet, importFdPrestatGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionFdPrestatGet, importFdPrestatGet, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -870,11 +806,10 @@ func TestSnapshotPreview1_FdPrestatGet_Errors(t *testing.T) {
 func TestSnapshotPreview1_FdPrestatDirName(t *testing.T) {
 	fd := uint32(3) // arbitrary fd after 0, 1, and 2, that are stdin/out/err
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{fd: {Path: "/tmp"}})
 	require.NoError(t, err)
 
-	a, mod, fn := instantiateModule(t, ctx, functionFdPrestatDirName, importFdPrestatDirName, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdPrestatDirName, importFdPrestatDirName, sysCtx)
 	defer mod.Close()
 
 	path := uint32(1)    // arbitrary offset
@@ -912,12 +847,10 @@ func TestSnapshotPreview1_FdPrestatDirName(t *testing.T) {
 
 func TestSnapshotPreview1_FdPrestatDirName_Errors(t *testing.T) {
 	fd := uint32(3) // arbitrary fd after 0, 1, and 2, that are stdin/out/err
-
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{fd: {Path: "/tmp"}})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionFdPrestatDirName, importFdPrestatDirName, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionFdPrestatDirName, importFdPrestatDirName, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -974,11 +907,7 @@ func TestSnapshotPreview1_FdPrestatDirName_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdPwrite only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdPwrite(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdPwrite, importFdPwrite, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdPwrite, importFdPwrite, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdPwrite", func(t *testing.T) {
@@ -995,8 +924,6 @@ func TestSnapshotPreview1_FdPwrite(t *testing.T) {
 }
 
 func TestSnapshotPreview1_FdRead(t *testing.T) {
-	ctx := context.Background()
-
 	fd := uint32(3)   // arbitrary fd after 0, 1, and 2, that are stdin/out/err
 	iovs := uint32(1) // arbitrary offset
 	initialMemory := []byte{
@@ -1023,12 +950,12 @@ func TestSnapshotPreview1_FdRead(t *testing.T) {
 	type fdReadFn func(ctx api.Module, fd, iovs, iovsCount, resultSize uint32) Errno
 	tests := []struct {
 		name   string
-		fdRead func(*wasiAPI, *wasm.ModuleContext, api.Function) fdReadFn
+		fdRead func(*wasiAPI, api.Module, api.Function) fdReadFn
 	}{
-		{"snapshotPreview1.FdRead", func(a *wasiAPI, _ *wasm.ModuleContext, _ api.Function) fdReadFn {
+		{"snapshotPreview1.FdRead", func(a *wasiAPI, _ api.Module, _ api.Function) fdReadFn {
 			return a.FdRead
 		}},
-		{functionFdRead, func(_ *wasiAPI, mod *wasm.ModuleContext, fn api.Function) fdReadFn {
+		{functionFdRead, func(_ *wasiAPI, mod api.Module, fn api.Function) fdReadFn {
 			return func(ctx api.Module, fd, iovs, iovsCount, resultSize uint32) Errno {
 				results, err := fn.Call(mod, uint64(fd), uint64(iovs), uint64(iovsCount), uint64(resultSize))
 				require.NoError(t, err)
@@ -1047,7 +974,7 @@ func TestSnapshotPreview1_FdRead(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			a, mod, fn := instantiateModule(t, ctx, functionFdRead, importFdRead, moduleName, sysCtx)
+			a, mod, fn := instantiateModule(t, functionFdRead, importFdRead, sysCtx)
 			defer mod.Close()
 
 			maskMemory(t, mod, len(expectedMemory))
@@ -1069,13 +996,12 @@ func TestSnapshotPreview1_FdRead_Errors(t *testing.T) {
 	validFD := uint32(3)                                 // arbitrary valid fd after 0, 1, and 2, that are stdin/out/err
 	file, testFS := createFile(t, "test_path", []byte{}) // file with empty contents
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 		validFD: {Path: "test_path", FS: testFS, File: file},
 	})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionFdRead, importFdRead, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionFdRead, importFdRead, sysCtx)
 	defer mod.Close()
 
 	tests := []struct {
@@ -1160,11 +1086,7 @@ func TestSnapshotPreview1_FdRead_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdReaddir only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdReaddir(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdReaddir, importFdReaddir, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdReaddir, importFdReaddir, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdReaddir", func(t *testing.T) {
@@ -1182,11 +1104,7 @@ func TestSnapshotPreview1_FdReaddir(t *testing.T) {
 
 // TestSnapshotPreview1_FdRenumber only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdRenumber(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdRenumber, importFdRenumber, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdRenumber, importFdRenumber, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdRenumber", func(t *testing.T) {
@@ -1207,13 +1125,12 @@ func TestSnapshotPreview1_FdSeek(t *testing.T) {
 	resultNewoffset := uint32(1)                                 // arbitrary offset in `ctx.Memory` for the new offset value
 	file, testFS := createFile(t, "test_path", []byte("wazero")) // arbitrary non-empty contents
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 		fd: {Path: "test_path", FS: testFS, File: file},
 	})
 	require.NoError(t, err)
 
-	a, mod, fn := instantiateModule(t, ctx, functionFdSeek, importFdSeek, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdSeek, importFdSeek, sysCtx)
 	defer mod.Close()
 
 	// TestSnapshotPreview1_FdSeek uses a matrix because setting up test files is complicated and has to be clean each time.
@@ -1313,13 +1230,13 @@ func TestSnapshotPreview1_FdSeek(t *testing.T) {
 func TestSnapshotPreview1_FdSeek_Errors(t *testing.T) {
 	validFD := uint32(3)                                         // arbitrary valid fd after 0, 1, and 2, that are stdin/out/err
 	file, testFS := createFile(t, "test_path", []byte("wazero")) // arbitrary valid file with non-empty contents
-	ctx := context.Background()
+
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 		validFD: {Path: "test_path", FS: testFS, File: file},
 	})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionFdSeek, importFdSeek, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionFdSeek, importFdSeek, sysCtx)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -1362,11 +1279,7 @@ func TestSnapshotPreview1_FdSeek_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_FdSync only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdSync(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdSync, importFdSync, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdSync, importFdSync, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdSync", func(t *testing.T) {
@@ -1384,11 +1297,7 @@ func TestSnapshotPreview1_FdSync(t *testing.T) {
 
 // TestSnapshotPreview1_FdTell only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_FdTell(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionFdTell, importFdTell, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionFdTell, importFdTell, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.FdTell", func(t *testing.T) {
@@ -1405,8 +1314,6 @@ func TestSnapshotPreview1_FdTell(t *testing.T) {
 }
 
 func TestSnapshotPreview1_FdWrite(t *testing.T) {
-	ctx := context.Background()
-
 	fd := uint32(3)   // arbitrary fd after 0, 1, and 2, that are stdin/out/err
 	iovs := uint32(1) // arbitrary offset
 	initialMemory := []byte{
@@ -1433,12 +1340,12 @@ func TestSnapshotPreview1_FdWrite(t *testing.T) {
 	type fdWriteFn func(ctx api.Module, fd, iovs, iovsCount, resultSize uint32) Errno
 	tests := []struct {
 		name    string
-		fdWrite func(*wasiAPI, *wasm.ModuleContext, api.Function) fdWriteFn
+		fdWrite func(*wasiAPI, api.Module, api.Function) fdWriteFn
 	}{
-		{"snapshotPreview1.FdWrite", func(a *wasiAPI, _ *wasm.ModuleContext, _ api.Function) fdWriteFn {
+		{"snapshotPreview1.FdWrite", func(a *wasiAPI, _ api.Module, _ api.Function) fdWriteFn {
 			return a.FdWrite
 		}},
-		{functionFdWrite, func(_ *wasiAPI, mod *wasm.ModuleContext, fn api.Function) fdWriteFn {
+		{functionFdWrite, func(_ *wasiAPI, mod api.Module, fn api.Function) fdWriteFn {
 			return func(ctx api.Module, fd, iovs, iovsCount, resultSize uint32) Errno {
 				results, err := fn.Call(mod, uint64(fd), uint64(iovs), uint64(iovsCount), uint64(resultSize))
 				require.NoError(t, err)
@@ -1460,7 +1367,7 @@ func TestSnapshotPreview1_FdWrite(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			a, mod, fn := instantiateModule(t, ctx, functionFdWrite, importFdWrite, moduleName, sysCtx)
+			a, mod, fn := instantiateModule(t, functionFdWrite, importFdWrite, sysCtx)
 			defer mod.Close()
 
 			maskMemory(t, mod, len(expectedMemory))
@@ -1490,13 +1397,12 @@ func TestSnapshotPreview1_FdWrite_Errors(t *testing.T) {
 	pathName := "test_path"
 	file, testFS := createWriteableFile(t, tmpDir, pathName, []byte{})
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 		validFD: {Path: pathName, FS: testFS, File: file},
 	})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionFdWrite, importFdWrite, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionFdWrite, importFdWrite, sysCtx)
 	defer mod.Close()
 
 	// Setup valid test memory
@@ -1564,11 +1470,7 @@ func TestSnapshotPreview1_FdWrite_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_PathCreateDirectory only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathCreateDirectory(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathCreateDirectory, importPathCreateDirectory, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathCreateDirectory, importPathCreateDirectory, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathCreateDirectory", func(t *testing.T) {
@@ -1586,11 +1488,7 @@ func TestSnapshotPreview1_PathCreateDirectory(t *testing.T) {
 
 // TestSnapshotPreview1_PathFilestatGet only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathFilestatGet(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathFilestatGet, importPathFilestatGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathFilestatGet, importPathFilestatGet, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathFilestatGet", func(t *testing.T) {
@@ -1608,11 +1506,7 @@ func TestSnapshotPreview1_PathFilestatGet(t *testing.T) {
 
 // TestSnapshotPreview1_PathFilestatSetTimes only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathFilestatSetTimes(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathFilestatSetTimes, importPathFilestatSetTimes, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathFilestatSetTimes, importPathFilestatSetTimes, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathFilestatSetTimes", func(t *testing.T) {
@@ -1630,11 +1524,7 @@ func TestSnapshotPreview1_PathFilestatSetTimes(t *testing.T) {
 
 // TestSnapshotPreview1_PathLink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathLink(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathLink, importPathLink, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathLink, importPathLink, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathLink", func(t *testing.T) {
@@ -1671,25 +1561,23 @@ func TestSnapshotPreview1_PathOpen(t *testing.T) {
 		'?',
 	)
 
-	ctx := context.Background()
-
 	// rights are ignored per https://github.com/WebAssembly/WASI/issues/469#issuecomment-1045251844
 	fsRightsBase, fsRightsInheriting := uint64(1), uint64(2)
 
-	setup := func() (*wasiAPI, *wasm.ModuleContext, api.Function) {
+	setup := func() (*wasiAPI, api.Module, api.Function) {
 		testFS := fstest.MapFS{pathName: &fstest.MapFile{Mode: os.ModeDir}}
 		sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 			workdirFD: {Path: ".", FS: testFS},
 		})
 		require.NoError(t, err)
-		a, mod, fn := instantiateModule(t, ctx, functionPathOpen, importPathOpen, moduleName, sysCtx)
+		a, mod, fn := instantiateModule(t, functionPathOpen, importPathOpen, sysCtx)
 		maskMemory(t, mod, len(expectedMemory))
 		ok := mod.Memory().Write(0, initialMemory)
 		require.True(t, ok)
 		return a, mod, fn
 	}
 
-	verify := func(errno Errno, mod *wasm.ModuleContext) {
+	verify := func(errno Errno, mod api.Module) {
 		require.Zero(t, errno, ErrnoName(errno))
 
 		actual, ok := mod.Memory().Read(0, uint32(len(expectedMemory)))
@@ -1722,13 +1610,12 @@ func TestSnapshotPreview1_PathOpen_Errors(t *testing.T) {
 	pathName := "wazero"
 	testFS := fstest.MapFS{pathName: &fstest.MapFile{Mode: os.ModeDir}}
 
-	ctx := context.Background()
 	sysCtx, err := newSysContext(nil, nil, map[uint32]*wasm.FileEntry{
 		validFD: {Path: ".", FS: testFS},
 	})
 	require.NoError(t, err)
 
-	a, mod, _ := instantiateModule(t, ctx, functionPathOpen, importPathOpen, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionPathOpen, importPathOpen, sysCtx)
 	defer mod.Close()
 
 	validPath := uint32(0)    // arbitrary offset
@@ -1787,11 +1674,7 @@ func TestSnapshotPreview1_PathOpen_Errors(t *testing.T) {
 
 // TestSnapshotPreview1_PathReadlink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathReadlink(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathReadlink, importPathReadlink, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathReadlink, importPathReadlink, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathLink", func(t *testing.T) {
@@ -1809,11 +1692,7 @@ func TestSnapshotPreview1_PathReadlink(t *testing.T) {
 
 // TestSnapshotPreview1_PathRemoveDirectory only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathRemoveDirectory(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathRemoveDirectory, importPathRemoveDirectory, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathRemoveDirectory, importPathRemoveDirectory, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathRemoveDirectory", func(t *testing.T) {
@@ -1831,11 +1710,7 @@ func TestSnapshotPreview1_PathRemoveDirectory(t *testing.T) {
 
 // TestSnapshotPreview1_PathRename only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathRename(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathRename, importPathRename, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathRename, importPathRename, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathRename", func(t *testing.T) {
@@ -1853,11 +1728,7 @@ func TestSnapshotPreview1_PathRename(t *testing.T) {
 
 // TestSnapshotPreview1_PathSymlink only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathSymlink(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathSymlink, importPathSymlink, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathSymlink, importPathSymlink, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathSymlink", func(t *testing.T) {
@@ -1875,11 +1746,7 @@ func TestSnapshotPreview1_PathSymlink(t *testing.T) {
 
 // TestSnapshotPreview1_PathUnlinkFile only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PathUnlinkFile(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPathUnlinkFile, importPathUnlinkFile, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPathUnlinkFile, importPathUnlinkFile, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PathUnlinkFile", func(t *testing.T) {
@@ -1897,11 +1764,7 @@ func TestSnapshotPreview1_PathUnlinkFile(t *testing.T) {
 
 // TestSnapshotPreview1_PollOneoff only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_PollOneoff(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionPollOneoff, importPollOneoff, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionPollOneoff, importPollOneoff, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.PollOneoff", func(t *testing.T) {
@@ -1918,10 +1781,6 @@ func TestSnapshotPreview1_PollOneoff(t *testing.T) {
 }
 
 func TestSnapshotPreview1_ProcExit(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
 	tests := []struct {
 		name     string
 		exitCode uint32
@@ -1930,7 +1789,6 @@ func TestSnapshotPreview1_ProcExit(t *testing.T) {
 			name:     "success (exitcode 0)",
 			exitCode: 0,
 		},
-
 		{
 			name:     "arbitrary non-zero exitcode",
 			exitCode: 42,
@@ -1943,11 +1801,11 @@ func TestSnapshotPreview1_ProcExit(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			// Note: Unlike most tests, this uses fn, not the 'a' result parameter. This is because currently, this function
 			// body panics, and we expect Call to unwrap the panic.
-			_, mod, fn := instantiateModule(t, ctx, functionProcExit, importProcExit, moduleName, sysCtx)
+			_, mod, fn := instantiateModule(t, functionProcExit, importProcExit, nil)
 			defer mod.Close()
 
 			// When ProcExit is called, store.Callfunction returns immediately, returning the exit code as the error.
-			_, err = fn.Call(mod, uint64(tc.exitCode))
+			_, err := fn.Call(nil, uint64(tc.exitCode))
 			require.Equal(t, tc.exitCode, err.(*sys.ExitError).ExitCode())
 		})
 	}
@@ -1955,11 +1813,7 @@ func TestSnapshotPreview1_ProcExit(t *testing.T) {
 
 // TestSnapshotPreview1_ProcRaise only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_ProcRaise(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionProcRaise, importProcRaise, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionProcRaise, importProcRaise, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.ProcRaise", func(t *testing.T) {
@@ -1977,11 +1831,7 @@ func TestSnapshotPreview1_ProcRaise(t *testing.T) {
 
 // TestSnapshotPreview1_SchedYield only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SchedYield(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionSchedYield, importSchedYield, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionSchedYield, importSchedYield, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.SchedYield", func(t *testing.T) {
@@ -2007,11 +1857,8 @@ func TestSnapshotPreview1_RandomGet(t *testing.T) {
 	length := uint32(5) // arbitrary length,
 	offset := uint32(1) // offset,
 	seed := int64(42)   // and seed value
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
 
-	a, mod, fn := instantiateModule(t, ctx, functionRandomGet, importRandomGet, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionRandomGet, importRandomGet, nil)
 	defer mod.Close()
 
 	a.randSource = func(p []byte) error {
@@ -2049,13 +1896,9 @@ func TestSnapshotPreview1_RandomGet(t *testing.T) {
 }
 
 func TestSnapshotPreview1_RandomGet_Errors(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
 	validAddress := uint32(0) // arbitrary valid address
 
-	a, mod, _ := instantiateModule(t, ctx, functionRandomGet, importRandomGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionRandomGet, importRandomGet, nil)
 	defer mod.Close()
 
 	memorySize := mod.Memory().Size()
@@ -2089,11 +1932,7 @@ func TestSnapshotPreview1_RandomGet_Errors(t *testing.T) {
 }
 
 func TestSnapshotPreview1_RandomGet_SourceError(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, _ := instantiateModule(t, ctx, functionRandomGet, importRandomGet, moduleName, sysCtx)
+	a, mod, _ := instantiateModule(t, functionRandomGet, importRandomGet, nil)
 	defer mod.Close()
 
 	a.randSource = func(p []byte) error {
@@ -2106,11 +1945,7 @@ func TestSnapshotPreview1_RandomGet_SourceError(t *testing.T) {
 
 // TestSnapshotPreview1_SockRecv only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockRecv(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionSockRecv, importSockRecv, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionSockRecv, importSockRecv, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.SockRecv", func(t *testing.T) {
@@ -2128,11 +1963,7 @@ func TestSnapshotPreview1_SockRecv(t *testing.T) {
 
 // TestSnapshotPreview1_SockSend only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockSend(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionSockSend, importSockSend, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionSockSend, importSockSend, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.SockSend", func(t *testing.T) {
@@ -2150,11 +1981,7 @@ func TestSnapshotPreview1_SockSend(t *testing.T) {
 
 // TestSnapshotPreview1_SockShutdown only tests it is stubbed for GrainLang per #271
 func TestSnapshotPreview1_SockShutdown(t *testing.T) {
-	ctx := context.Background()
-	sysCtx, err := newSysContext(nil, nil, nil)
-	require.NoError(t, err)
-
-	a, mod, fn := instantiateModule(t, ctx, functionSockShutdown, importSockShutdown, moduleName, sysCtx)
+	a, mod, fn := instantiateModule(t, functionSockShutdown, importSockShutdown, nil)
 	defer mod.Close()
 
 	t.Run("snapshotPreview1.SockShutdown", func(t *testing.T) {
@@ -2179,34 +2006,30 @@ func maskMemory(t *testing.T, mod api.Module, size int) {
 	}
 }
 
-func instantiateModule(t *testing.T, ctx context.Context, wasifunction, wasiimport, moduleName string, sysCtx *wasm.SysContext) (*wasiAPI, *wasm.ModuleContext, api.Function) {
-	enabledFeatures := wasm.Features20191205
-	store := wasm.NewStore(interpreter.NewEngine(), enabledFeatures)
+func instantiateModule(t *testing.T, wasifunction, wasiimport string, sysCtx *wasm.SysContext) (*wasiAPI, api.Module, api.Function) {
+	r := wazero.NewRuntimeWithConfig(wazero.NewRuntimeConfigInterpreter())
 
 	// The package `wazero` has a simpler interface for adding host modules, but we can't use that as it would create an
 	// import cycle. Instead, we export wasm.NewHostModule and use it here.
 	a, fns := snapshotPreview1Functions()
-	m, err := wasm.NewHostModule("wasi_snapshot_preview1", fns)
+	_, err := r.NewModuleBuilder("wasi_snapshot_preview1").ExportFunctions(fns).Instantiate()
 	require.NoError(t, err)
 
-	// Double-check what we created passes same validity as module-defined modules.
-	require.NoError(t, m.Validate(enabledFeatures))
-
-	w, err := store.Instantiate(ctx, m, m.NameSection.ModuleName, nil)
-	require.NoError(t, err)
-	defer w.Close()
-
-	m, err = text.DecodeModule([]byte(fmt.Sprintf(`(module
+	m, err := r.CompileModule([]byte(fmt.Sprintf(`(module
   %[2]s
-  (memory 1)  ;; just an arbitrary size big enough for tests
+  (memory 1 1)  ;; just an arbitrary size big enough for tests
   (export "memory" (memory 0))
   (export "%[1]s" (func $wasi.%[1]s))
-)`, wasifunction, wasiimport)), enabledFeatures, wasm.MemoryMaxPages)
+)`, wasifunction, wasiimport)))
 	require.NoError(t, err)
 
-	mod, err := store.Instantiate(ctx, m, moduleName, sysCtx)
+	mod, err := r.InstantiateModuleWithConfig(m, wazero.NewModuleConfig().WithName(t.Name()))
 	require.NoError(t, err)
 
+	if sysCtx != nil {
+		mod.(*wasm.ModuleContext).Sys = sysCtx
+	}
+
 	fn := mod.ExportedFunction(wasifunction)
 	require.NotNil(t, fn)
 	return a, mod, fn