experimental

This shows how to implement a custom cleanup task on close.

package main

import (
	"context"

	"github.com/tetratelabs/wazero/experimental"
)

var ctx context.Context

// This shows how to implement a custom cleanup task on close.
func main() {
	closeCh := make(chan struct{})
	ctx = experimental.WithCloseNotifier(
		ctx,
		experimental.CloseNotifyFunc(func(context.Context, uint32) { close(closeCh) }),
	)

	// ... create module, do some work. Sometime later in another goroutine:

	select {
	case <-closeCh:
		// do some cleanup
	default:
		// do some more work with the module
	}
}

Output:

This shows how to make a listener that counts go function calls.

package main

import (
	"context"
	"fmt"
	"log"
	"sort"

	_ "embed"
	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/experimental"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
)

// listenerWasm was generated by the following:
//
//	cd logging/testdata; wat2wasm --debug-names listener.wat
//
//go:embed logging/testdata/listener.wasm
var listenerWasm []byte

// uniqGoFuncs implements both FunctionListenerFactory and FunctionListener
type uniqGoFuncs map[string]struct{}

// callees returns the go functions called.
func (u uniqGoFuncs) callees() []string {
	ret := make([]string, 0, len(u))
	for k := range u {
		ret = append(ret, k)
	}

	sort.Strings(ret)
	return ret
}

// NewFunctionListener implements FunctionListenerFactory.NewFunctionListener
func (u uniqGoFuncs) NewFunctionListener(def api.FunctionDefinition) experimental.FunctionListener {
	if def.GoFunction() == nil {
		return nil
	}
	return u
}

// Before implements FunctionListener.Before
func (u uniqGoFuncs) Before(ctx context.Context, _ api.Module, def api.FunctionDefinition, _ []uint64, _ experimental.StackIterator) {
	u[def.DebugName()] = struct{}{}
}

// After implements FunctionListener.After
func (u uniqGoFuncs) After(context.Context, api.Module, api.FunctionDefinition, []uint64) {}

// Abort implements FunctionListener.Abort
func (u uniqGoFuncs) Abort(context.Context, api.Module, api.FunctionDefinition, error) {}

func main() {
	u := uniqGoFuncs{}

	// Set context to one that has an experimental listener
	ctx := experimental.WithFunctionListenerFactory(context.Background(), u)

	r := wazero.NewRuntime(ctx)
	defer r.Close(ctx) // This closes everything this Runtime created.

	wasi_snapshot_preview1.MustInstantiate(ctx, r)

	mod, err := r.Instantiate(ctx, listenerWasm)
	if err != nil {
		log.Panicln(err)
	}

	for i := 0; i < 5; i++ {
		if _, err = mod.ExportedFunction("rand").Call(ctx, 4); err != nil {
			log.Panicln(err)
		}
	}

	// A Go function was called multiple times, but we should only see it once.
	for _, f := range u.callees() {
		fmt.Println(f)
	}

}

Output:

wasi_snapshot_preview1.fd_write
wasi_snapshot_preview1.random_get

package main

import (
	"context"
	_ "embed"
	"fmt"
	"log"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/experimental"
)

// snapshotWasm was generated by the following:
//
//	cd testdata; wat2wasm snapshot.wat
//
//go:embed testdata/snapshot.wasm
var snapshotWasm []byte

type snapshotsKey struct{}

