From 63a31f835c14714ecd706d41e8d8f9ab19fcffa5 Mon Sep 17 00:00:00 2001 From: Keegan Smith Date: Tue, 20 Oct 2026 11:21:58 +0000 Subject: [PATCH] feat/diff: stream printed diffs to an io.Writer Printing a large diff used to hold several copies of it at once: PrintHunks built a growing buffer per file, PrintFileDiff copied it into a second buffer and PrintMultiFileDiff copied every file into a third, each growing by doubling. Sourcegraph's frontend prints whole repository comparisons this way, and diffs of hundreds of MiB have contributed to OOMs there. WriteMultiFileDiff and WriteFileDiff now write directly to an io.Writer, with hunk bodies passed through without copying and header lines assembled in one reused scratch slice instead of through fmt. Callers that can stream (such as an HTTP response) no longer need the printed diff in memory at all. The Print functions are built on the same code and keep their output byte-for-byte, including returning nil for empty output. They first measure the output with a counting writer and then fill a buffer of exactly that size, so the result is allocated once. For a 148 MiB diff of 200k files this cuts allocation from 816 MiB in 3.4M mallocs to 148 MiB in 5, and peak RSS growth from ~475 MiB to 148 MiB; for a single 121 MiB hunk, from 363 MiB to 121 MiB. Despite formatting twice, printing is about 10x faster. Amp-Thread-ID: https://ampcode.com/threads/T-01a15907-a942-77c9-ab83-3ff2545b0c9d Co-authored-by: Amp --- diff/diff_test.go | 46 +++++++++++ diff/print.go | 201 +++++++++++++++++++++++++++------------------- 2 files changed, 166 insertions(+), 81 deletions(-) diff --git a/diff/diff_test.go b/diff/diff_test.go index fc42e0b..673f076 100644 --- a/diff/diff_test.go +++ b/diff/diff_test.go @@ -5,6 +5,7 @@ import ( "github.com/google/go-cmp/cmp" "io" "io/ioutil" + "math" "path/filepath" "reflect" "strings" @@ -952,6 +953,51 @@ func TestParseMultiFileDiffAndPrintMultiFileDiff(t *testing.T) { if !bytes.Equal(printed, diffData) { t.Errorf("%s: printed multi-file diff != original multi-file diff\n\n# PrintMultiFileDiff output - Original:\n%s", test.filename, cmp.Diff(diffData, printed)) } + + var written bytes.Buffer + if err := WriteMultiFileDiff(&written, diffs); err != nil { + t.Errorf("%s: WriteMultiFileDiff: %s", test.filename, err) + } + if !bytes.Equal(written.Bytes(), diffData) { + t.Errorf("%s: written multi-file diff != original multi-file diff\n\n# WriteMultiFileDiff output - Original:\n%s", test.filename, cmp.Diff(diffData, written.Bytes())) + } + } +} + +// failingWriter fails the write that goes past the first n bytes and accepts +// all other writes, so a caller that ignores that one error goes unnoticed. +type failingWriter struct{ n int } + +func (w *failingWriter) Write(p []byte) (int, error) { + if len(p) > w.n { + n := w.n + w.n = math.MaxInt + return n, io.ErrShortWrite + } + w.n -= len(p) + return len(p), nil +} + +func TestWriteMultiFileDiffReturnsWriteError(t *testing.T) { + // Covers extended headers, "Only in" lines, timestamps, sections and + // missing newlines, so that every write in the printer can fail. + diffData, err := ioutil.ReadFile(filepath.Join("testdata", "sample_contains_added_deleted_files.diff")) + if err != nil { + t.Fatal(err) + } + diffData = append(diffData, "--- a\t2009-10-11 15:12:20.000000000 -0700\n+++ b\n@@ -1,1 +1,1 @@ section\n-a\n\\ No newline at end of file\n+b\n\\ No newline at end of file\n"...) + diffs, err := ParseMultiFileDiff(diffData) + if err != nil { + t.Fatal(err) + } + printed, err := PrintMultiFileDiff(diffs) + if err != nil || !bytes.Equal(printed, diffData) { + t.Fatalf("PrintMultiFileDiff did not round-trip: %v", err) + } + for n := 0; n < len(printed); n++ { + if err := WriteMultiFileDiff(&failingWriter{n: n}, diffs); err != io.ErrShortWrite { + t.Fatalf("failing after %d bytes: got err %v, want %v", n, err, io.ErrShortWrite) + } } } diff --git a/diff/print.go b/diff/print.go index 012651a..c0cebcb 100644 --- a/diff/print.go +++ b/diff/print.go @@ -5,137 +5,176 @@ import ( "fmt" "io" "path/filepath" + "strconv" "time" ) // PrintMultiFileDiff prints a multi-file diff in unified diff format. +// +// To avoid holding the whole diff in memory, use WriteMultiFileDiff. func PrintMultiFileDiff(ds []*FileDiff) ([]byte, error) { - var buf bytes.Buffer + return printExact(func(w io.Writer) error { return WriteMultiFileDiff(w, ds) }) +} + +// PrintFileDiff prints a FileDiff in unified diff format. +// +// To avoid holding the whole diff in memory, use WriteFileDiff. +func PrintFileDiff(d *FileDiff) ([]byte, error) { + return printExact(func(w io.Writer) error { return WriteFileDiff(w, d) }) +} + +// PrintHunks prints diff hunks in unified diff format. +func PrintHunks(hunks []*Hunk) ([]byte, error) { + return printExact(func(w io.Writer) error { return (&printer{w: w}).hunks(hunks) }) +} + +// printExact calls write twice: first to measure the output, then to fill a +// buffer of exactly that size. This keeps the peak memory of printing a large +// diff at the size of the output, instead of the up to 2x that a growing buffer +// needs while it copies itself into a larger allocation. +func printExact(write func(io.Writer) error) ([]byte, error) { + var n countingWriter + if err := write(&n); err != nil { + return nil, err + } + if n == 0 { + return nil, nil + } + buf := bytes.NewBuffer(make([]byte, 0, n)) + if err := write(buf); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// countingWriter counts the bytes written to it and discards them. +type countingWriter int + +func (c *countingWriter) Write(p []byte) (int, error) { + *c += countingWriter(len(p)) + return len(p), nil +} + +// WriteMultiFileDiff writes a multi-file diff in unified diff format to w. +// +// Hunk bodies are written to w directly, without intermediate copies, so many +// small writes go to w. Wrap w in a bufio.Writer if those writes are costly. +func WriteMultiFileDiff(w io.Writer, ds []*FileDiff) error { + p := printer{w: w} for _, d := range ds { - diff, err := PrintFileDiff(d) - if err != nil { - return nil, err - } - if _, err := buf.Write(diff); err != nil { - return nil, err + if err := p.fileDiff(d); err != nil { + return err } } - return buf.Bytes(), nil + return nil } -// PrintFileDiff prints a FileDiff in unified diff format. +// WriteFileDiff writes a FileDiff in unified diff format to w. +// +// Hunk bodies are written to w directly, without intermediate copies, so many +// small writes go to w. Wrap w in a bufio.Writer if those writes are costly. // // TODO(sqs): handle escaping whitespace/etc. chars in filenames -func PrintFileDiff(d *FileDiff) ([]byte, error) { - var buf bytes.Buffer +func WriteFileDiff(w io.Writer, d *FileDiff) error { + return (&printer{w: w}).fileDiff(d) +} + +// printer writes diffs to w. Header lines are assembled in line, which is +// reused across lines so that printing does not allocate per line. +type printer struct { + w io.Writer + line []byte +} +func (p *printer) fileDiff(d *FileDiff) error { for _, xheader := range d.Extended { - if _, err := fmt.Fprintln(&buf, xheader); err != nil { - return nil, err + p.line = append(append(p.line[:0], xheader...), '\n') + if err := p.writeLine(); err != nil { + return err } } // FileDiff is added/deleted file // No further hunks printing needed if d.NewName == "" { - _, err := fmt.Fprintf(&buf, onlyInMessage, filepath.Dir(d.OrigName), filepath.Base(d.OrigName)) - if err != nil { - return nil, err - } - return buf.Bytes(), nil + _, err := fmt.Fprintf(p.w, onlyInMessage, filepath.Dir(d.OrigName), filepath.Base(d.OrigName)) + return err } if d.Hunks == nil { - return buf.Bytes(), nil + return nil } - if err := printFileHeader(&buf, "--- ", d.OrigName, d.OrigTime); err != nil { - return nil, err - } - if err := printFileHeader(&buf, "+++ ", d.NewName, d.NewTime); err != nil { - return nil, err - } - - ph, err := PrintHunks(d.Hunks) - if err != nil { - return nil, err + if err := p.fileHeader("--- ", d.OrigName, d.OrigTime); err != nil { + return err } - - if _, err := buf.Write(ph); err != nil { - return nil, err + if err := p.fileHeader("+++ ", d.NewName, d.NewTime); err != nil { + return err } - return buf.Bytes(), nil + return p.hunks(d.Hunks) } -func printFileHeader(w io.Writer, prefix string, filename string, timestamp *time.Time) error { - if _, err := fmt.Fprint(w, prefix, filename); err != nil { - return err - } +func (p *printer) fileHeader(prefix string, filename string, timestamp *time.Time) error { + p.line = append(append(p.line[:0], prefix...), filename...) if timestamp != nil { - if _, err := fmt.Fprint(w, "\t", timestamp.Format(diffTimeFormatLayout)); err != nil { - return err - } + p.line = timestamp.AppendFormat(append(p.line, '\t'), diffTimeFormatLayout) } - if _, err := fmt.Fprintln(w); err != nil { - return err - } - return nil + p.line = append(p.line, '\n') + return p.writeLine() } -// PrintHunks prints diff hunks in unified diff format. -func PrintHunks(hunks []*Hunk) ([]byte, error) { - var buf bytes.Buffer +func (p *printer) hunks(hunks []*Hunk) error { for _, hunk := range hunks { - _, err := fmt.Fprintf(&buf, - "@@ -%d,%d +%d,%d @@", hunk.OrigStartLine, hunk.OrigLines, hunk.NewStartLine, hunk.NewLines, - ) - if err != nil { - return nil, err - } + // Equivalent to fmt.Sprintf(hunkHeader, ...) plus the section. + p.line = append(p.line[:0], "@@ -"...) + p.line = strconv.AppendInt(p.line, int64(hunk.OrigStartLine), 10) + p.line = append(p.line, ',') + p.line = strconv.AppendInt(p.line, int64(hunk.OrigLines), 10) + p.line = append(p.line, " +"...) + p.line = strconv.AppendInt(p.line, int64(hunk.NewStartLine), 10) + p.line = append(p.line, ',') + p.line = strconv.AppendInt(p.line, int64(hunk.NewLines), 10) + p.line = append(p.line, " @@"...) if hunk.Section != "" { - _, err := fmt.Fprint(&buf, " ", hunk.Section) - if err != nil { - return nil, err - } + p.line = append(append(p.line, ' '), hunk.Section...) } - if _, err := fmt.Fprintln(&buf); err != nil { - return nil, err + p.line = append(p.line, '\n') + if err := p.writeLine(); err != nil { + return err } if hunk.OrigNoNewlineAt == 0 { - if _, err := buf.Write(hunk.Body); err != nil { - return nil, err + if _, err := p.w.Write(hunk.Body); err != nil { + return err } } else { - if _, err := buf.Write(hunk.Body[:hunk.OrigNoNewlineAt]); err != nil { - return nil, err + if _, err := p.w.Write(hunk.Body[:hunk.OrigNoNewlineAt]); err != nil { + return err } - if err := printNoNewlineMessage(&buf); err != nil { - return nil, err + if err := p.noNewlineMessage(); err != nil { + return err } - if _, err := buf.Write(hunk.Body[hunk.OrigNoNewlineAt:]); err != nil { - return nil, err + if _, err := p.w.Write(hunk.Body[hunk.OrigNoNewlineAt:]); err != nil { + return err } } if !bytes.HasSuffix(hunk.Body, []byte{'\n'}) { - if _, err := fmt.Fprintln(&buf); err != nil { - return nil, err - } - if err := printNoNewlineMessage(&buf); err != nil { - return nil, err + p.line = append(append(append(p.line[:0], '\n'), noNewlineMessage...), '\n') + if err := p.writeLine(); err != nil { + return err } } } - return buf.Bytes(), nil + return nil } -func printNoNewlineMessage(w io.Writer) error { - if _, err := w.Write([]byte(noNewlineMessage)); err != nil { - return err - } - if _, err := fmt.Fprintln(w); err != nil { - return err - } - return nil +func (p *printer) noNewlineMessage() error { + p.line = append(append(p.line[:0], noNewlineMessage...), '\n') + return p.writeLine() +} + +func (p *printer) writeLine() error { + _, err := p.w.Write(p.line) + return err }