ResultBuilder and ResultsWithContext take a parent context to cancel results

Author: gammazero

query/query.go

@@ -221,7 +221,7 @@ func (rb *ResultBuilder) Results() Results {
 const NormalBufSize = 1
 const KeysOnlyBufSize = 128
 
-func NewResultBuilder(q Query) *ResultBuilder {
+func NewResultBuilder(ctx context.Context, q Query) *ResultBuilder {
 	bufSize := NormalBufSize
 	if q.KeysOnly {
 		bufSize = KeysOnlyBufSize
@@ -230,7 +230,7 @@ func NewResultBuilder(q Query) *ResultBuilder {
 		Query:  q,
 		Output: make(chan Result, bufSize),
 	}
-	b.ctx, b.cancel = context.WithCancel(context.Background())
+	b.ctx, b.cancel = context.WithCancel(ctx)
 	context.AfterFunc(b.ctx, func() {
 		close(b.Output)
 	})
@@ -243,7 +243,7 @@ func NewResultBuilder(q Query) *ResultBuilder {
 // DEPRECATED: This iterator is impossible to cancel correctly. Canceling it
 // will leave anything trying to write to the result channel hanging.
 func ResultsWithChan(q Query, res <-chan Result) Results {
-	return ResultsWithContext(q, func(ctx context.Context, out chan<- Result) {
+	return ResultsWithContext(context.Background(), q, func(ctx context.Context, out chan<- Result) {
 		for {
 			select {
 			case <-ctx.Done(): // client told us to close early
@@ -263,10 +263,12 @@ func ResultsWithChan(q Query, res <-chan Result) Results {
 	})
 }
 
-// ResultsWithCtxs returns a Results object with the results generated by the
-// passed proc function called in a separate goroutine.
-func ResultsWithContext(q Query, proc func(context.Context, chan<- Result)) Results {
-	b := NewResultBuilder(q)
+// ResultsWithContext returns a Results object with the results generated by
+// the passed proc function called in a separate goroutine. The context passed
+// into the proc function is cancelled when the returned Results is closed, or
+// when the parent context, ctx, is canceled.
+func ResultsWithContext(ctx context.Context, q Query, proc func(context.Context, chan<- Result)) Results {
+	b := NewResultBuilder(ctx, q)
 
 	go func() {
 		defer b.cancel()
@@ -386,7 +388,7 @@ func (r *resultsIter) useLegacyResults() {
 		return
 	}
 
-	b := NewResultBuilder(r.query)
+	b := NewResultBuilder(context.Background(), r.query)
 
 	// go consume all the entries and add them to the results.
 	go func(ctx context.Context, cancel context.CancelFunc, out chan<- Result) {