REesultsBuilder waits for results goroutine to finish before closing output channel

Author: gammazero

query/query.go

@@ -3,6 +3,7 @@ package query
 import (
 	"context"
 	"fmt"
+	"sync"
 	"time"
 )
 
@@ -149,6 +150,7 @@ type Results interface {
 	NextSync() (Result, bool) // blocks and waits to return the next result, second parameter returns false when results are exhausted
 	Rest() ([]Entry, error)   // waits till processing finishes, returns all entries at once.
 	Close()                   // client may call Close to signal early exit
+	Done() <-chan struct{}    // signals that Results is closed
 }
 
 // results implements Results
@@ -189,6 +191,10 @@ func (r *results) Query() Query {
 	return r.query
 }
 
+func (r *results) Done() <-chan struct{} {
+	return r.ctx.Done()
+}
+
 // ResultBuilder is what implementors use to construct results
 // Implementors of datastores and their clients must respect the
 // Process of the Request:
@@ -205,6 +211,7 @@ type ResultBuilder struct {
 
 	ctx    context.Context
 	cancel context.CancelFunc
+	wg     sync.WaitGroup
 }
 
 // Results returns a Results to to this builder.
@@ -232,6 +239,7 @@ func NewResultBuilder(q Query) *ResultBuilder {
 	}
 	b.ctx, b.cancel = context.WithCancel(context.Background())
 	context.AfterFunc(b.ctx, func() {
+		b.wg.Wait()
 		close(b.Output)
 	})
 	return b
@@ -240,8 +248,11 @@ func NewResultBuilder(q Query) *ResultBuilder {
 // ResultsWithChan returns a Results object from a channel
 // of Result entries.
 //
-// DEPRECATED: This iterator is impossible to cancel correctly. Canceling it
-// will leave anything trying to write to the result channel hanging.
+// DEPRECATED: This iterator takes sepcial care to cancel correctly. Canceling
+// it will leave anything trying to write to the result channel hanging, unless
+// that write can select the result channel and Results.Done(). This requires
+// creating the result channel, calline ResultsWithChan, and then writing to
+// the results channel.
 func ResultsWithChan(q Query, res <-chan Result) Results {
 	return ResultsWithContext(q, func(ctx context.Context, out chan<- Result) {
 		for {
@@ -263,14 +274,15 @@ 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.
+// ResultsWithContext 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)
-
+	b.wg.Add(1)
 	go func() {
-		defer b.cancel()
 		proc(b.ctx, b.Output)
+		b.cancel()
+		b.wg.Done()
 	}()
 
 	return b.Results()
@@ -381,6 +393,11 @@ func (r *resultsIter) Query() Query {
 	return r.query
 }
 
+func (r *resultsIter) Done() <-chan struct{} {
+	r.useLegacyResults()
+	return r.legacyResults.Done()
+}
+
 func (r *resultsIter) useLegacyResults() {
 	if r.legacyResults != nil {
 		return