rego: expose QueryTracers, tracing.Options and Cancel from QueryContext

QueryTracers is required for parameterized subtests to work in a
different rego plugin.

tracing.Options are needed to have http.send and friends be wired up
with OTel when using a different rego plugin.

Cancellation is useful when the evaluation scenario is different from
the usual, like in bulk requests

Co-authored-by: Philip Conrad <philip@chariot-chaser.net>
Signed-off-by: Stephan Renatus <stephan@styra.com>

v1/rego/rego.go

@@ -126,6 +126,8 @@ type EvalContext struct {
 	strictBuiltinErrors         bool
 	virtualCache                topdown.VirtualCache
 	baseCache                   topdown.BaseCache
+	tracing                     tracing.Options
+	externalCancel              topdown.Cancel // Note(philip): If non-nil, the cancellation is handled outside of this package.
 }
 
 func (e *EvalContext) RawInput() *any {
@@ -180,6 +182,18 @@ func (e *EvalContext) Transaction() storage.Transaction {
 	return e.txn
 }
 
+func (e *EvalContext) TracingOpts() tracing.Options {
+	return e.tracing
+}
+
+func (e *EvalContext) ExternalCancel() topdown.Cancel {
+	return e.externalCancel
+}
+
+func (e *EvalContext) QueryTracers() []topdown.QueryTracer {
+	return e.queryTracers
+}
+
 // EvalOption defines a function to set an option on an EvalConfig
 type EvalOption func(*EvalContext)
 
@@ -388,6 +402,14 @@ func EvalNondeterministicBuiltins(yes bool) EvalOption {
 	}
 }
 
+// EvalExternalCancel sets an external topdown.Cancel for the interpreter to use
+// for cancellation. This is useful for batch-evaluation of many rego queries.
+func EvalExternalCancel(ec topdown.Cancel) EvalOption {
+	return func(e *EvalContext) {
+		e.externalCancel = ec
+	}
+}
+
 func (pq preparedQuery) Modules() map[string]*ast.Module {
 	mods := make(map[string]*ast.Module)
 
@@ -2214,13 +2236,19 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
 	}
 
 	// Cancel query if context is cancelled or deadline is reached.
-	c := topdown.NewCancel()
-	q = q.WithCancel(c)
-	exit := make(chan struct{})
-	defer close(exit)
-	go waitForDone(ctx, exit, func() {
-		c.Cancel()
-	})
+	if ectx.externalCancel == nil {
+		// Create a one-off goroutine to handle cancellation for this query.
+		c := topdown.NewCancel()
+		q = q.WithCancel(c)
+		exit := make(chan struct{})
+		defer close(exit)
+		go waitForDone(ctx, exit, func() {
+			c.Cancel()
+		})
+	} else {
+		// Query cancellation is being handled elsewhere.
+		q = q.WithCancel(ectx.externalCancel)
+	}
 
 	var rs ResultSet
 	err := q.Iter(ctx, func(qr topdown.QueryResult) error {
@@ -2502,13 +2530,19 @@ func (r *Rego) partial(ctx context.Context, ectx *EvalContext) (*PartialQueries,
 	}
 
 	// Cancel query if context is cancelled or deadline is reached.
-	c := topdown.NewCancel()
-	q = q.WithCancel(c)
-	exit := make(chan struct{})
-	defer close(exit)
-	go waitForDone(ctx, exit, func() {
-		c.Cancel()
-	})
+	if ectx.externalCancel == nil {
+		// Create a one-off goroutine to handle cancellation for this query.
+		c := topdown.NewCancel()
+		q = q.WithCancel(c)
+		exit := make(chan struct{})
+		defer close(exit)
+		go waitForDone(ctx, exit, func() {
+			c.Cancel()
+		})
+	} else {
+		// Query cancellation is being handled elsewhere.
+		q = q.WithCancel(ectx.externalCancel)
+	}
 
 	queries, support, err := q.PartialRun(ctx)
 	if err != nil {