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

     1  package experimental_test
     2  
     3  import (
     4  	"context"
     5  	_ "embed"
     6  	"testing"
     7  
     8  	"wa-lang.org/wazero"
     9  	"wa-lang.org/wazero/api"
    10  	. "wa-lang.org/wazero/experimental"
    11  	"wa-lang.org/wazero/internal/testing/require"
    12  	"wa-lang.org/wazero/internal/wasm"
    13  	"wa-lang.org/wazero/internal/wasm/binary"
    14  )
    15  
    16  // compile-time check to ensure recorder implements FunctionListenerFactory
    17  var _ FunctionListenerFactory = &recorder{}
    18  
    19  type recorder struct {
    20  	m                       map[string]struct{}
    21  	beforeNames, afterNames []string
    22  }
    23  
    24  func (r *recorder) Before(ctx context.Context, def api.FunctionDefinition, _ []uint64) context.Context {
    25  	r.beforeNames = append(r.beforeNames, def.DebugName())
    26  	return ctx
    27  }
    28  
    29  func (r *recorder) After(_ context.Context, def api.FunctionDefinition, _ error, _ []uint64) {
    30  	r.afterNames = append(r.afterNames, def.DebugName())
    31  }
    32  
    33  func (r *recorder) NewListener(definition api.FunctionDefinition) FunctionListener {
    34  	r.m[definition.Name()] = struct{}{}
    35  	return r
    36  }
    37  
    38  func TestFunctionListenerFactory(t *testing.T) {
    39  	// Set context to one that has an experimental listener
    40  	factory := &recorder{m: map[string]struct{}{}}
    41  	ctx := context.WithValue(context.Background(), FunctionListenerFactoryKey{}, factory)
    42  
    43  	// Define a module with two functions
    44  	bin := binary.EncodeModule(&wasm.Module{
    45  		TypeSection:     []*wasm.FunctionType{{}},
    46  		ImportSection:   []*wasm.Import{{}},
    47  		FunctionSection: []wasm.Index{0, 0},
    48  		CodeSection: []*wasm.Code{
    49  			// fn1
    50  			{Body: []byte{
    51  				// call fn2 twice
    52  				wasm.OpcodeCall, 2,
    53  				wasm.OpcodeCall, 2,
    54  				wasm.OpcodeEnd,
    55  			}},
    56  			// fn2
    57  			{Body: []byte{wasm.OpcodeEnd}},
    58  		},
    59  		ExportSection: []*wasm.Export{{Name: "fn1", Type: wasm.ExternTypeFunc, Index: 1}},
    60  		NameSection: &wasm.NameSection{
    61  			ModuleName: "test",
    62  			FunctionNames: wasm.NameMap{
    63  				{Index: 0, Name: "import"}, // should skip for building listeners.
    64  				{Index: 1, Name: "fn1"},
    65  				{Index: 2, Name: "fn2"},
    66  			},
    67  		},
    68  	})
    69  
    70  	r := wazero.NewRuntime(ctx)
    71  	defer r.Close(ctx) // This closes everything this Runtime created.
    72  
    73  	_, err := r.NewHostModuleBuilder("").NewFunctionBuilder().WithFunc(func() {}).Export("").Instantiate(ctx, r)
    74  	require.NoError(t, err)
    75  
    76  	// Ensure the imported function was converted to a listener.
    77  	require.Equal(t, map[string]struct{}{"": {}}, factory.m)
    78  
    79  	compiled, err := r.CompileModule(ctx, bin)
    80  	require.NoError(t, err)
    81  
    82  	// Ensure each function was converted to a listener eagerly
    83  	require.Equal(t, map[string]struct{}{
    84  		"":    {},
    85  		"fn1": {},
    86  		"fn2": {},
    87  	}, factory.m)
    88  
    89  	// Ensures that FunctionListener is a compile-time option, so passing context.Background here
    90  	// is ok to use listeners at runtime.
    91  	m, err := r.InstantiateModule(context.Background(), compiled, wazero.NewModuleConfig())
    92  	require.NoError(t, err)
    93  
    94  	fn1 := m.ExportedFunction("fn1")
    95  	require.NotNil(t, fn1)
    96  
    97  	_, err = fn1.Call(context.Background())
    98  	require.NoError(t, err)
    99  
   100  	require.Equal(t, []string{"test.fn1", "test.fn2", "test.fn2"}, factory.beforeNames)
   101  	require.Equal(t, []string{"test.fn2", "test.fn2", "test.fn1"}, factory.afterNames) // after is in the reverse order.
   102  }