From a987873218e728dc8cf3f2ebf33740ac3f732736 Mon Sep 17 00:00:00 2001 From: Hammad Majid Date: Sat, 5 Sep 2026 13:07:05 +0500 Subject: [PATCH] fix(aggregator): drain resultsCh on early return and stop polling a closed errCh AggregateChannel had three early returns (context cancellation, an error on errCh, and a fatal res.Error) that all walked away from a still-running producer with resultsCh open and undrained. Producers stream over a bounded channel, so once the buffer fills every worker blocks on its send, the task dispatcher blocks behind them, wg.Wait never returns and neither channel is ever closed: the worker pool, its dispatcher and the object store they pin leak for the life of the process, and a later range over the same channel hangs too. A named-result defer now drains resultsCh whenever the function exits with an error or a cancelled context, so the producer always gets the receives it needs to run itself down. The drain is skipped when the loop already observed the close, since ranging over the nilled channel would block forever. A receive on a closed channel is ready forever, so the exhausted errCh arm re-armed itself on every pass and spun the select at full CPU until the producer caught up (measured: ~214k passes across 100ms of streaming). Exhausted arms are now disabled by nilling their channel, which is what the pre-existing "if errCh != nil" guard had always intended. Finalization no longer depends on an undocumented ordering contract. The old post-close check was a non-blocking peek, so a producer that closed resultsCh before publishing its terminal error had that error silently swallowed and the search reported as successful. Aggregation now disables the results arm on close and keeps looping until errCh is closed as well, so a late error is surfaced. The producer and caller contracts this relies on - close both channels, share a cancellable context - are documented on AggregateChannel. Tests: a real search.Pipeline is driven into both early-return paths and the goroutine count is polled back to baseline; a custom context counts select passes to prove the closed errCh is no longer polled; a producer that closes resultsCh before publishing its error must still surface it. The existing channel tests now close both channels, as the documented producer contract requires. Closes #15 --- internal/aggregator/aggregator.go | 92 +++++--- internal/aggregator/aggregator_test.go | 294 ++++++++++++++++++++++++- 2 files changed, 353 insertions(+), 33 deletions(-) diff --git a/internal/aggregator/aggregator.go b/internal/aggregator/aggregator.go index 41bb1d3..3ebf1ef 100644 --- a/internal/aggregator/aggregator.go +++ b/internal/aggregator/aggregator.go @@ -90,54 +90,90 @@ func (a *Aggregator) Aggregate(results []*search.BlobResult) *AggregatedResults } // AggregateChannel organizes streamed BlobResult items into structured, ordered FileMatches. -// It terminates cleanly when resultsCh is closed, an error is received on errCh, or ctx is cancelled. -func (a *Aggregator) AggregateChannel(ctx context.Context, resultsCh <-chan *search.BlobResult, errCh <-chan error) (*AggregatedResults, error) { +// +// It consumes resultsCh and errCh until both are closed and then returns the aggregated +// results. It returns early, with a nil result, when errCh yields a non-nil error, when a +// result carries a fatal error, or when ctx is cancelled. +// +// The producer owns both channels and must close both when it stops. It may publish its +// terminal error before or after closing resultsCh: aggregation finalizes only once errCh +// is closed, so a late error is still surfaced rather than silently dropped. +// +// The caller must pass a context that the producer also observes, and that is cancelled +// when the caller loses interest. On every early return this function drains resultsCh so +// the producer is never left blocked on a send, and that drain only terminates once the +// producer closes resultsCh; a shared cancellable context is what guarantees the producer +// gets there. +func (a *Aggregator) AggregateChannel(ctx context.Context, resultsCh <-chan *search.BlobResult, errCh <-chan error) (out *AggregatedResults, err error) { if ctx == nil { ctx = context.Background() } + // Every early return below walks away from a producer that is still sending. Producers + // use a bounded results channel, so once its buffer fills every worker blocks on its + // send, the task dispatcher blocks behind them, and the object store they pin is never + // released - the producer wedges permanently and never closes its channels. Draining + // hands the producer the receives it is waiting for so it can run itself down. + // + // Termination: a nil resultsCh means the loop below already observed the close and + // consumed the channel to completion, so there is nothing left to unblock - and + // ranging over a nil channel would block forever, so return instead. Otherwise every + // receive advances the producer, so the drain ends as soon as the producer closes + // resultsCh. A producer honouring the contract above always gets there, whether it + // finishes its work or unwinds because the shared context was cancelled. + defer func() { + if resultsCh == nil || (err == nil && ctx.Err() == nil) { + return + } + for range resultsCh { + } + }() + var fileOrder []string fileMap := make(map[string]*fileBuilder) - for { + // A receive on a closed channel is ready forever, so an exhausted arm must be disabled + // by nilling its channel: a nil channel is never ready and its arm is never chosen + // again. Leaving a closed channel in place would make the select spin on it at full + // CPU. Both channels closed means the producer is done and the results are complete. + for resultsCh != nil || errCh != nil { select { case <-ctx.Done(): return nil, ctx.Err() - case err, ok := <-errCh: - if ok && err != nil { - return nil, err + case perr, ok := <-errCh: + if !ok { + errCh = nil + continue + } + if perr != nil { + return nil, perr } case res, ok := <-resultsCh: if !ok { - // Channel closed: verify if any pending error remains on errCh - if errCh != nil { - select { - case err, ok := <-errCh: - if ok && err != nil { - return nil, err - } - default: - } - } - return a.finalizeResults(fileMap, fileOrder), nil + resultsCh = nil + continue } - if res != nil { - if res.Error != nil { - // A blob the pipeline could not read is a soft failure: it carries no - // matches and is skipped so the remaining blobs still aggregate. - var bre *search.BlobReadError - if errors.As(res.Error, &bre) { - continue - } - return nil, res.Error + if res == nil { + continue + } + if res.Error != nil { + // A blob the pipeline could not read is a soft failure: it carries no + // matches and is skipped so the remaining blobs still aggregate. + var bre *search.BlobReadError + if errors.As(res.Error, &bre) { + continue } - a.processBlobResult(res, fileMap, &fileOrder) + return nil, res.Error } + a.processBlobResult(res, fileMap, &fileOrder) } } + + return a.finalizeResults(fileMap, fileOrder), nil } -// AggregateStream is an alias to AggregateChannel for streaming API flexibility. +// AggregateStream is an alias to AggregateChannel for streaming API flexibility. The +// producer and caller contracts documented on AggregateChannel apply unchanged. func (a *Aggregator) AggregateStream(ctx context.Context, resultsCh <-chan *search.BlobResult, errCh <-chan error) (*AggregatedResults, error) { return a.AggregateChannel(ctx, resultsCh, errCh) } diff --git a/internal/aggregator/aggregator_test.go b/internal/aggregator/aggregator_test.go index 02effec..8e99bde 100644 --- a/internal/aggregator/aggregator_test.go +++ b/internal/aggregator/aggregator_test.go @@ -3,6 +3,9 @@ package aggregator import ( "context" "errors" + "fmt" + "runtime" + "sync/atomic" "testing" "time" @@ -237,6 +240,9 @@ func TestAggregator_AggregateChannel_Success(t *testing.T) { resultsCh <- res1 resultsCh <- res2 close(resultsCh) + // A producer that has finished closes both of its channels; aggregation only + // finalizes once errCh is closed, so that a late error is not lost. + close(errCh) actual, err := agg.AggregateChannel(context.Background(), resultsCh, errCh) if err != nil { @@ -260,15 +266,21 @@ func TestAggregator_AggregateChannel_Cancellation(t *testing.T) { agg := New(&model.Config{}) resultsCh := make(chan *search.BlobResult) errCh := make(chan error) + // A producer that has already stopped streaming but never reports on errCh: + // cancellation is the only way out, and the closed results channel is what lets + // the drain on the way out terminate. + close(resultsCh) ctx, cancel := context.WithCancel(context.Background()) cancel() // pre-cancel - out, err := agg.AggregateChannel(ctx, resultsCh, errCh) + var out *AggregatedResults + var err error + runWithin(t, 5*time.Second, func() { out, err = agg.AggregateChannel(ctx, resultsCh, errCh) }) if out != nil { - t.Errorf("expected nil results on cancelled context") + t.Errorf("expected nil results on cancelled context, got %+v", out) } - if err != context.Canceled { + if !errors.Is(err, context.Canceled) { t.Errorf("expected context.Canceled, got %v", err) } } @@ -280,8 +292,11 @@ func TestAggregator_AggregateChannel_Error(t *testing.T) { expectedErr := errors.New("pipeline failed") errCh <- expectedErr + close(resultsCh) - out, err := agg.AggregateChannel(context.Background(), resultsCh, errCh) + var out *AggregatedResults + var err error + runWithin(t, 5*time.Second, func() { out, err = agg.AggregateChannel(context.Background(), resultsCh, errCh) }) if out != nil { t.Errorf("expected nil results on error") } @@ -297,8 +312,11 @@ func TestAggregator_AggregateChannel_BlobError(t *testing.T) { blobErr := errors.New("blob decompression error") resultsCh <- &search.BlobResult{Error: blobErr} + close(resultsCh) - out, err := agg.AggregateChannel(context.Background(), resultsCh, errCh) + var out *AggregatedResults + var err error + runWithin(t, 5*time.Second, func() { out, err = agg.AggregateChannel(context.Background(), resultsCh, errCh) }) if out != nil { t.Errorf("expected nil results on blob error") } @@ -331,6 +349,7 @@ func TestAggregator_AggregateChannel_SkipsBlobReadError(t *testing.T) { Occurrences: []model.BlobOccurrence{{Path: "b.txt", CommitSHA: "c1", CommitDate: t1}}, } close(resultsCh) + close(errCh) out, err := agg.AggregateChannel(context.Background(), resultsCh, errCh) if err != nil { @@ -345,3 +364,268 @@ func TestAggregator_AggregateChannel_SkipsBlobReadError(t *testing.T) { } } } + +// runWithin runs fn on another goroutine and fails if it has not returned by the +// deadline, so a regression that blocks forever surfaces as a test failure instead of +// hanging the whole binary. +func runWithin(t *testing.T, d time.Duration, fn func()) { + t.Helper() + + done := make(chan struct{}) + go func() { + defer close(done) + fn() + }() + + select { + case <-done: + case <-time.After(d): + t.Fatalf("AggregateChannel did not return within %s", d) + } +} + +// requireGoroutineBaseline polls the goroutine count back down to baseline, the way +// test/leak_test.go polls for leaks, and fails if a producer is still parked. +func requireGoroutineBaseline(t *testing.T, baseline int) { + t.Helper() + + deadline := time.Now().Add(5 * time.Second) + for { + n := runtime.NumGoroutine() + if n <= baseline { + return + } + if time.Now().After(deadline) { + t.Fatalf("goroutine count did not return to baseline %d (still %d): the producer is still blocked on an abandoned results channel", baseline, n) + } + time.Sleep(10 * time.Millisecond) + } +} + +// requireDrained asserts that resultsCh was consumed to close before AggregateChannel +// returned, i.e. that the producer was released synchronously rather than abandoned. +func requireDrained(t *testing.T, resultsCh <-chan *search.BlobResult) { + t.Helper() + + select { + case _, ok := <-resultsCh: + if ok { + t.Error("results channel still carries buffered results: producer was abandoned mid-stream") + } + default: + t.Error("results channel was neither drained nor closed before returning") + } +} + +// stubObjectReader serves blobs from memory so a real search.Pipeline can be driven +// without a repository on disk. +type stubObjectReader struct { + objects map[string]*gitengine.Object +} + +func (s *stubObjectReader) ReadObject(oid string) (*gitengine.Object, error) { + if obj, ok := s.objects[oid]; ok { + return obj, nil + } + return nil, gitengine.ErrObjectNotFound +} + +func (s *stubObjectReader) HasObject(oid string) bool { + _, ok := s.objects[oid] + return ok +} + +func (s *stubObjectReader) Close() error { return nil } + +// newStreamingPipeline builds a real pipeline over n in-memory matching blobs. n is +// chosen by the caller to exceed the pipeline's bounded results buffer so its workers +// must block on their sends, which is what makes abandoning the channel fatal. +func newStreamingPipeline(t *testing.T, n int) (*search.Pipeline, []model.BlobOccurrence, *model.Config) { + t.Helper() + + objects := make(map[string]*gitengine.Object, n) + occurrences := make([]model.BlobOccurrence, 0, n) + commitDate := time.Date(2026, 1, 1, 10, 0, 0, 0, time.UTC) + + for i := range n { + oid := fmt.Sprintf("%040d", i) + data := fmt.Sprintf("first line\nNEEDLE_AGG %d\nlast line\n", i) + objects[oid] = &gitengine.Object{ + OID: oid, + Type: gitengine.TypeBlob, + Size: int64(len(data)), + Data: []byte(data), + } + occurrences = append(occurrences, model.BlobOccurrence{ + BlobOID: oid, + Path: fmt.Sprintf("dir/file_%d.txt", i), + CommitSHA: fmt.Sprintf("%040x", i), + CommitDate: commitDate, + }) + } + + cfg := &model.Config{Pattern: "NEEDLE_AGG", CaseSensitive: true} + matcher, err := search.NewMatcher(cfg) + if err != nil { + t.Fatalf("failed creating matcher: %v", err) + } + + return search.NewPipeline(&stubObjectReader{objects: objects}, matcher, cfg), occurrences, cfg +} + +// Returning early must not abandon the producer. A real pipeline blocks its workers on +// a bounded results channel, so a consumer that walks away without draining wedges the +// worker pool, its dispatcher, and the object store they hold forever (regression #15). +func TestAggregator_AggregateChannel_ReleasesProducerOnEarlyReturn(t *testing.T) { + const blobs = 256 // well above max(32, NumCPU*4), the pipeline's results buffer + + t.Run("producer error", func(t *testing.T) { + pipeline, occurrences, cfg := newStreamingPipeline(t, blobs) + baseline := runtime.NumGoroutine() + + // The pipeline keeps its own context: it is not cancelled when the consumer + // gives up, so only draining can unblock it. + resultsCh, _ := pipeline.ExecuteContext(context.Background(), occurrences) + + fatal := errors.New("repository closed under the search") + errCh := make(chan error, 1) + errCh <- fatal + close(errCh) + + var out *AggregatedResults + var err error + runWithin(t, 30*time.Second, func() { + out, err = New(cfg).AggregateChannel(context.Background(), resultsCh, errCh) + }) + + if !errors.Is(err, fatal) { + t.Fatalf("expected %v, got %v", fatal, err) + } + if out != nil { + t.Errorf("expected nil results alongside the error, got %+v", out) + } + requireDrained(t, resultsCh) + requireGoroutineBaseline(t, baseline) + }) + + t.Run("consumer cancellation", func(t *testing.T) { + pipeline, occurrences, cfg := newStreamingPipeline(t, blobs) + baseline := runtime.NumGoroutine() + + resultsCh, errCh := pipeline.ExecuteContext(context.Background(), occurrences) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var err error + runWithin(t, 30*time.Second, func() { + _, err = New(cfg).AggregateChannel(ctx, resultsCh, errCh) + }) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + requireDrained(t, resultsCh) + requireGoroutineBaseline(t, baseline) + }) +} + +// selectCountingContext counts how often its Done channel is requested. AggregateChannel +// evaluates ctx.Done() exactly once per pass through its select, which turns "the loop +// busy-polls an exhausted channel" into an observable number. +type selectCountingContext struct { + context.Context + done chan struct{} + passes atomic.Int64 +} + +func newSelectCountingContext() *selectCountingContext { + return &selectCountingContext{Context: context.Background(), done: make(chan struct{})} +} + +func (c *selectCountingContext) Done() <-chan struct{} { + c.passes.Add(1) + return c.done +} + +// A producer is free to close errCh before it has finished streaming results. A receive +// on a closed channel is ready forever, so the exhausted arm has to be disabled or the +// loop burns a core until the producer catches up (regression #15). +func TestAggregator_AggregateChannel_ClosedErrChIsNotPolled(t *testing.T) { + const ( + results = 5 + gap = 20 * time.Millisecond + ) + + resultsCh := make(chan *search.BlobResult) + errCh := make(chan error) + close(errCh) // producer reports "no failures" up front, then keeps streaming + + commitDate := time.Date(2026, 1, 1, 10, 0, 0, 0, time.UTC) + go func() { + defer close(resultsCh) + for i := range results { + time.Sleep(gap) + resultsCh <- &search.BlobResult{ + BlobOID: fmt.Sprintf("blob%d", i), + Matches: []model.SearchMatch{{LineNum: 1, LineText: "needle"}}, + Occurrences: []model.BlobOccurrence{{Path: fmt.Sprintf("f%d.txt", i), CommitSHA: "c1", CommitDate: commitDate}}, + } + } + }() + + ctx := newSelectCountingContext() + agg := New(&model.Config{}) + + var out *AggregatedResults + var err error + runWithin(t, 30*time.Second, func() { out, err = agg.AggregateChannel(ctx, resultsCh, errCh) }) + + if err != nil { + t.Fatalf("a closed error channel reports success, got: %v", err) + } + if out == nil || out.TotalFiles != results { + t.Fatalf("expected %d files aggregated after errCh closed, got %+v", results, out) + } + // One pass per delivered result plus a handful for the two closes; a spinning loop + // racks up millions over the same 100ms of streaming. + if passes := ctx.passes.Load(); passes > 1000 { + t.Errorf("select was re-entered %d times for %d results: the closed error channel is being polled", passes, results) + } +} + +// Nothing obliges a producer to publish its terminal error before closing resultsCh, so +// an error that arrives afterwards must still be surfaced rather than reported as a +// successful search (regression #15). +func TestAggregator_AggregateChannel_ErrorAfterResultsClosed(t *testing.T) { + agg := New(&model.Config{}) + resultsCh := make(chan *search.BlobResult, 1) + errCh := make(chan error) + + commitDate := time.Date(2026, 1, 1, 10, 0, 0, 0, time.UTC) + resultsCh <- &search.BlobResult{ + BlobOID: "blob1", + Matches: []model.SearchMatch{{LineNum: 1, LineText: "needle"}}, + Occurrences: []model.BlobOccurrence{{Path: "a.txt", CommitSHA: "c1", CommitDate: commitDate}}, + } + close(resultsCh) + + fatal := errors.New("history walk aborted") + go func() { + defer close(errCh) + // Ordered the other way round from the pipeline: results first, error second. + time.Sleep(20 * time.Millisecond) + errCh <- fatal + }() + + var out *AggregatedResults + var err error + runWithin(t, 30*time.Second, func() { out, err = agg.AggregateChannel(context.Background(), resultsCh, errCh) }) + + if !errors.Is(err, fatal) { + t.Fatalf("error published after resultsCh closed must be surfaced, got err=%v out=%+v", err, out) + } + if out != nil { + t.Errorf("expected nil results alongside the error, got %+v", out) + } +}