From 0d0953188c8b2084fad76df7b8981f5dffd6bcbf Mon Sep 17 00:00:00 2001 From: Mattias Jonsson Date: Sun, 12 Jul 2026 13:37:03 +0200 Subject: [PATCH] casblob: reuse compressed-output buffer to cut upload memory churn Under a burst of concurrent uploads, bazel-remote could be OOM-killed by a transient Go heap overshoot even though idle/steady-state memory is small. Heap profiling (alloc_space) pointed at zstd.(*Encoder).encodeAll: the write path compresses each blob in 1 MiB chunks and called EncodeAll(in, nil) per chunk, so a fresh output slice was allocated for every chunk of every upload. The chunk size is exactly 1<<20, and klauspost's EncodeAll only pre-allocates an output buffer when len(src) < 1<<20 (strictly less-than). A full chunk therefore starts from a nil dst and grows it by repeated append-doubling, so each 1 MiB chunk churned several MiB of transient garbage. With no memory backpressure on concurrent Puts, the allocation rate outran the GC and the process was killed. Thread a reusable dst through zstdimpl.EncodeAll (both the pure-Go and cgo backends already accept one) and have casblob.WriteAndClose allocate a single output buffer, sized to compressBound(chunkSize), reused across all chunks of the blob. Sizing to the ZSTD_compressBound worst case ensures EncodeAll never has to grow (and reallocate) the buffer, even for incompressible chunks whose output is slightly larger than the input. A microbenchmark of the write path (16 MiB blob, incompressible data) shows per-upload allocations drop from ~79 MB/op to ~1.1 MB/op. --- cache/disk/casblob/casblob.go | 16 ++++++- cache/disk/casblob/casblob_test.go | 67 ++++++++++++++++++++++++++++++ cache/disk/zstdimpl/cgozstd.go | 4 +- cache/disk/zstdimpl/gozstd.go | 4 +- cache/disk/zstdimpl/zstdimpl.go | 6 ++- 5 files changed, 90 insertions(+), 7 deletions(-) diff --git a/cache/disk/casblob/casblob.go b/cache/disk/casblob/casblob.go index 31a5953fe..995fd3493 100644 --- a/cache/disk/casblob/casblob.go +++ b/cache/disk/casblob/casblob.go @@ -22,6 +22,8 @@ const ( Zstandard CompressionType = 1 ) +// If changed to < 128 KiB, WriteAndClose's output-buffer sizing must be updated +// (see the compressedChunkBuffer comment there). const defaultChunkSize = 1024 * 1024 * 1 // 1M // 4 bytes, to be written to disk in little-endian format. @@ -396,7 +398,7 @@ func GetZstdReadCloser(zstd zstdimpl.ZstdImpl, f *os.File, expectedSize int64, o } chunkToRecompress := uncompressedFirstChunk[remainder:] - recompressedChunk := zstd.EncodeAll(chunkToRecompress) + recompressedChunk := zstd.EncodeAll(chunkToRecompress, nil) br := bytes.NewReader(recompressedChunk) if chunkNum == int64(len(h.chunkOffsets)-2) { @@ -592,6 +594,16 @@ func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t Compressio }() uncompressedChunk := *chunkBufferPtr + // Output buffer reused for every chunk, sized to zstd's ZSTD_COMPRESSBOUND + // (srcSize + srcSize>>8 for a >= 128 KiB input, the incompressible worst + // case) so EncodeAll never grows it. This also bounds the pure-Go + // github.com/klauspost/compress/zstd backend for any such size: + // Encoder.MaxEncodedSize is srcSize + a <=14-byte frame header + 3 bytes per + // 64 KiB block (65 B for a 1 MiB chunk), far under the srcSize>>8 margin. + // An undersized buffer would only cost a reallocation, never fail, so this + // bound is best-effort, not a correctness requirement. + compressedChunkBuffer := make([]byte, 0, int(chunkSize+chunkSize>>8)) + hasher := sha256.New() for nextChunk < len(h.chunkOffsets)-1 { @@ -609,7 +621,7 @@ func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t Compressio return -1, fmt.Errorf("only managed to read %d of %d bytes: %w", numRead, chunkEnd, err) } - compressedChunk := zstd.EncodeAll(uncompressedChunk[0:chunkEnd]) + compressedChunk := zstd.EncodeAll(uncompressedChunk[0:chunkEnd], compressedChunkBuffer[:0]) hasher.Write(uncompressedChunk[0:chunkEnd]) diff --git a/cache/disk/casblob/casblob_test.go b/cache/disk/casblob/casblob_test.go index 32906fec8..bd17a71cf 100644 --- a/cache/disk/casblob/casblob_test.go +++ b/cache/disk/casblob/casblob_test.go @@ -82,3 +82,70 @@ func TestZstdFromLegacy(t *testing.T) { t.Fatalf("Unexpected content sha %s, expected %s", hs, hash) } } + +// blobSizeForBenchmark spans several 1 MiB chunks so WriteAndClose compresses +// in a loop, exercising the per-chunk output-buffer reuse. +// See https://github.com/buchgr/bazel-remote/pull/907. +const blobSizeForBenchmark = 16 * 1024 * 1024 // 16 MiB => 16 chunks + +// writeBlob is the benchmarks' unit of work: one WriteAndClose to a fresh temp +// file, then remove it. +func writeBlob(tb testing.TB, zstd zstdimpl.ZstdImpl, dir string, data []byte, hash string) { + f, err := os.CreateTemp(dir, "blob-") + if err != nil { + tb.Fatal(err) + } + name := f.Name() + _, err = casblob.WriteAndClose(zstd, bytes.NewReader(data), f, + casblob.Zstandard, hash, int64(len(data))) + if err != nil { + tb.Fatal(err) + } + if err := os.Remove(name); err != nil { + tb.Fatal(err) + } +} + +// BenchmarkWriteAndCloseZstd measures allocations of the zstd write path for a +// single upload. Run with -benchmem; B/op is the regression metric. +func BenchmarkWriteAndCloseZstd(b *testing.B) { + zstd, err := zstdimpl.Get("go") + if err != nil { + b.Fatal(err) + } + + // Incompressible data is the worst case: each chunk's output stays near the + // full 1 MiB. + data, hash := testutils.RandomDataAndHash(blobSizeForBenchmark) + dir := b.TempDir() + + b.SetBytes(blobSizeForBenchmark) + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + writeBlob(b, zstd, dir, data, hash) + } +} + +// BenchmarkWriteAndCloseZstdParallel reproduces a concurrent upload burst: many +// Puts compressing at once. Run with -benchmem for the aggregate alloc rate. +func BenchmarkWriteAndCloseZstdParallel(b *testing.B) { + zstd, err := zstdimpl.Get("go") + if err != nil { + b.Fatal(err) + } + + data, hash := testutils.RandomDataAndHash(blobSizeForBenchmark) + dir := b.TempDir() + + b.SetBytes(blobSizeForBenchmark) + b.ReportAllocs() + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + writeBlob(b, zstd, dir, data, hash) + } + }) +} diff --git a/cache/disk/zstdimpl/cgozstd.go b/cache/disk/zstdimpl/cgozstd.go index 55b773cd7..bab5a03ad 100644 --- a/cache/disk/zstdimpl/cgozstd.go +++ b/cache/disk/zstdimpl/cgozstd.go @@ -35,8 +35,8 @@ func (cgoZstd) DecodeAll(in []byte) ([]byte, error) { return gozstd.Decompress(nil, in) } -func (cgoZstd) EncodeAll(in []byte) []byte { - return gozstd.CompressLevel(nil, in, compressionLevel) +func (cgoZstd) EncodeAll(src, dst []byte) []byte { + return gozstd.CompressLevel(dst, src, compressionLevel) } // -- Reader pool diff --git a/cache/disk/zstdimpl/gozstd.go b/cache/disk/zstdimpl/gozstd.go index 720da9ee8..9b78a3bea 100644 --- a/cache/disk/zstdimpl/gozstd.go +++ b/cache/disk/zstdimpl/gozstd.go @@ -65,6 +65,6 @@ func (goZstd) DecodeAll(in []byte) ([]byte, error) { return decoder.DecodeAll(in, nil) } -func (goZstd) EncodeAll(in []byte) []byte { - return encoder.EncodeAll(in, nil) +func (goZstd) EncodeAll(src, dst []byte) []byte { + return encoder.EncodeAll(src, dst) } diff --git a/cache/disk/zstdimpl/zstdimpl.go b/cache/disk/zstdimpl/zstdimpl.go index cc0102031..9b50c3992 100644 --- a/cache/disk/zstdimpl/zstdimpl.go +++ b/cache/disk/zstdimpl/zstdimpl.go @@ -39,7 +39,11 @@ type ZstdImpl interface { GetDecoder(in io.ReadCloser) (io.ReadCloser, error) GetEncoder(out io.WriteCloser) (zstdEncoder, error) DecodeAll(in []byte) ([]byte, error) - EncodeAll(in []byte) []byte + + // EncodeAll compresses src and appends the result to dst, returning the + // updated slice (like github.com/klauspost/compress/zstd's EncodeAll). A dst + // with spare capacity is reused instead of allocating; pass nil to allocate. + EncodeAll(src, dst []byte) []byte } type zstdEncoder interface {