From 4c5c9309be3a6c4f5ac70c62a04f88c28b9cac9d Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 16 Oct 2023 16:19:37 +0900 Subject: [PATCH 01/10] Add experimental support for snapshot/restore Signed-off-by: Anuraag Agrawal --- experimental/checkpoint.go | 27 +++++++++++++++++++ internal/engine/compiler/engine.go | 43 ++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) create mode 100644 experimental/checkpoint.go diff --git a/experimental/checkpoint.go b/experimental/checkpoint.go new file mode 100644 index 0000000000..8791abd29b --- /dev/null +++ b/experimental/checkpoint.go @@ -0,0 +1,27 @@ +package experimental + +// Snapshot holds the execution state at the time of a Snapshotter.Snapshot call. +type Snapshot interface { + // Restore sets the Wasm execution state to the capture. Because a host function + // calling this is resetting the pointer to the executation stack, the host function + // will not be able to return values in the normal way. ret is a slice of values the + // host function intends to return from the restored function. + Restore(ret []uint64) +} + +// Snapshotter allows host functions to snapshot the WebAssembly execution environment. +// Currently, only the Wasm stack is captured, but in the future, this may be expanded +// to things like globals. +type Snapshotter interface { + // Snapshot captures the current execution state. + Snapshot() Snapshot +} + +// EnableSnapshotterKey is a context key to indicate that snapshotting should be enabled. +// The context.Context passed to a exported function invocation should have this key set +// to a non-nil value, and host functions will be able to retrieve it using SnapshotterKey. +type EnableSnapshotterKey struct{} + +// SnapshotterKey is a context key to access a Snapshotter from a host function. +// It is only present if EnableSnapshotter was set in the function invocation context. +type SnapshotterKey struct{} diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index d4e8dcd06f..efaf6cc565 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -773,6 +773,10 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u defer done() } + if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { + ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, ce) + } + ce.execWasmFunction(ctx, m) // This returns a safe copy of the results, instead of a slice view. If we @@ -1141,6 +1145,45 @@ func (ce *callEngine) builtinFunctionTableGrow(tables []*wasm.TableInstance) { ce.pushValue(uint64(res)) } +// snapshot implements experimental.Snapshot +type snapshot struct { + stackPointer uint64 + stackBasePointerInBytes uint64 + returnAddress uint64 + hostBase int + stack []uint64 + + ce *callEngine +} + +// Snapshot implements the same method as documented on experimental.Snapshotter. +func (ce *callEngine) Snapshot() experimental.Snapshot { + hostBase := int(ce.stackBasePointerInBytes >> 3) + + stackTop := int(ce.stackTopIndex()) + stack := make([]uint64, stackTop) + copy(stack, ce.stack[:stackTop]) + + return &snapshot{ + stackPointer: ce.stackContext.stackPointer, + stackBasePointerInBytes: ce.stackBasePointerInBytes, + returnAddress: uint64(ce.returnAddress), + hostBase: hostBase, + stack: stack, + ce: ce, + } +} + +// Restore implements the same method as documented on experimental.Snapshot. +func (s *snapshot) Restore(ret []uint64) { + ce := s.ce + ce.stackContext.stackPointer = s.stackPointer + ce.stackContext.stackBasePointerInBytes = s.stackBasePointerInBytes + copy(ce.stack, s.stack) + ce.returnAddress = uintptr(s.returnAddress) + copy(ce.stack[s.hostBase:], ret) +} + // stackIterator implements experimental.StackIterator. type stackIterator struct { stack []uint64 From 3444522228932170d8250984b95defa06f47d6c3 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 16 Oct 2023 16:19:37 +0900 Subject: [PATCH 02/10] Implement snapshot/restore Signed-off-by: Anuraag Agrawal --- internal/engine/compiler/engine.go | 74 +++++++++++++++++++++++++++--- 1 file changed, 68 insertions(+), 6 deletions(-) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index efaf6cc565..a0dc0118df 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -1052,12 +1052,27 @@ entry: stack := ce.stack[base : base+stackLen] fn := calleeHostFunction.parent.goFunc - switch fn := fn.(type) { - case api.GoModuleFunction: - fn.Call(ctx, ce.callerModuleInstance, stack) - case api.GoFunction: - fn.Call(ctx, stack) - } + func() { + defer func() { + if r := recover(); r != nil { + if s, ok := r.(*snapshot); ok { + if s.ce == ce { + s.doRestore() + } else { + panic(r) + } + } else { + panic(r) + } + } + }() + switch fn := fn.(type) { + case api.GoModuleFunction: + fn.Call(ctx, ce.callerModuleInstance, stack) + case api.GoFunction: + fn.Call(ctx, stack) + } + }() codeAddr, modAddr = ce.returnAddress, ce.moduleInstance goto entry @@ -1184,6 +1199,53 @@ func (s *snapshot) Restore(ret []uint64) { copy(ce.stack[s.hostBase:], ret) } +// snapshot implements experimental.Snapshot +type snapshot struct { + stackPointer uint64 + stackBasePointerInBytes uint64 + returnAddress uint64 + hostBase int + stack []uint64 + + ret []uint64 + + ce *callEngine +} + +// Snapshot implements the same method as documented on experimental.Snapshotter. +func (ce *callEngine) Snapshot() experimental.Snapshot { + hostBase := int(ce.stackBasePointerInBytes >> 3) + + stackTop := int(ce.stackTopIndex()) + stack := make([]uint64, stackTop) + copy(stack, ce.stack[:stackTop]) + + return &snapshot{ + stackPointer: ce.stackContext.stackPointer, + stackBasePointerInBytes: ce.stackBasePointerInBytes, + returnAddress: uint64(ce.returnAddress), + hostBase: hostBase, + stack: stack, + ce: ce, + } +} + +// Restore implements the same method as documented on experimental.Snapshot. +func (s *snapshot) Restore(ret []uint64) { + s.ret = ret + panic(s) +} + +// Restore implements the same method as documented on experimental.Snapshot. +func (s *snapshot) doRestore() { + ce := s.ce + ce.stackContext.stackPointer = s.stackPointer + ce.stackContext.stackBasePointerInBytes = s.stackBasePointerInBytes + copy(ce.stack, s.stack) + ce.returnAddress = uintptr(s.returnAddress) + copy(ce.stack[s.hostBase:], s.ret) +} + // stackIterator implements experimental.StackIterator. type stackIterator struct { stack []uint64 From 21eb098915011c1f14714a56ef16486bfb9317e1 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Tue, 17 Oct 2023 11:17:02 +0900 Subject: [PATCH 03/10] Fix merge Signed-off-by: Anuraag Agrawal --- internal/engine/compiler/engine.go | 39 ------------------------------ 1 file changed, 39 deletions(-) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index a0dc0118df..e7527a2433 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -1160,45 +1160,6 @@ func (ce *callEngine) builtinFunctionTableGrow(tables []*wasm.TableInstance) { ce.pushValue(uint64(res)) } -// snapshot implements experimental.Snapshot -type snapshot struct { - stackPointer uint64 - stackBasePointerInBytes uint64 - returnAddress uint64 - hostBase int - stack []uint64 - - ce *callEngine -} - -// Snapshot implements the same method as documented on experimental.Snapshotter. -func (ce *callEngine) Snapshot() experimental.Snapshot { - hostBase := int(ce.stackBasePointerInBytes >> 3) - - stackTop := int(ce.stackTopIndex()) - stack := make([]uint64, stackTop) - copy(stack, ce.stack[:stackTop]) - - return &snapshot{ - stackPointer: ce.stackContext.stackPointer, - stackBasePointerInBytes: ce.stackBasePointerInBytes, - returnAddress: uint64(ce.returnAddress), - hostBase: hostBase, - stack: stack, - ce: ce, - } -} - -// Restore implements the same method as documented on experimental.Snapshot. -func (s *snapshot) Restore(ret []uint64) { - ce := s.ce - ce.stackContext.stackPointer = s.stackPointer - ce.stackContext.stackBasePointerInBytes = s.stackBasePointerInBytes - copy(ce.stack, s.stack) - ce.returnAddress = uintptr(s.returnAddress) - copy(ce.stack[s.hostBase:], ret) -} - // snapshot implements experimental.Snapshot type snapshot struct { stackPointer uint64 From 9d169f53f7bc8e6456f8834093318f06216bb78a Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Oct 2023 14:07:23 +0900 Subject: [PATCH 04/10] Tests Signed-off-by: Anuraag Agrawal --- experimental/checkpoint_example_test.go | 100 ++++++++++++++++++++ experimental/checkpoint_test.go | 118 ++++++++++++++++++++++++ experimental/listener_example_test.go | 2 +- experimental/testdata/snapshot.wasm | Bin 0 -> 163 bytes experimental/testdata/snapshot.wat | 34 +++++++ internal/engine/compiler/engine.go | 10 ++ 6 files changed, 263 insertions(+), 1 deletion(-) create mode 100644 experimental/checkpoint_example_test.go create mode 100644 experimental/checkpoint_test.go create mode 100644 experimental/testdata/snapshot.wasm create mode 100644 experimental/testdata/snapshot.wat diff --git a/experimental/checkpoint_example_test.go b/experimental/checkpoint_example_test.go new file mode 100644 index 0000000000..d995befea3 --- /dev/null +++ b/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 +} diff --git a/experimental/checkpoint_test.go b/experimental/checkpoint_test.go new file mode 100644 index 0000000000..14bb4629fc --- /dev/null +++ b/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) +} diff --git a/experimental/listener_example_test.go b/experimental/listener_example_test.go index 57d346ccdf..0d2ece4102 100644 --- a/experimental/listener_example_test.go +++ b/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 diff --git a/experimental/testdata/snapshot.wasm b/experimental/testdata/snapshot.wasm new file mode 100644 index 0000000000000000000000000000000000000000..b07f5f2a1b27d5db6e0e84730c466baa644817cd GIT binary patch literal 163 zcmYLAy$ZrG7`*$EG^GK_-~-e_Hy@&J(hLCy)0U)yPI+|$qTXxy2k|1 XX>dofHj5KR=sLo9ajt8@<0Eu$px+-v literal 0 HcmV?d00001 diff --git a/experimental/testdata/snapshot.wat b/experimental/testdata/snapshot.wat new file mode 100644 index 0000000000..687142376f --- /dev/null +++ b/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) +) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index e7527a2433..ef42db7425 100644 --- a/internal/engine/compiler/engine.go +++ b/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 From b3d34ad3200852c8fa378803132ab4f0e2a9766d Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Oct 2023 14:41:21 +0900 Subject: [PATCH 05/10] Implement interpreter Signed-off-by: Anuraag Agrawal --- experimental/checkpoint_test.go | 10 ++- internal/engine/compiler/engine.go | 2 +- internal/engine/interpreter/interpreter.go | 77 +++++++++++++++++++++- 3 files changed, 84 insertions(+), 5 deletions(-) diff --git a/experimental/checkpoint_test.go b/experimental/checkpoint_test.go index 14bb4629fc..868933b1a5 100644 --- a/experimental/checkpoint_test.go +++ b/experimental/checkpoint_test.go @@ -112,7 +112,11 @@ func TestSnapshotMultipleWasmInvocations(t *testing.T) { // 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) + // Fails, snapshot and restore are called from different wasm invocations. Currently, this + // results in a panic. + err = require.CapturePanic(func() { + _, _ = mod.ExportedFunction("restore").Call(ctx, snapshotPtr) + }) + require.EqualError(t, err, "unhandled snapshot restore, this generally indicates restore was called from a different "+ + "exported function invocation than snapshot") } diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index ef42db7425..e29bd09898 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -1202,7 +1202,6 @@ func (s *snapshot) Restore(ret []uint64) { panic(s) } -// Restore implements the same method as documented on experimental.Snapshot. func (s *snapshot) doRestore() { ce := s.ce ce.stackContext.stackPointer = s.stackPointer @@ -1212,6 +1211,7 @@ func (s *snapshot) doRestore() { copy(ce.stack[s.hostBase:], s.ret) } +// Error implements the same method on error. func (s *snapshot) Error() string { return "unhandled snapshot restore, this generally indicates restore was called from a different " + "exported function invocation than snapshot" diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go index d47e7c3129..cddca1038c 100644 --- a/internal/engine/interpreter/interpreter.go +++ b/internal/engine/interpreter/interpreter.go @@ -216,6 +216,53 @@ func functionFromUintptr(ptr uintptr) *function { return *(**function)(unsafe.Pointer(wrapped)) } +type snapshot struct { + stack []uint64 + frames []*callFrame + pc uint64 + + ret []uint64 + + ce *callEngine +} + +// Snapshot implements the same method as documented on experimental.Snapshotter. +func (ce *callEngine) Snapshot() experimental.Snapshot { + stack := make([]uint64, len(ce.stack)) + copy(stack, ce.stack) + + frames := make([]*callFrame, len(ce.frames)) + copy(frames, ce.frames) + + return &snapshot{ + stack: stack, + frames: frames, + ce: ce, + } +} + +// Restore implements the same method as documented on experimental.Snapshot. +func (s *snapshot) Restore(ret []uint64) { + s.ret = ret + panic(s) +} + +func (s *snapshot) doRestore() { + ce := s.ce + + ce.stack = s.stack + ce.frames = s.frames + ce.frames[len(ce.frames)-1].pc = s.pc + + copy(ce.stack[len(ce.stack)-len(s.ret):], s.ret) +} + +// Error implements the same method on error. +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 @@ -512,6 +559,10 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u } } + if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { + ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, ce) + } + defer func() { // If the module closed during the call, and the call didn't err for another reason, set an ExitError. if err == nil { @@ -555,6 +606,12 @@ type functionListenerInvocation struct { // with the call frame stack traces. Also, reset the state of callEngine // so that it can be used for the subsequent calls. func (ce *callEngine) recoverOnCall(ctx context.Context, m *wasm.ModuleInstance, v interface{}) (err error) { + if s, ok := v.(*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) + } + builder := wasmdebug.NewErrorBuilder() frameCount := len(ce.frames) functionListeners := make([]functionListenerInvocation, 0, 16) @@ -669,7 +726,25 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance ce.drop(op.Us[v+1]) frame.pc = op.Us[v] case wazeroir.OperationKindCall: - ce.callFunction(ctx, f.moduleInstance, &functions[op.U1]) + func() { + defer func() { + if r := recover(); r != nil { + if s, ok := r.(*snapshot); ok { + if s.ce == ce { + s.doRestore() + frame = ce.frames[len(ce.frames)-1] + body = frame.f.parent.body + bodyLen = uint64(len(body)) + } else { + panic(r) + } + } else { + panic(r) + } + } + }() + ce.callFunction(ctx, f.moduleInstance, &functions[op.U1]) + }() frame.pc++ case wazeroir.OperationKindCallIndirect: offset := ce.popValue() From 88ab51110d291f4f61e1f2bf78addbea41698c10 Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Oct 2023 14:52:48 +0900 Subject: [PATCH 06/10] Stack only Signed-off-by: Anuraag Agrawal --- experimental/checkpoint.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/experimental/checkpoint.go b/experimental/checkpoint.go index 8791abd29b..1560e6ba49 100644 --- a/experimental/checkpoint.go +++ b/experimental/checkpoint.go @@ -10,8 +10,6 @@ type Snapshot interface { } // Snapshotter allows host functions to snapshot the WebAssembly execution environment. -// Currently, only the Wasm stack is captured, but in the future, this may be expanded -// to things like globals. type Snapshotter interface { // Snapshot captures the current execution state. Snapshot() Snapshot From f7dfa0b4f82e71dd9f3dc2264f6bcfe9b4b25f7b Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Fri, 20 Oct 2023 14:54:50 +0900 Subject: [PATCH 07/10] Drift Signed-off-by: Anuraag Agrawal --- experimental/checkpoint_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/experimental/checkpoint_test.go b/experimental/checkpoint_test.go index 868933b1a5..3abf921f9b 100644 --- a/experimental/checkpoint_test.go +++ b/experimental/checkpoint_test.go @@ -2,7 +2,6 @@ package experimental_test import ( "context" - _ "embed" "testing" "github.com/tetratelabs/wazero" From 04cf900f1edc84b16c0194f9013da9b19c70d12e Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Mon, 22 Jan 2024 12:11:38 +0900 Subject: [PATCH 08/10] Only defer when snapshotting Signed-off-by: Anuraag Agrawal --- internal/engine/compiler/engine.go | 26 +++++++++++++--------- internal/engine/interpreter/interpreter.go | 26 ++++++++++++---------- 2 files changed, 29 insertions(+), 23 deletions(-) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index 5ddeeb9acb..8779ada3e8 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -787,11 +787,13 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u defer done() } + snapshotEnabled := false if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, ce) + snapshotEnabled = true } - ce.execWasmFunction(ctx, m) + ce.execWasmFunction(ctx, m, snapshotEnabled) // This returns a safe copy of the results, instead of a slice view. If we // returned a re-slice, the caller could accidentally or purposefully @@ -1048,7 +1050,7 @@ const ( builtinFunctionMemoryNotify ) -func (ce *callEngine) execWasmFunction(ctx context.Context, m *wasm.ModuleInstance) { +func (ce *callEngine) execWasmFunction(ctx context.Context, m *wasm.ModuleInstance, snapshotEnabled bool) { codeAddr := ce.initialFn.codeInitialAddress modAddr := ce.initialFn.moduleInstance @@ -1075,19 +1077,21 @@ entry: fn := calleeHostFunction.parent.goFunc func() { - defer func() { - if r := recover(); r != nil { - if s, ok := r.(*snapshot); ok { - if s.ce == ce { - s.doRestore() + if snapshotEnabled { + defer func() { + if r := recover(); r != nil { + if s, ok := r.(*snapshot); ok { + if s.ce == ce { + s.doRestore() + } else { + panic(r) + } } else { panic(r) } - } else { - panic(r) } - } - }() + }() + } switch fn := fn.(type) { case api.GoModuleFunction: fn.Call(ctx, ce.callerModuleInstance, stack) diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go index 71b1bad096..4f1fae2d85 100644 --- a/internal/engine/interpreter/interpreter.go +++ b/internal/engine/interpreter/interpreter.go @@ -735,22 +735,24 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance frame.pc = op.Us[v] case wazeroir.OperationKindCall: func() { - defer func() { - if r := recover(); r != nil { - if s, ok := r.(*snapshot); ok { - if s.ce == ce { - s.doRestore() - frame = ce.frames[len(ce.frames)-1] - body = frame.f.parent.body - bodyLen = uint64(len(body)) + if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { + defer func() { + if r := recover(); r != nil { + if s, ok := r.(*snapshot); ok { + if s.ce == ce { + s.doRestore() + frame = ce.frames[len(ce.frames)-1] + body = frame.f.parent.body + bodyLen = uint64(len(body)) + } else { + panic(r) + } } else { panic(r) } - } else { - panic(r) } - } - }() + }() + } ce.callFunction(ctx, f.moduleInstance, &functions[op.U1]) }() frame.pc++ From 6abe8602f0dd5b663d81201ccef17e4c87ed7f9c Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Tue, 23 Jan 2024 10:26:38 +0900 Subject: [PATCH 09/10] cleanup Signed-off-by: Anuraag Agrawal --- internal/engine/compiler/engine.go | 8 ++------ internal/engine/interpreter/interpreter.go | 14 +++++--------- 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index 8779ada3e8..d83cd10406 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -1080,12 +1080,8 @@ entry: if snapshotEnabled { defer func() { if r := recover(); r != nil { - if s, ok := r.(*snapshot); ok { - if s.ce == ce { - s.doRestore() - } else { - panic(r) - } + if s, ok := r.(*snapshot); ok && s.ce == ce { + s.doRestore() } else { panic(r) } diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go index 4f1fae2d85..29c44c39bb 100644 --- a/internal/engine/interpreter/interpreter.go +++ b/internal/engine/interpreter/interpreter.go @@ -738,15 +738,11 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { defer func() { if r := recover(); r != nil { - if s, ok := r.(*snapshot); ok { - if s.ce == ce { - s.doRestore() - frame = ce.frames[len(ce.frames)-1] - body = frame.f.parent.body - bodyLen = uint64(len(body)) - } else { - panic(r) - } + if s, ok := r.(*snapshot); ok && s.ce == ce { + s.doRestore() + frame = ce.frames[len(ce.frames)-1] + body = frame.f.parent.body + bodyLen = uint64(len(body)) } else { panic(r) } From 2d2d62c8d17c1a82c1c125672bb0b220ac867f8e Mon Sep 17 00:00:00 2001 From: Anuraag Agrawal Date: Wed, 24 Jan 2024 08:35:22 +0900 Subject: [PATCH 10/10] cleanup Signed-off-by: Anuraag Agrawal --- internal/engine/compiler/engine.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go index d83cd10406..babf5c0d32 100644 --- a/internal/engine/compiler/engine.go +++ b/internal/engine/compiler/engine.go @@ -787,10 +787,9 @@ func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []u defer done() } - snapshotEnabled := false - if ctx.Value(experimental.EnableSnapshotterKey{}) != nil { + snapshotEnabled := ctx.Value(experimental.EnableSnapshotterKey{}) != nil + if snapshotEnabled { ctx = context.WithValue(ctx, experimental.SnapshotterKey{}, ce) - snapshotEnabled = true } ce.execWasmFunction(ctx, m, snapshotEnabled)