wa-lang.org/wazero@v1.0.2/experimental/listener_example_test.go (about)

     1  package experimental_test
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"fmt"
     7  	"log"
     8  	"sort"
     9  
    10  	"wa-lang.org/wazero"
    11  	"wa-lang.org/wazero/api"
    12  	. "wa-lang.org/wazero/experimental"
    13  	"wa-lang.org/wazero/imports/wasi_snapshot_preview1"
    14  )
    15  
    16  // listenerWasm was generated by the following:
    17  //
    18  //	cd testdata; wat2wasm --debug-names listener.wat
    19  //
    20  //go:embed logging/testdata/listener.wasm
    21  var listenerWasm []byte
    22  
    23  // uniqGoFuncs implements both FunctionListenerFactory and FunctionListener
    24  type uniqGoFuncs map[string]struct{}
    25  
    26  // callees returns the go functions called.
    27  func (u uniqGoFuncs) callees() []string {
    28  	ret := make([]string, 0, len(u))
    29  	for k := range u {
    30  		ret = append(ret, k)
    31  	}
    32  	// Sort names for consistent iteration
    33  	sort.Strings(ret)
    34  	return ret
    35  }
    36  
    37  // NewListener implements FunctionListenerFactory.NewListener
    38  func (u uniqGoFuncs) NewListener(def api.FunctionDefinition) FunctionListener {
    39  	if def.GoFunction() == nil {
    40  		return nil // only track go funcs
    41  	}
    42  	return u
    43  }
    44  
    45  // Before implements FunctionListener.Before
    46  func (u uniqGoFuncs) Before(ctx context.Context, def api.FunctionDefinition, _ []uint64) context.Context {
    47  	u[def.DebugName()] = struct{}{}
    48  	return ctx
    49  }
    50  
    51  // After implements FunctionListener.After
    52  func (u uniqGoFuncs) After(context.Context, api.FunctionDefinition, error, []uint64) {}
    53  
    54  // This shows how to make a listener that counts go function calls.
    55  func Example_customListenerFactory() {
    56  	u := uniqGoFuncs{}
    57  
    58  	// Set context to one that has an experimental listener
    59  	ctx := context.WithValue(context.Background(), FunctionListenerFactoryKey{}, u)
    60  
    61  	r := wazero.NewRuntime(ctx)
    62  	defer r.Close(ctx) // This closes everything this Runtime created.
    63  
    64  	wasi_snapshot_preview1.MustInstantiate(ctx, r)
    65  
    66  	// Compile the WebAssembly module using the default configuration.
    67  	code, err := r.CompileModule(ctx, listenerWasm)
    68  	if err != nil {
    69  		log.Panicln(err)
    70  	}
    71  
    72  	mod, err := r.InstantiateModule(ctx, code, wazero.NewModuleConfig())
    73  	if err != nil {
    74  		log.Panicln(err)
    75  	}
    76  
    77  	for i := 0; i < 5; i++ {
    78  		if _, err = mod.ExportedFunction("rand").Call(ctx, 4); err != nil {
    79  			log.Panicln(err)
    80  		}
    81  	}
    82  
    83  	// A Go function was called multiple times, but we should only see it once.
    84  	for _, f := range u.callees() {
    85  		fmt.Println(f)
    86  	}
    87  
    88  	// Output:
    89  	// wasi_snapshot_preview1.random_get
    90  }