Propagates context to all api interface methods that aren't constant (#…

…502)

This prepares for exposing operations like Memory.Grow while keeping the
ability to trace what did that, by adding a `context.Context` initial
parameter. This adds this to all API methods that mutate or return
mutated data.

Before, we made a change to trace functions and general lifecycle
commands, but we missed this part. Ex. We track functions, but can't
track what closed the module, changed memory or a mutable constant.
Changing to do this now is not only more consistent, but helps us
optimize at least the interpreter to help users identify otherwise
opaque code that can cause harm. This is critical before we add more
functions that can cause harm, such as Memory.Grow.

Signed-off-by: Adrian Cole <adrian@tetrate.io>

README.md

@@ -26,7 +26,7 @@ func main() {
 
 	// Instantiate the module and return its exported functions
 	module, _ := wazero.NewRuntime().InstantiateModuleFromCode(ctx, source)
-	defer module.Close()
+	defer module.Close(ctx)
 
 	// Discover 7! is 5040
 	fmt.Println(module.ExportedFunction("fac").Call(ctx, 7))
@@ -63,7 +63,7 @@ env, err := r.NewModuleBuilder("env").
 if err != nil {
 	log.Fatal(err)
 }
-defer env.Close()
+defer env.Close(ctx)
 ```
 
 While not a standards body like W3C, there is another dominant community in the
@@ -77,11 +77,11 @@ For example, here's how you can allow WebAssembly modules to read
 "/work/home/a.txt" as "/a.txt" or "./a.txt":
 ```go
 wm, err := wasi.InstantiateSnapshotPreview1(ctx, r)
-defer wm.Close()
+defer wm.Close(ctx)
 
 config := wazero.ModuleConfig().WithFS(os.DirFS("/work/home"))
 module, err := r.InstantiateModule(ctx, binary, config)
-defer module.Close()
+defer module.Close(ctx)
 ...
 ```
 
@@ -302,7 +302,7 @@ top-level project. That said, Takeshi's original motivation is as relevant
 today as when he started the project, and worthwhile reading:
 
 If you want to provide Wasm host environments in your Go programs, currently
-there's no other choice than using CGO andleveraging the state-of-the-art
+there's no other choice than using CGO leveraging the state-of-the-art
 runtimes written in C++/Rust (e.g. V8, Wasmtime, Wasmer, WAVM, etc.), and
 there's no pure Go Wasm runtime out there. (There's only one exception named
 [wagon](https://github.com/go-interpreter/wagon), but it was archived with the
@@ -313,7 +313,7 @@ plugin systems in your Go project and want these plugin systems to be
 safe/fast/flexible, and enable users to write plugins in their favorite
 languages. That's where Wasm comes into play. You write your own Wasm host
 environments and embed Wasm runtime in your projects, and now users are able to
-write plugins in their own favorite lanugages (AssembyScript, C, C++, Rust,
+write plugins in their own favorite languages (AssemblyScript, C, C++, Rust,
 Zig, etc.). As a specific example, you maybe write proxy severs in Go and want
 to allow users to extend the proxy via [Proxy-Wasm ABI](https://github.com/proxy-wasm/spec).
 Maybe you are writing server-side rendering applications via Wasm, or

api/wasm.go

@@ -70,8 +70,8 @@ type Module interface {
 	Name() string
 
 	// Close is a convenience that invokes CloseWithExitCode with zero.
-	Close() error
-	// ^^ not io.Closer as the rationale (static analysis of leaks) is invalid when there are multiple close methods.
+	// Note: When the context is nil, it defaults to context.Background.
+	Close(context.Context) error
 
 	// CloseWithExitCode releases resources allocated for this Module. Use a non-zero exitCode parameter to indicate a
 	// failure to ExportedFunction callers.
@@ -82,7 +82,8 @@ type Module interface {
 	//
 	// Calling this inside a host function is safe, and may cause ExportedFunction callers to receive a sys.ExitError
 	// with the exitCode.
-	CloseWithExitCode(exitCode uint32) error
+	// Note: When the context is nil, it defaults to context.Background.
+	CloseWithExitCode(ctx context.Context, exitCode uint32) error
 
 	// Memory returns a memory defined in this module or nil if there are none wasn't.
 	Memory() Memory
@@ -121,7 +122,7 @@ type Function interface {
 	// encoded according to ResultTypes. An error is returned for any failure looking up or invoking the function
 	// including signature mismatch.
 	//
-	// Note: when `ctx` is nil, it defaults to context.Background.
+	// Note: When the context is nil, it defaults to context.Background.
 	// Note: If Module.Close or Module.CloseWithExitCode were invoked during this call, the error returned may be a
 	// sys.ExitError. Interpreting this is specific to the module. For example, some "main" functions always call a
 	// function that exits.
@@ -153,7 +154,9 @@ type Global interface {
 
 	// Get returns the last known value of this global.
 	// See Type for how to encode this value from a Go type.
-	Get() uint64
+	//
+	// Note: When the context is nil, it defaults to context.Background.
+	Get(context.Context) uint64
 }
 
 // MutableGlobal is a Global whose value can be updated at runtime (variable).
@@ -162,11 +165,14 @@ type MutableGlobal interface {
 
 	// Set updates the value of this global.
 	// See Global.Type for how to decode this value to a Go type.
-	Set(v uint64)
+	//
+	// Note: When the context is nil, it defaults to context.Background.
+	Set(ctx context.Context, v uint64)
 }
 
 // Memory allows restricted access to a module's memory. Notably, this does not allow growing.
 //
+// Note: All functions accept a context.Context, which when nil, default to context.Background.
 // Note: This is an interface for decoupling, not third-party implementations. All implementations are in wazero.
 // Note: This includes all value types available in WebAssembly 1.0 (20191205) and all are encoded little-endian.
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#storage%E2%91%A0
@@ -177,59 +183,67 @@ type Memory interface {
 	// memory has min 0 and max 2 pages, this returns zero.
 	//
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#-hrefsyntax-instr-memorymathsfmemorysize%E2%91%A0
-	Size() uint32
+	Size(context.Context) uint32
 
 	// IndexByte returns the index of the first instance of c in the underlying buffer at the offset or returns false if
 	// not found or out of range.
-	IndexByte(offset uint32, c byte) (uint32, bool)
+	IndexByte(ctx context.Context, offset uint32, c byte) (uint32, bool)
 
 	// ReadByte reads a single byte from the underlying buffer at the offset or returns false if out of range.
-	ReadByte(offset uint32) (byte, bool)
+	ReadByte(ctx context.Context, offset uint32) (byte, bool)
+
+	// ReadUint16Le reads a uint16 in little-endian encoding from the underlying buffer at the offset in or returns
+	// false if out of range.
+	ReadUint16Le(ctx context.Context, offset uint32) (uint16, bool)
 
 	// ReadUint32Le reads a uint32 in little-endian encoding from the underlying buffer at the offset in or returns
 	// false if out of range.
-	ReadUint32Le(offset uint32) (uint32, bool)
+	ReadUint32Le(ctx context.Context, offset uint32) (uint32, bool)
 
 	// ReadFloat32Le reads a float32 from 32 IEEE 754 little-endian encoded bits in the underlying buffer at the offset
 	// or returns false if out of range.
 	// See math.Float32bits
-	ReadFloat32Le(offset uint32) (float32, bool)
+	ReadFloat32Le(ctx context.Context, offset uint32) (float32, bool)
 
 	// ReadUint64Le reads a uint64 in little-endian encoding from the underlying buffer at the offset or returns false
 	// if out of range.
-	ReadUint64Le(offset uint32) (uint64, bool)
+	ReadUint64Le(ctx context.Context, offset uint32) (uint64, bool)
 
 	// ReadFloat64Le reads a float64 from 64 IEEE 754 little-endian encoded bits in the underlying buffer at the offset
 	// or returns false if out of range.
 	// See math.Float64bits
-	ReadFloat64Le(offset uint32) (float64, bool)
+	ReadFloat64Le(ctx context.Context, offset uint32) (float64, bool)
 
 	// Read reads byteCount bytes from the underlying buffer at the offset or returns false if out of range.
-	Read(offset, byteCount uint32) ([]byte, bool)
+	Read(ctx context.Context, offset, byteCount uint32) ([]byte, bool)
 
 	// WriteByte writes a single byte to the underlying buffer at the offset in or returns false if out of range.
-	WriteByte(offset uint32, v byte) bool
+	WriteByte(ctx context.Context, offset uint32, v byte) bool
+
+	// WriteUint16Le writes the value in little-endian encoding to the underlying buffer at the offset in or returns
+	// false if out of range.
+	WriteUint16Le(ctx context.Context, offset uint32, v uint16) bool
 
 	// WriteUint32Le writes the value in little-endian encoding to the underlying buffer at the offset in or returns
 	// false if out of range.
-	WriteUint32Le(offset, v uint32) bool
+	WriteUint32Le(ctx context.Context, offset, v uint32) bool
 
 	// WriteFloat32Le writes the value in 32 IEEE 754 little-endian encoded bits to the underlying buffer at the offset
 	// or returns false if out of range.
 	// See math.Float32bits
-	WriteFloat32Le(offset uint32, v float32) bool
+	WriteFloat32Le(ctx context.Context, offset uint32, v float32) bool
 
 	// WriteUint64Le writes the value in little-endian encoding to the underlying buffer at the offset in or returns
 	// false if out of range.
-	WriteUint64Le(offset uint32, v uint64) bool
+	WriteUint64Le(ctx context.Context, offset uint32, v uint64) bool
 
 	// WriteFloat64Le writes the value in 64 IEEE 754 little-endian encoded bits to the underlying buffer at the offset
 	// or returns false if out of range.
 	// See math.Float64bits
-	WriteFloat64Le(offset uint32, v float64) bool
+	WriteFloat64Le(ctx context.Context, offset uint32, v float64) bool
 
 	// Write writes the slice to the underlying buffer at the offset or returns false if out of range.
-	Write(offset uint32, v []byte) bool
+	Write(ctx context.Context, offset uint32, v []byte) bool
 }
 
 // EncodeI32 encodes the input as a ValueTypeI32.

builder.go

@@ -25,10 +25,10 @@ import (
 //	env, _ := r.NewModuleBuilder("env").ExportFunction("get_random_string", getRandomString).Build(ctx)
 //
 //	env1, _ := r.InstantiateModuleWithConfig(ctx, env, NewModuleConfig().WithName("env.1"))
-//	defer env1.Close()
+//	defer env1.Close(ctx)
 //
 //	env2, _ := r.InstantiateModuleWithConfig(ctx, env, NewModuleConfig().WithName("env.2"))
-//	defer env2.Close()
+//	defer env2.Close(ctx)
 //
 // Note: Builder methods do not return errors, to allow chaining. Any validation errors are deferred until Build.
 // Note: Insertion order is not retained. Anything defined by this builder is sorted lexicographically on Build.
@@ -53,17 +53,17 @@ type ModuleBuilder interface {
 	//
 	// Ex. This uses a Go Context:
 	//
-	//	addInts := func(m context.Context, x uint32, uint32) uint32 {
+	//	addInts := func(ctx context.Context, x uint32, uint32) uint32 {
 	//		// add a little extra if we put some in the context!
-	//		return x + y + m.Value(extraKey).(uint32)
+	//		return x + y + ctx.Value(extraKey).(uint32)
 	//	}
 	//
 	// Ex. This uses an api.Module to reads the parameters from memory. This is important because there are only numeric
 	// types in Wasm. The only way to share other data is via writing memory and sharing offsets.
 	//
-	//	addInts := func(m api.Module, offset uint32) uint32 {
-	//		x, _ := m.Memory().ReadUint32Le(offset)
-	//		y, _ := m.Memory().ReadUint32Le(offset + 4) // 32 bits == 4 bytes!
+	//	addInts := func(ctx context.Context, m api.Module, offset uint32) uint32 {
+	//		x, _ := m.Memory().ReadUint32Le(ctx, offset)
+	//		y, _ := m.Memory().ReadUint32Le(ctx, offset + 4) // 32 bits == 4 bytes!
 	//		return x + y
 	//	}
 	//
@@ -151,12 +151,12 @@ type ModuleBuilder interface {
 	ExportGlobalF64(name string, v float64) ModuleBuilder
 
 	// Build returns a module to instantiate, or returns an error if any of the configuration is invalid.
-	Build(ctx context.Context) (*CompiledCode, error)
+	Build(context.Context) (*CompiledCode, error)
 
 	// Instantiate is a convenience that calls Build, then Runtime.InstantiateModule
 	//
 	// Note: Fields in the builder are copied during instantiation: Later changes do not affect the instantiated result.
-	Instantiate(ctx context.Context) (api.Module, error)
+	Instantiate(context.Context) (api.Module, error)
 }
 
 // moduleBuilder implements ModuleBuilder
@@ -274,8 +274,8 @@ func (b *moduleBuilder) Instantiate(ctx context.Context) (api.Module, error) {
 		if err = b.r.store.Engine.CompileModule(ctx, module.module); err != nil {
 			return nil, err
 		}
-		// *wasm.ModuleInstance cannot be tracked, so we release the cache inside of this function.
-		defer module.Close()
+		// *wasm.ModuleInstance cannot be tracked, so we release the cache inside this function.
+		defer module.Close(ctx)
 		return b.r.InstantiateModuleWithConfig(ctx, module, NewModuleConfig().WithName(b.moduleName))
 	}
 }

config.go

@@ -1,6 +1,7 @@
 package wazero
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -151,13 +152,12 @@ type CompiledCode struct {
 	compiledEngine wasm.Engine
 }
 
-// compile-time check to ensure CompiledCode implements io.Closer (consistent with api.Module)
-var _ io.Closer = &CompiledCode{}
-
 // Close releases all the allocated resources for this CompiledCode.
 //
 // Note: It is safe to call Close while having outstanding calls from Modules instantiated from this *CompiledCode.
-func (c *CompiledCode) Close() error {
+func (c *CompiledCode) Close(_ context.Context) error {
+	// Note: If you use the context.Context param, don't forget to coerce nil to context.Background()!
+
 	c.compiledEngine.DeleteCompiledModule(c.module)
 	// It is possible the underlying may need to return an error later, but in any case this matches api.Module.Close.
 	return nil

config_test.go

@@ -1,6 +1,7 @@
 package wazero
 
 import (
+	"context"
 	"io"
 	"math"
 	"testing"
@@ -777,3 +778,27 @@ func requireSysContext(t *testing.T, max uint32, args, environ []string, stdin i
 	require.NoError(t, err)
 	return sys
 }
+
+func TestCompiledCode_Close(t *testing.T) {
+	for _, ctx := range []context.Context{nil, testCtx} { // Ensure it doesn't crash on nil!
+		e := &mockEngine{name: "1", cachedModules: map[*wasm.Module]struct{}{}}
+
+		var cs []*CompiledCode
+		for i := 0; i < 10; i++ {
+			m := &wasm.Module{}
+			err := e.CompileModule(ctx, m)
+			require.NoError(t, err)
+			cs = append(cs, &CompiledCode{module: m, compiledEngine: e})
+		}
+
+		// Before Close.
+		require.Equal(t, 10, len(e.cachedModules))
+
+		for _, c := range cs {
+			require.NoError(t, c.Close(ctx))
+		}
+
+		// After Close.
+		require.Zero(t, len(e.cachedModules))
+	}
+}

example_test.go

@@ -29,7 +29,7 @@ func Example() {
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer mod.Close()
+	defer mod.Close(ctx)
 
 	// Get a function that can be reused until its module is closed:
 	add := mod.ExportedFunction("add")

examples/allocation/rust/greet.go

@@ -34,15 +34,15 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer env.Close()
+	defer env.Close(ctx)
 
 	// Instantiate a WebAssembly module that imports the "log" function defined
 	// in "env" and exports "memory" and functions we'll use in this example.
 	mod, err := r.InstantiateModuleFromCode(ctx, greetWasm)
 	if err != nil {
 		log.Fatal(err)
 	}
-	defer mod.Close()
+	defer mod.Close(ctx)
 
 	// Get references to WebAssembly functions we'll use in this example.
 	greet := mod.ExportedFunction("greet")
@@ -67,9 +67,9 @@ func main() {
 	defer deallocate.Call(ctx, namePtr, nameSize)
 
 	// The pointer is a linear memory offset, which is where we write the name.
-	if !mod.Memory().Write(uint32(namePtr), []byte(name)) {
+	if !mod.Memory().Write(ctx, uint32(namePtr), []byte(name)) {
 		log.Fatalf("Memory.Write(%d, %d) out of range of memory size %d",
-			namePtr, nameSize, mod.Memory().Size())
+			namePtr, nameSize, mod.Memory().Size(ctx))
 	}
 
 	// Now, we can call "greet", which reads the string we wrote to memory!
@@ -91,16 +91,16 @@ func main() {
 	defer deallocate.Call(ctx, uint64(greetingPtr), uint64(greetingSize))
 
 	// The pointer is a linear memory offset, which is where we write the name.
-	if bytes, ok := mod.Memory().Read(greetingPtr, greetingSize); !ok {
+	if bytes, ok := mod.Memory().Read(ctx, greetingPtr, greetingSize); !ok {
 		log.Fatalf("Memory.Read(%d, %d) out of range of memory size %d",
-			greetingPtr, greetingSize, mod.Memory().Size())
+			greetingPtr, greetingSize, mod.Memory().Size(ctx))
 	} else {
 		fmt.Println("go >>", string(bytes))
 	}
 }
 
-func logString(m api.Module, offset, byteCount uint32) {
-	buf, ok := m.Memory().Read(offset, byteCount)
+func logString(ctx context.Context, m api.Module, offset, byteCount uint32) {
+	buf, ok := m.Memory().Read(ctx, offset, byteCount)
 	if !ok {
 		log.Fatalf("Memory.Read(%d, %d) out of range", offset, byteCount)
 	}