Tests

Signed-off-by: Anuraag Agrawal <anuraaga@gmail.com>

experimental/checkpoint_example_test.go

@@ -0,0 +1,100 @@
+package experimental_test
+
+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 Example_enableSnapshotterKey() {
+	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 = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+
+	// 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 set EnableSnapshotterKey to context, this is non-nil.
+			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).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
+}

experimental/checkpoint_test.go

@@ -0,0 +1,118 @@
+package experimental_test
+
+import (
+	"context"
+	_ "embed"
+	"testing"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestSnapshotNestedWasmInvocation(t *testing.T) {
+	ctx := context.Background()
+
+	rt := wazero.NewRuntime(ctx)
+	defer rt.Close(ctx)
+
+	sidechannel := 0
+
+	_, err := rt.NewHostModuleBuilder("example").
+		NewFunctionBuilder().
+		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
+			defer func() {
+				sidechannel = 10
+			}()
+			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
+			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
+			idx := len(*snapshots)
+			*snapshots = append(*snapshots, snapshot)
+			ok := mod.Memory().WriteUint32Le(snapshotPtr, uint32(idx))
+			require.True(t, ok)
+
+			_, err := mod.ExportedFunction("restore").Call(ctx, uint64(snapshotPtr))
+			require.NoError(t, err)
+
+			return 2
+		}).
+		Export("snapshot").
+		NewFunctionBuilder().
+		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) {
+			idx, ok := mod.Memory().ReadUint32Le(snapshotPtr)
+			require.True(t, ok)
+			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
+			snapshot := (*snapshots)[idx]
+
+			snapshot.Restore([]uint64{12})
+		}).
+		Export("restore").
+		Instantiate(ctx)
+	require.NoError(t, err)
+
+	mod, err := rt.Instantiate(ctx, snapshotWasm)
+	require.NoError(t, err)
+
+	var snapshots []experimental.Snapshot
+	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)
+	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+
+	snapshotPtr := uint64(0)
+	res, err := mod.ExportedFunction("snapshot").Call(ctx, snapshotPtr)
+	require.NoError(t, err)
+	// return value from restore
+	require.Equal(t, uint64(12), res[0])
+	// Host function defers within the call stack work fine
+	require.Equal(t, 10, sidechannel)
+}
+
+func TestSnapshotMultipleWasmInvocations(t *testing.T) {
+	ctx := context.Background()
+
+	rt := wazero.NewRuntime(ctx)
+	defer rt.Close(ctx)
+
+	_, err := rt.NewHostModuleBuilder("example").
+		NewFunctionBuilder().
+		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
+			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
+			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
+			idx := len(*snapshots)
+			*snapshots = append(*snapshots, snapshot)
+			ok := mod.Memory().WriteUint32Le(snapshotPtr, uint32(idx))
+			require.True(t, ok)
+
+			return 0
+		}).
+		Export("snapshot").
+		NewFunctionBuilder().
+		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) {
+			idx, ok := mod.Memory().ReadUint32Le(snapshotPtr)
+			require.True(t, ok)
+			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
+			snapshot := (*snapshots)[idx]
+
+			snapshot.Restore([]uint64{12})
+		}).
+		Export("restore").
+		Instantiate(ctx)
+	require.NoError(t, err)
+
+	mod, err := rt.Instantiate(ctx, snapshotWasm)
+	require.NoError(t, err)
+
+	var snapshots []experimental.Snapshot
+	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)
+	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
+
+	snapshotPtr := uint64(0)
+	res, err := mod.ExportedFunction("snapshot").Call(ctx, snapshotPtr)
+	require.NoError(t, err)
+	// snapshot returned zero
+	require.Equal(t, uint64(0), res[0])
+
+	res, err = mod.ExportedFunction("restore").Call(ctx, snapshotPtr)
+	// Fails, snapshot and restore are called from different wasm invocations.
+	require.Error(t, err)
+}

experimental/listener_example_test.go

@@ -16,7 +16,7 @@ import (
 
 // listenerWasm was generated by the following:
 //
-//	cd testdata; wat2wasm --debug-names listener.wat
+//	cd logging/testdata; wat2wasm --debug-names listener.wat
 //
 //go:embed logging/testdata/listener.wasm
 var listenerWasm []byte

experimental/testdata/snapshot.wat

@@ -0,0 +1,34 @@
+(module
+  (import "example" "snapshot" (func $snapshot (param i32) (result i32)))
+  (import "example" "restore" (func $restore (param i32)))
+
+  (func $helper (result i32)
+    (call $restore (i32.const 0))
+    ;; Not executed
+    i32.const 10
+  )
+
+  (func (export "run") (result i32) (local i32)
+    (call $snapshot (i32.const 0))
+    local.set 0
+    local.get 0
+    (if (result i32)
+      (then ;; restore return, finish with the value returned by it
+        local.get 0
+      )
+      (else ;; snapshot return, call heloer
+        (call $helper)
+      )
+    )
+  )
+
+  (func (export "snapshot") (param i32) (result i32)
+    (call $snapshot (local.get 0))
+  )
+
+  (func (export "restore") (param i32)
+    (call $restore (local.get 0))
+  )
+
+  (memory (export "memory") 1 1)
+)

internal/engine/compiler/engine.go

@@ -843,6 +843,11 @@ func callFrameOffset(funcType *wasm.FunctionType) (ret int) {
 //
 // This is defined for testability.
 func (ce *callEngine) deferredOnCall(ctx context.Context, m *wasm.ModuleInstance, recovered interface{}) (err error) {
+	if s, ok := recovered.(*snapshot); ok {
+		// A snapshot that wasn't handled was created by a different call engine possibly from a nested wasm invocation,
+		// let it propagate up to be handled by the caller.
+		panic(s)
+	}
 	if recovered != nil {
 		builder := wasmdebug.NewErrorBuilder()
 
@@ -1207,6 +1212,11 @@ func (s *snapshot) doRestore() {
 	copy(ce.stack[s.hostBase:], s.ret)
 }
 
+func (s *snapshot) Error() string {
+	return "unhandled snapshot restore, this generally indicates restore was called from a different " +
+		"exported function invocation than snapshot"
+}
+
 // stackIterator implements experimental.StackIterator.
 type stackIterator struct {
 	stack   []uint64