wasm: bridge ctx cancellation once per eval, not per builtin call

The ctx->topdown.Cancel bridge goroutine was spawned per opa_builtinN
dispatch. Since the BuiltinContext (and its topdown.Cancel) is created
once per eval in Reset, a single bridge per eval suffices. Reset now
returns a stop func tying the goroutine's lifetime to the eval; the
per-dispatch goroutine is removed. A builtin-heavy policy now spawns one
cancellation goroutine per eval instead of one per builtin call.


Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

internal/wasm/sdk/internal/wasm/bindings.go

@@ -63,14 +63,18 @@ func (d *builtinDispatcher) SetMap(m map[int32]topdown.BuiltinFunc) {
 	d.builtins = m
 }
 
-// Reset is called in Eval before using the builtinDispatcher.
+// Reset is called in Eval before using the builtinDispatcher. It (re)builds the
+// BuiltinContext and starts a single goroutine that bridges context
+// cancellation into topdown.Cancel, so builtins that cooperate via topdown.Cancel
+// (e.g. net.cidr_expand) are aborted when ctx is done. The returned stop func
+// must be called when the eval completes to tear that goroutine down.
 func (d *builtinDispatcher) Reset(ctx context.Context,
 	seed io.Reader,
 	ns time.Time,
 	iqbCache cache.InterQueryCache,
 	ndbCache builtins.NDBCache,
 	ph print.Hook,
-	capabilities *ast.Capabilities) {
+	capabilities *ast.Capabilities) (stop func()) {
 	if ns.IsZero() {
 		ns = time.Now()
 	}
@@ -95,6 +99,19 @@ func (d *builtinDispatcher) Reset(ctx context.Context,
 		PrintHook:              ph,
 		Capabilities:           capabilities,
 	}
+
+	// WithCloseOnContextDone interrupts wasm-native execution, but it cannot
+	// preempt a running Go builtin; those cooperate via topdown.Cancel instead.
+	// A single bridge per eval suffices since d.ctx.Cancel lives for the eval.
+	done := make(chan struct{})
+	go func() {
+		select {
+		case <-done:
+		case <-ctx.Done():
+			d.ctx.Cancel.Cancel()
+		}
+	}()
+	return func() { close(done) }
 }
 
 func (d *builtinDispatcher) opaAbort(_ context.Context, addr int32) {
@@ -161,18 +178,6 @@ func (d *builtinDispatcher) dispatch(ctx context.Context, id int32, argAddrs ...
 		panic(abortError{message: "unreachable: uninitialized builtin dispatcher index"})
 	}
 
-	// Bridge context cancellation into topdown.Cancel so that builtins that
-	// use it (e.g. net.cidr_expand) are properly aborted.
-	done := make(chan struct{})
-	defer close(done)
-	go func() {
-		select {
-		case <-done:
-		case <-d.ctx.Context.Done():
-			d.ctx.Cancel.Cancel()
-		}
-	}()
-
 	convertedArgs := make([]*ast.Term, 0, len(argAddrs))
 	for _, addr := range argAddrs {
 		x, err := d.fromWasmValue(ctx, addr)

internal/wasm/sdk/internal/wasm/vm.go

@@ -291,7 +291,8 @@ func (i *VM) Eval(ctx context.Context,
 	// make use of it (e.g. `http.send`); and it will spawn a go routine
 	// cancelling the builtins that use topdown.Cancel, when the context is
 	// cancelled.
-	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
+	stop := i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
+	defer stop()
 
 	metrics.Timer("wasm_vm_eval_call").Start()
 	resultAddr, err := call(ctx, i, "opa_eval", 0 /* reserved */, entrypoint, i.dataAddr, inputAddr, inputLen, heapPtr, 1 /* value output */)
@@ -330,7 +331,8 @@ func (i *VM) evalCompat(ctx context.Context,
 
 	metrics.Timer("wasm_vm_eval_prepare_input").Start()
 
-	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
+	stop := i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
+	defer stop()
 
 	if err := i.setHeapState(ctx, i.evalHeapPtr); err != nil {
 		return nil, err