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

     1  package wazero_test
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"fmt"
     7  	"log"
     8  	"time"
     9  
    10  	wazero "github.com/wasilibs/wazerox"
    11  )
    12  
    13  // infiniteLoopWasm exports a function named "infinite_loop" that never exits.
    14  //
    15  //go:embed internal/integration_test/engine/testdata/infinite_loop.wasm
    16  var infiniteLoopWasm []byte
    17  
    18  // ExampleRuntimeConfig_WithCloseOnContextDone_context_timeout demonstrates how to ensure the termination
    19  // of infinite loop function with context.Context created by context.WithTimeout powered by
    20  // RuntimeConfig.WithEnsureTermination configuration.
    21  func ExampleRuntimeConfig_WithCloseOnContextDone_context_timeout() {
    22  	ctx := context.Background()
    23  
    24  	r := wazero.NewRuntimeWithConfig(ctx,
    25  		// Enables the WithCloseOnContextDone option.
    26  		wazero.NewRuntimeConfig().WithCloseOnContextDone(true))
    27  	defer r.Close(ctx)
    28  
    29  	moduleInstance, err := r.InstantiateWithConfig(ctx, infiniteLoopWasm,
    30  		wazero.NewModuleConfig().WithName("malicious_wasm"))
    31  	if err != nil {
    32  		log.Panicln(err)
    33  	}
    34  
    35  	infiniteLoop := moduleInstance.ExportedFunction("infinite_loop")
    36  
    37  	// Create the context.Context to be passed to the invocation of infinite_loop.
    38  	ctx, cancel := context.WithTimeout(ctx, time.Second)
    39  	defer cancel()
    40  
    41  	// Invoke the infinite loop with the timeout context.
    42  	_, err = infiniteLoop.Call(ctx)
    43  
    44  	// Timeout is correctly handled and triggers the termination of infinite loop.
    45  	fmt.Println(err)
    46  
    47  	// Output:
    48  	//	module closed with context deadline exceeded
    49  }
    50  
    51  // ExampleRuntimeConfig_WithCloseOnContextDone_context_cancel demonstrates how to ensure the termination
    52  // of infinite loop function with context.Context created by context.WithCancel powered by
    53  // RuntimeConfig.WithEnsureTermination configuration.
    54  func ExampleRuntimeConfig_WithCloseOnContextDone_context_cancel() {
    55  	ctx := context.Background()
    56  
    57  	r := wazero.NewRuntimeWithConfig(ctx,
    58  		// Enables the WithCloseOnContextDone option.
    59  		wazero.NewRuntimeConfig().WithCloseOnContextDone(true))
    60  	defer r.Close(ctx)
    61  
    62  	moduleInstance, err := r.InstantiateWithConfig(ctx, infiniteLoopWasm,
    63  		wazero.NewModuleConfig().WithName("malicious_wasm"))
    64  	if err != nil {
    65  		log.Panicln(err)
    66  	}
    67  
    68  	infiniteLoop := moduleInstance.ExportedFunction("infinite_loop")
    69  
    70  	// Create the context.Context to be passed to the invocation of infinite_loop.
    71  	ctx, cancel := context.WithCancel(ctx)
    72  	go func() {
    73  		// After 2 seconds, cancel the invocation of infinite loop.
    74  		time.Sleep(2 * time.Second)
    75  		cancel()
    76  	}()
    77  
    78  	// Invoke the infinite loop with the timeout context.
    79  	_, err = infiniteLoop.Call(ctx)
    80  
    81  	// context Cancellation is correctly handled and triggers the termination of infinite loop.
    82  	fmt.Println(err)
    83  
    84  	// Output:
    85  	//	module closed with context canceled
    86  }
    87  
    88  // ExampleRuntimeConfig_WithCloseOnContextDone_moduleClose demonstrates how to ensure the termination
    89  // of infinite loop function with api.Module's CloseWithExitCode method powered by
    90  // RuntimeConfig.WithEnsureTermination configuration.
    91  func ExampleRuntimeConfig_WithCloseOnContextDone_moduleClose() {
    92  	ctx := context.Background()
    93  
    94  	r := wazero.NewRuntimeWithConfig(ctx,
    95  		// Enables the WithCloseOnContextDone option.
    96  		wazero.NewRuntimeConfig().WithCloseOnContextDone(true))
    97  	defer r.Close(ctx)
    98  
    99  	moduleInstance, err := r.InstantiateWithConfig(ctx, infiniteLoopWasm,
   100  		wazero.NewModuleConfig().WithName("malicious_wasm"))
   101  	if err != nil {
   102  		log.Panicln(err)
   103  	}
   104  
   105  	infiniteLoop := moduleInstance.ExportedFunction("infinite_loop")
   106  
   107  	go func() {
   108  		// After 2 seconds, close the module instance with CloseWithExitCode, which triggers the termination
   109  		// of infinite loop.
   110  		time.Sleep(2 * time.Second)
   111  		if err := moduleInstance.CloseWithExitCode(ctx, 1); err != nil {
   112  			log.Panicln(err)
   113  		}
   114  	}()
   115  
   116  	// Invoke the infinite loop with the timeout context.
   117  	_, err = infiniteLoop.Call(ctx)
   118  
   119  	// The exit code is correctly handled and triggers the termination of infinite loop.
   120  	fmt.Println(err)
   121  
   122  	// Output:
   123  	//	module closed with exit_code(1)
   124  }