add support for re-entering subqueries in recursive functions (#6325)

This commit addresses an unfinished aspect of subqueries when they
are invoked from recursive functions.  Previously, the subquery would
incorrectly block as Eval() is called concurrently violating our invariant.

The fix is to keep a stack of non-shared subquery instances that can
be used when recursion is detected.  This involves calling the rungen
builder from the runtime whenever we need a new instance, which turned
out to be surprisingly easy to implement.

Author: mccanne

compiler/rungen/expr.go

@@ -398,13 +398,17 @@ func (b *Builder) compileSubquery(query *dag.SubqueryExpr) (expr.Evaluator, erro
 		}
 		return subquery.NewCachedSubquery(b.rctx, body), nil
 	}
-	subquery := subquery.NewSubquery(b.rctx)
-	body, err := b.compileSeqAndCombine(query.Body, []sbuf.Puller{subquery})
-	if err != nil {
-		return nil, err
+	var create func() *subquery.Subquery
+	create = func() *subquery.Subquery {
+		subquery := subquery.NewSubquery(b.rctx, create)
+		body, err := b.compileSeqAndCombine(query.Body, []sbuf.Puller{subquery})
+		if err != nil {
+			panic(err)
+		}
+		subquery.SetBody(body)
+		return subquery
 	}
-	subquery.SetBody(body)
-	return subquery, nil
+	return create(), nil
 }
 
 func (b *Builder) compileSeqAndCombine(seq dag.Seq, parents []sbuf.Puller) (sbuf.Puller, error) {

runtime/sam/op/subquery/subquery.go

@@ -51,14 +51,19 @@ type Subquery struct {
 	sctx    *super.Context
 	batchCh chan sbuf.Batch
 	eos     bool
+	tos     int
+	stack   []*Subquery
+	create  func() *Subquery
 
 	body sbuf.Puller
 }
 
-func NewSubquery(rctx *runtime.Context) *Subquery {
+func NewSubquery(rctx *runtime.Context, create func() *Subquery) *Subquery {
 	return &Subquery{
 		ctx:     rctx.Context,
 		sctx:    rctx.Sctx,
+		tos:     -2,
+		create:  create,
 		batchCh: make(chan sbuf.Batch, 1),
 	}
 }
@@ -81,22 +86,43 @@ func (s *Subquery) Pull(done bool) (sbuf.Batch, error) {
 	}
 }
 
-func (q *Subquery) Eval(this super.Value) super.Value {
+const MaxSubqueryRecursion = 10000
+
+func (s *Subquery) Eval(this super.Value) super.Value {
+	s.tos++
+	defer func() {
+		s.tos--
+	}()
+	if s.tos >= 0 {
+		// We're re-entering this subquery instance before it's done evaluating
+		// the previous invocation.  This happens when a subquery is invoked
+		// inside of a recursive function so the same instance ends up being
+		// called by different call frames.  To deal with this, we keep a stack
+		// of Subquery duplicates where each duplicate is not shared and extend
+		// the stack as needed.  If the stack overflows, we return an error.
+		if s.tos >= MaxSubqueryRecursion {
+			return s.sctx.WrapError("subquery recursion depth exceeded", this)
+		}
+		if s.tos >= len(s.stack) {
+			s.stack = append(s.stack, s.create())
+		}
+		return s.stack[s.tos].Eval(this)
+	}
 	select {
-	case q.batchCh <- sbuf.NewArray([]super.Value{this}):
-	case <-q.ctx.Done():
-		return q.sctx.NewError(q.ctx.Err())
+	case s.batchCh <- sbuf.NewArray([]super.Value{this}):
+	case <-s.ctx.Done():
+		return s.sctx.NewError(s.ctx.Err())
 	}
 	val := super.Null
 	var count int
 	for {
-		b, err := q.body.Pull(false)
+		b, err := s.body.Pull(false)
 		if err != nil {
 			panic(err)
 		}
 		if b == nil {
 			if count > 1 {
-				return q.sctx.NewErrorf("query expression produced multiple values (consider [subquery])")
+				return s.sctx.NewErrorf("query expression produced multiple values (consider [subquery])")
 			}
 			return val
 		}

runtime/sam/op/ztests/recursive-subquery.yaml

@@ -0,0 +1,13 @@
+spq: |
+  fn walk(v):
+    case kind(v)
+    when "array" then 
+      [unnest v | walk(this)]
+    when "record" then 
+      unflatten([unnest flatten(v) | {key,value:walk(value)}])
+    else v+1
+    end
+  values walk([{x:[1,2]},{y:3}])
+
+output: |
+  [{x:[2,3]},{y:4}]