From e485a09d0da682dfa450527fa304f3bc8dd7b724 Mon Sep 17 00:00:00 2001 From: Hammad Majid Date: Sat, 5 Sep 2026 13:04:35 +0500 Subject: [PATCH] fix(output): make rendering cancellable mid-stream Formatter.Format took no context, so after emitResults' single pre-flight ctx.Err() check a multi-million-line result set streamed to stdout with no way to stop it. If stdout was a slow pipe, Ctrl-C was inert for the entire render. Add a leading ctx to the Formatter interface and thread it through all four concrete formatters (grouped, single-line, count, files-with-matches). Each consults the context through a cancelGuard, which reads it at file-group boundaries and once every 1024 emitted lines; the per-line cost is a decrement and a branch, keeping the atomic load off the innermost write path. Checks sit between line writes, so the output produced before an abort is always a whole-line prefix of the complete rendering. In emitResults the now-redundant pre-flight check is folded into the render itself, and a cancellation return maps to cancelError using the same predicate as the walker call site above it. Also updates test/leak_test.go's single Format call site to the new signature (agreed with the lead; no agent owns that file in this wave). Verified end-to-end: rendering 34962 lines into a slow pipe, SIGINT is ignored for >5s before this change and aborts after 0.32s with only 285 lines emitted after it. Closes #17 --- cmd/grg/main.go | 11 +- internal/output/count.go | 14 ++- internal/output/files_with_matches.go | 14 ++- internal/output/grouped.go | 17 ++- internal/output/output_test.go | 149 ++++++++++++++++++++++++-- internal/output/single.go | 17 ++- internal/output/writer.go | 45 +++++++- test/leak_test.go | 2 +- 8 files changed, 246 insertions(+), 23 deletions(-) diff --git a/cmd/grg/main.go b/cmd/grg/main.go index 7a509dc..c7c456f 100644 --- a/cmd/grg/main.go +++ b/cmd/grg/main.go @@ -202,12 +202,13 @@ func emitResults(ctx context.Context, cfg *model.Config, results []*search.BlobR return noMatchError{} } - if err := ctx.Err(); err != nil { - return cancelError{err: err, quiet: cfg.Quiet} - } - + // Rendering is cancellable: Format observes ctx at coarse boundaries, so the + // pre-flight check is folded into the render itself. formatter := output.NewFormatter(cfg) - if err := formatter.Format(stdout, aggregated); err != nil { + if err := formatter.Format(ctx, stdout, aggregated); err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || ctx.Err() != nil { + return cancelError{err: err, quiet: cfg.Quiet} + } return err } diff --git a/internal/output/count.go b/internal/output/count.go index 354b25f..f81a32a 100644 --- a/internal/output/count.go +++ b/internal/output/count.go @@ -1,6 +1,7 @@ package output import ( + "context" "fmt" "io" "strconv" @@ -24,14 +25,20 @@ func NewCountFormatter(cfg *model.Config) *CountFormatter { } } -// Format writes :: lines to w. -func (f *CountFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error { +// Format writes :: lines to w, aborting with +// ctx.Err() if ctx is cancelled mid-render. +func (f *CountFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error { if results == nil || len(results.Files) == 0 { return nil } c := f.color + guard := newCancelGuard(ctx) for _, file := range results.Files { + if err := guard.boundary(); err != nil { + return err + } + for _, commit := range file.Commits { count := len(commit.Matches) if commit.IsBinary { @@ -57,6 +64,9 @@ func (f *CountFormatter) Format(w io.Writer, results *aggregator.AggregatedResul formattedCount); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } diff --git a/internal/output/files_with_matches.go b/internal/output/files_with_matches.go index 27bcddc..ecceae1 100644 --- a/internal/output/files_with_matches.go +++ b/internal/output/files_with_matches.go @@ -1,6 +1,7 @@ package output import ( + "context" "fmt" "io" @@ -23,16 +24,22 @@ func NewFilesWithMatchesFormatter(cfg *model.Config) *FilesWithMatchesFormatter } } -// Format writes distinct : lines to w. -func (f *FilesWithMatchesFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error { +// Format writes distinct : lines to w, aborting with +// ctx.Err() if ctx is cancelled mid-render. +func (f *FilesWithMatchesFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error { if results == nil || len(results.Files) == 0 { return nil } c := f.color + guard := newCancelGuard(ctx) seen := make(map[string]bool) for _, file := range results.Files { + if err := guard.boundary(); err != nil { + return err + } + for _, commit := range file.Commits { if len(commit.Matches) == 0 && !commit.IsBinary { continue @@ -56,6 +63,9 @@ func (f *FilesWithMatchesFormatter) Format(w io.Writer, results *aggregator.Aggr if _, err := fmt.Fprintf(w, "%s%s%s\n", formattedCommit, formattedSep, formattedPath); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } diff --git a/internal/output/grouped.go b/internal/output/grouped.go index 72b21ca..070a901 100644 --- a/internal/output/grouped.go +++ b/internal/output/grouped.go @@ -1,6 +1,7 @@ package output import ( + "context" "fmt" "io" "strconv" @@ -24,16 +25,22 @@ func NewGroupedFormatter(cfg *model.Config) *GroupedFormatter { } } -// Format writes the grouped results to w. -func (g *GroupedFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error { +// Format writes the grouped results to w, aborting with ctx.Err() if ctx is +// cancelled mid-render. +func (g *GroupedFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error { if results == nil || len(results.Files) == 0 { return nil } c := g.color + guard := newCancelGuard(ctx) firstFile := true for _, file := range results.Files { + if err := guard.boundary(); err != nil { + return err + } + if !firstFile { if _, err := fmt.Fprintln(w); err != nil { return err @@ -71,6 +78,9 @@ func (g *GroupedFormatter) Format(w io.Writer, results *aggregator.AggregatedRes if err := g.writeLine(w, line.LineNum, line.LineText, line.IsMatch, line.Submatches); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } } else { @@ -78,6 +88,9 @@ func (g *GroupedFormatter) Format(w io.Writer, results *aggregator.AggregatedRes if err := g.writeLine(w, match.LineNum, match.LineText, true, match.Submatches); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } } diff --git a/internal/output/output_test.go b/internal/output/output_test.go index 6cb70c0..59261b6 100644 --- a/internal/output/output_test.go +++ b/internal/output/output_test.go @@ -2,6 +2,9 @@ package output import ( "bytes" + "context" + "errors" + "fmt" "strings" "testing" "time" @@ -60,7 +63,7 @@ func TestGroupedFormatter(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -84,7 +87,7 @@ func TestGroupedFormatter_NoLineNumber(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -140,7 +143,7 @@ func TestGroupedFormatter_ContextLines(t *testing.T) { fmtter := NewFormatter(cfg) var buf bytes.Buffer - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -167,7 +170,7 @@ func TestSingleLineFormatter(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -189,7 +192,7 @@ func TestSingleLineFormatter_NoLineNumber(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -210,7 +213,7 @@ func TestFilesWithMatchesFormatter(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -229,7 +232,7 @@ func TestCountFormatter(t *testing.T) { var buf bytes.Buffer results := makeTestAggregatedResults() - if err := fmtter.Format(&buf, results); err != nil { + if err := fmtter.Format(context.Background(), &buf, results); err != nil { t.Fatalf("unexpected error: %v", err) } @@ -270,7 +273,7 @@ func TestBinaryFormatter(t *testing.T) { Color: model.ColorNever, } var buf bytes.Buffer - if err := NewFormatter(cfg).Format(&buf, results); err != nil { + if err := NewFormatter(cfg).Format(context.Background(), &buf, results); err != nil { t.Fatal(err) } @@ -320,3 +323,133 @@ func TestColorHighlighting(t *testing.T) { t.Errorf("expected %q, got %q", expected, highlighted) } } + +// makeLargeAggregatedResults builds a result set whose rendering spans several +// cancelCheckInterval windows, so that a mid-render cancellation is observable. +func makeLargeAggregatedResults(files, matchesPerFile int) *aggregator.AggregatedResults { + date := time.Date(2026, 9, 4, 12, 0, 0, 0, time.UTC) + res := &aggregator.AggregatedResults{TotalFiles: files} + for f := range files { + matches := make([]model.SearchMatch, matchesPerFile) + for m := range matches { + matches[m] = model.SearchMatch{ + LineNum: m + 1, + LineText: fmt.Sprintf("needle in file %d line %d", f, m+1), + } + } + res.Files = append(res.Files, aggregator.FileMatches{ + Path: fmt.Sprintf("pkg/file%04d.go", f), + Commits: []aggregator.CommitMatches{{ + CommitSHA: fmt.Sprintf("%040x", f), + ShortSHA: fmt.Sprintf("%07x", f), + CommitDate: date, + Author: "Alice", + AuthorName: "Alice", + Summary: "bulk commit", + Matches: matches, + }}, + }) + res.TotalMatches += matchesPerFile + } + return res +} + +// cancelAfterWriter cancels the render's context once after lines have been +// written, and keeps recording whatever the formatter emits afterwards so the +// test can prove the remainder was never produced. +type cancelAfterWriter struct { + buf bytes.Buffer + cancel context.CancelFunc + after int + lines int +} + +func (w *cancelAfterWriter) Write(p []byte) (int, error) { + n, err := w.buf.Write(p) + w.lines += bytes.Count(p, []byte("\n")) + if w.cancel != nil && w.lines >= w.after { + w.cancel() + w.cancel = nil + } + return n, err +} + +func TestFormatterCancelledMidRender(t *testing.T) { + tests := []struct { + name string + cfg *model.Config + results *aggregator.AggregatedResults + }{ + { + name: "grouped", + cfg: &model.Config{Heading: true, LineNumber: true, Color: model.ColorNever}, + results: makeLargeAggregatedResults(8, 512), + }, + { + name: "single", + cfg: &model.Config{Heading: false, LineNumber: true, Color: model.ColorNever}, + results: makeLargeAggregatedResults(8, 512), + }, + { + name: "count", + cfg: &model.Config{Count: true, Color: model.ColorNever}, + results: makeLargeAggregatedResults(3000, 1), + }, + { + name: "files-with-matches", + cfg: &model.Config{FilesWithMatches: true, Color: model.ColorNever}, + results: makeLargeAggregatedResults(3000, 1), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var full bytes.Buffer + if err := NewFormatter(tc.cfg).Format(context.Background(), &full, tc.results); err != nil { + t.Fatalf("uncancelled render failed: %v", err) + } + want := full.String() + if lines := strings.Count(want, "\n"); lines <= 2*cancelCheckInterval { + t.Fatalf("fixture renders %d lines, too few to span several check intervals", lines) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + w := &cancelAfterWriter{cancel: cancel, after: 10} + + err := NewFormatter(tc.cfg).Format(ctx, w, tc.results) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + + got := w.buf.String() + if got == "" { + t.Fatal("expected the output written before cancellation to be retained") + } + if len(got) >= len(want) { + t.Fatalf("cancellation emitted %d of %d bytes: the remainder was not skipped", len(got), len(want)) + } + if !strings.HasPrefix(want, got) { + t.Fatalf("output after cancellation is not a prefix of the full render (%d bytes written)", len(got)) + } + if !strings.HasSuffix(got, "\n") { + t.Fatalf("cancellation truncated mid-line: %q", got[max(0, len(got)-64):]) + } + }) + } +} + +func TestFormatterAlreadyCancelledEmitsNothing(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + cfg := &model.Config{Heading: true, LineNumber: true, Color: model.ColorNever} + var buf bytes.Buffer + err := NewFormatter(cfg).Format(ctx, &buf, makeTestAggregatedResults()) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if buf.Len() != 0 { + t.Fatalf("expected no output for an already-cancelled context, got %q", buf.String()) + } +} diff --git a/internal/output/single.go b/internal/output/single.go index 32af7a7..647153d 100644 --- a/internal/output/single.go +++ b/internal/output/single.go @@ -1,6 +1,7 @@ package output import ( + "context" "fmt" "io" "strconv" @@ -24,14 +25,20 @@ func NewSingleLineFormatter(cfg *model.Config) *SingleLineFormatter { } } -// Format writes single-line matches to w. -func (s *SingleLineFormatter) Format(w io.Writer, results *aggregator.AggregatedResults) error { +// Format writes single-line matches to w, aborting with ctx.Err() if ctx is +// cancelled mid-render. +func (s *SingleLineFormatter) Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error { if results == nil || len(results.Files) == 0 { return nil } c := s.color + guard := newCancelGuard(ctx) for _, file := range results.Files { + if err := guard.boundary(); err != nil { + return err + } + for _, commit := range file.Commits { shortSHA := commit.ShortSHA if shortSHA == "" { @@ -56,6 +63,9 @@ func (s *SingleLineFormatter) Format(w io.Writer, results *aggregator.Aggregated if err := s.writeEntry(w, shortSHA, file.Path, line.LineNum, line.LineText, line.IsMatch, line.Submatches); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } } else { @@ -63,6 +73,9 @@ func (s *SingleLineFormatter) Format(w io.Writer, results *aggregator.Aggregated if err := s.writeEntry(w, shortSHA, file.Path, match.LineNum, match.LineText, true, match.Submatches); err != nil { return err } + if err := guard.lines(1); err != nil { + return err + } } } } diff --git a/internal/output/writer.go b/internal/output/writer.go index 5494c7e..3c1291d 100644 --- a/internal/output/writer.go +++ b/internal/output/writer.go @@ -1,15 +1,58 @@ package output import ( + "context" "io" "github.com/kryft-dev/grg/internal/aggregator" "github.com/kryft-dev/grg/internal/model" ) +// cancelCheckInterval is the number of emitted lines between context +// cancellation checks performed by a cancelGuard. +const cancelCheckInterval = 1024 + // Formatter provides a unified interface for rendering aggregated search results. +// +// Format renders results to w and aborts with ctx.Err() if ctx is cancelled +// mid-render. Cancellation is observed at coarse boundaries, so output written +// before the abort is always a whole-line prefix of the complete rendering. type Formatter interface { - Format(w io.Writer, results *aggregator.AggregatedResults) error + Format(ctx context.Context, w io.Writer, results *aggregator.AggregatedResults) error +} + +// cancelGuard amortises cancellation checks over emitted output. Consulting +// the context per line would put an atomic load on the innermost write path, +// so the guard only reads it once every cancelCheckInterval lines and at +// structural boundaries; the per-line cost is a decrement and a branch. +type cancelGuard struct { + ctx context.Context + left int +} + +func newCancelGuard(ctx context.Context) cancelGuard { + if ctx == nil { + ctx = context.Background() + } + return cancelGuard{ctx: ctx, left: cancelCheckInterval} +} + +// lines records n freshly emitted lines, returning the context error once the +// check interval elapses and the context is done. +func (g *cancelGuard) lines(n int) error { + g.left -= n + if g.left > 0 { + return nil + } + g.left = cancelCheckInterval + return g.ctx.Err() +} + +// boundary forces a check at a coarse structural boundary, such as the start +// of a file group, and restarts the line interval. +func (g *cancelGuard) boundary() error { + g.left = cancelCheckInterval + return g.ctx.Err() } // NewFormatter constructs the appropriate Formatter based on configuration flags. diff --git a/test/leak_test.go b/test/leak_test.go index 82a67a8..efe5db2 100644 --- a/test/leak_test.go +++ b/test/leak_test.go @@ -402,7 +402,7 @@ func TestGoroutineLeak_EndToEndLifecycle(t *testing.T) { var outBuf bytes.Buffer formatter := output.NewFormatter(cfg) - if err := formatter.Format(&outBuf, aggregated); err != nil { + if err := formatter.Format(context.Background(), &outBuf, aggregated); err != nil { t.Fatalf("formatter.Format failed: %v", err) }