github.com/wasilibs/wazerox@v0.0.0-20240124024944-4923be63ab5f/experimental/checkpoint_example_test.go (about)

     1  package experimental_test
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"fmt"
     7  	"log"
     8  
     9  	wazero "github.com/wasilibs/wazerox"
    10  	"github.com/wasilibs/wazerox/api"
    11  	"github.com/wasilibs/wazerox/experimental"
    12  )
    13  
    14  // snapshotWasm was generated by the following:
    15  //
    16  //	cd testdata; wat2wasm snapshot.wat
    17  //
    18  //go:embed testdata/snapshot.wasm
    19  var snapshotWasm []byte
    20  
    21  type snapshotsKey struct{}
    22  
    23  func Example_enableSnapshotterKey() {
    24  	ctx := context.Background()
    25  
    26  	rt := wazero.NewRuntime(ctx)
    27  	defer rt.Close(ctx) // This closes everything this Runtime created.
    28  
    29  	// Enable experimental snapshotting functionality by setting it to context. We use this
    30  	// context when invoking functions, indicating to wazero to enable it.
    31  	ctx = context.WithValue(ctx, experimental.EnableSnapshotterKey{}, struct{}{})
    32  
    33  	// Also place a mutable holder of snapshots to be referenced during restore.
    34  	var snapshots []experimental.Snapshot
    35  	ctx = context.WithValue(ctx, snapshotsKey{}, &snapshots)
    36  
    37  	// Register host functions using snapshot and restore. Generally snapshot is saved
    38  	// into a mutable location in context to be referenced during restore.
    39  	_, err := rt.NewHostModuleBuilder("example").
    40  		NewFunctionBuilder().
    41  		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) int32 {
    42  			// Because we set EnableSnapshotterKey to context, this is non-nil.
    43  			snapshot := ctx.Value(experimental.SnapshotterKey{}).(experimental.Snapshotter).Snapshot()
    44  
    45  			// Get our mutable snapshots holder to be able to add to it. Our example only calls snapshot
    46  			// and restore once but real programs will often call them at multiple layers within a call
    47  			// stack with various e.g., try/catch statements.
    48  			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
    49  			idx := len(*snapshots)
    50  			*snapshots = append(*snapshots, snapshot)
    51  
    52  			// Write a value to be passed back to restore. This is meant to be opaque to the guest
    53  			// and used to re-reference the snapshot.
    54  			ok := mod.Memory().WriteUint32Le(snapshotPtr, uint32(idx))
    55  			if !ok {
    56  				log.Panicln("failed to write snapshot index")
    57  			}
    58  
    59  			return 0
    60  		}).
    61  		Export("snapshot").
    62  		NewFunctionBuilder().
    63  		WithFunc(func(ctx context.Context, mod api.Module, snapshotPtr uint32) {
    64  			// Read the value written by snapshot to re-reference the snapshot.
    65  			idx, ok := mod.Memory().ReadUint32Le(snapshotPtr)
    66  			if !ok {
    67  				log.Panicln("failed to read snapshot index")
    68  			}
    69  
    70  			// Get the snapshot
    71  			snapshots := ctx.Value(snapshotsKey{}).(*[]experimental.Snapshot)
    72  			snapshot := (*snapshots)[idx]
    73  
    74  			// Restore! The invocation of this function will end as soon as we invoke
    75  			// Restore, so we also pass in our return value. The guest function run
    76  			// will finish with this return value.
    77  			snapshot.Restore([]uint64{5})
    78  		}).
    79  		Export("restore").
    80  		Instantiate(ctx)
    81  	if err != nil {
    82  		log.Panicln(err)
    83  	}
    84  
    85  	mod, err := rt.Instantiate(ctx, snapshotWasm) // Instantiate the actual code
    86  	if err != nil {
    87  		log.Panicln(err)
    88  	}
    89  
    90  	// Call the guest entrypoint.
    91  	res, err := mod.ExportedFunction("run").Call(ctx)
    92  	if err != nil {
    93  		log.Panicln(err)
    94  	}
    95  	// We restored and returned the restore value, so it's our result. If restore
    96  	// was instead a no-op, we would have returned 10 from normal code flow.
    97  	fmt.Println(res[0])
    98  	// Output:
    99  	// 5
   100  }