func main() {
	ctx := context.Background()

	rt := wazero.NewRuntime(ctx)
	defer rt.Close(ctx) // This closes everything this Runtime created.

	// Enable experimental snapshotting functionality by setting it to context. We use this
	// context when invoking functions, indicating to wazero to enable it.
	ctx = experimental.WithSnapshotter(ctx)

	// Also place a mutable holder of snapshots to be referenced during restore.
	var snapshots []experimental.Snapshot
	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)

	// Register host functions using snapshot and restore. Generally snapshot is saved
	// into a mutable location in context to be referenced during restore.
	_, err := rt.NewHostModuleBuilder("example").
		NewFunctionBuilder().
		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
			// Because we enabled snapshots with WithSnapshotter, this is non-nil.
			snapshot := experimental.GetSnapshotter(ctx).Snapshot()

			// Get our mutable snapshots holder to be able to add to it. Our example only calls snapshot
			// and restore once but real programs will often call them at multiple layers within a call
			// stack with various e.g., try/catch statements.
			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
			idx := len(*snapshots)
			*snapshots = append(*snapshots, snapshot)

			// Write a value to be passed back to restore. This is meant to be opaque to the guest
			// and used to re-reference the snapshot.
			ok := mod.Memory().WriteUint32Le(snapshotPtr, uint32(idx))
			if !ok {
				log.Panicln("failed to write snapshot index")
			}

			return 0
		}).
		Export("snapshot").
		NewFunctionBuilder().
		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) {
			// Read the value written by snapshot to re-reference the snapshot.
			idx, ok := mod.Memory().ReadUint32Le(snapshotPtr)
			if !ok {
				log.Panicln("failed to read snapshot index")
			}

			// Get the snapshot
			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
			snapshot := (*snapshots)[idx]

			// Restore! The invocation of this function will end as soon as we invoke
			// Restore, so we also pass in our return value. The guest function run
			// will finish with this return value.
			snapshot.Restore([]uint64{5})
		}).
		Export("restore").
		Instantiate(ctx)
	if err != nil {
		log.Panicln(err)
	}

	mod, err := rt.Instantiate(ctx, snapshotWasm) // Instantiate the actual code
	if err != nil {
		log.Panicln(err)
	}

	// Call the guest entrypoint.
	res, err := mod.ExportedFunction("run").Call(ctx)
	if err != nil {
		log.Panicln(err)
	}
	// We restored and returned the restore value, so it's our result. If restore
	// was instead a no-op, we would have returned 10 from normal code flow.
	fmt.Println(res[0])
}

Output:

5

package main

import (
	"fmt"

	_ "embed"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/experimental"
	"github.com/tetratelabs/wazero/internal/wasm"
)

func main() {
	it := &fakeStackIterator{}

	for it.Next() {
		fn := it.Function()
		pc := it.ProgramCounter()
		fmt.Println("function:", fn.Definition().DebugName())
		fmt.Println("\tprogram counter:", pc)
		fmt.Println("\tsource offset:", fn.SourceOffsetForPC(pc))
	}

}

type fakeStackIterator struct {
	iteration    int
	def          api.FunctionDefinition
	args         []uint64
	pc           uint64
	sourceOffset uint64
}

func (s *fakeStackIterator) Next() bool {
	switch s.iteration {
	case 0:
		s.def = &mockFunctionDefinition{debugName: "fn0"}
		s.args = []uint64{1, 2, 3}
		s.pc = 5890831
		s.sourceOffset = 1234
	case 1:
		s.def = &mockFunctionDefinition{debugName: "fn1"}
		s.args = []uint64{}
		s.pc = 5899822
		s.sourceOffset = 7286
	case 2:
		s.def = &mockFunctionDefinition{debugName: "fn2"}
		s.args = []uint64{4}
		s.pc = 6820312
		s.sourceOffset = 935891
	case 3:
		return false
	}
	s.iteration++
	return true
}

func (s *fakeStackIterator) Function() experimental.InternalFunction {
	return internalFunction{
		definition:   s.def,
		sourceOffset: s.sourceOffset,
	}
}

func (s *fakeStackIterator) ProgramCounter() experimental.ProgramCounter {
	return experimental.ProgramCounter(s.pc)
}

type internalFunction struct {
	definition   api.FunctionDefinition
	sourceOffset uint64
}

func (f internalFunction) Definition() api.FunctionDefinition {
	return f.definition
}

func (f internalFunction) SourceOffsetForPC(pc experimental.ProgramCounter) uint64 {
	return f.sourceOffset
}

type mockFunctionDefinition struct {
	debugName string
	*wasm.FunctionDefinition
}

func (f *mockFunctionDefinition) DebugName() string {
	return f.debugName
}

func (f *mockFunctionDefinition) ParamTypes() []wasm.ValueType {
	return []wasm.ValueType{}
}

func (f *mockFunctionDefinition) ResultTypes() []wasm.ValueType {
	return []wasm.ValueType{}
}

Output:

function: fn0
	program counter: 5890831
	source offset: 1234
function: fn1
	program counter: 5899822
	source offset: 7286
function: fn2
	program counter: 6820312
	source offset: 935891