From 8f35ed8e7edee7946c6a06191c05774b4cfabb39 Mon Sep 17 00:00:00 2001 From: Mattias Jonsson Date: Sun, 12 Jul 2026 13:37:03 +0200 Subject: [PATCH 1/2] 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 | 31 ++++++++++++- cache/disk/casblob/casblob_test.go | 73 ++++++++++++++++++++++++++++++ cache/disk/zstdimpl/cgozstd.go | 4 +- cache/disk/zstdimpl/gozstd.go | 4 +- cache/disk/zstdimpl/zstdimpl.go | 8 +++- 5 files changed, 113 insertions(+), 7 deletions(-) diff --git a/cache/disk/casblob/casblob.go b/cache/disk/casblob/casblob.go index 31a5953fe..2fb6c4400 100644 --- a/cache/disk/casblob/casblob.go +++ b/cache/disk/casblob/casblob.go @@ -396,7 +396,8 @@ func GetZstdReadCloser(zstd zstdimpl.ZstdImpl, f *os.File, expectedSize int64, o } chunkToRecompress := uncompressedFirstChunk[remainder:] - recompressedChunk := zstd.EncodeAll(chunkToRecompress) + dst := make([]byte, 0, compressBound(len(chunkToRecompress))) + recompressedChunk := zstd.EncodeAll(dst, chunkToRecompress) br := bytes.NewReader(recompressedChunk) if chunkNum == int64(len(h.chunkOffsets)-2) { @@ -517,6 +518,27 @@ var chunkBufferPool = &sync.Pool{ }, } +// compressBound returns the guaranteed maximum compressed size of srcSize +// bytes. It is a Go port of the ZSTD_COMPRESSBOUND() C macro from libzstd: +// +// #define ZSTD_COMPRESSBOUND(srcSize) \ +// ((srcSize) + ((srcSize)>>8) + \ +// (((srcSize) < (128<<10)) ? (((128<<10) - (srcSize)) >> 11) : 0)) +// +// zstd's own C.ZSTD_compressBound() is used by the cgo implementation; this +// keeps the pure-Go path from depending on any single backend's internals. +// We use it to size the reusable compressed-output buffer so that EncodeAll +// never has to grow (and thus reallocate) it, even for incompressible chunks +// whose output is slightly larger than the input. +func compressBound(srcSize int) int { + const lowLimit = 128 << 10 // 128 KiB + margin := 0 + if srcSize < lowLimit { + margin = (lowLimit - srcSize) >> 11 + } + return srcSize + (srcSize >> 8) + margin +} + // Read from r and write to f, using CompressionType t. // Return the size on disk or an error if something went wrong. func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t CompressionType, hash string, size int64) (int64, error) { @@ -592,6 +614,11 @@ func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t Compressio }() uncompressedChunk := *chunkBufferPtr + // Reusable output buffer for the compressed chunks. Sized to the maximum + // compressed size of a chunk so EncodeAll writes into it without growing, + // and reused for every chunk of this blob. + compressedChunkBuffer := make([]byte, 0, compressBound(int(chunkSize))) + hasher := sha256.New() for nextChunk < len(h.chunkOffsets)-1 { @@ -609,7 +636,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(compressedChunkBuffer[:0], uncompressedChunk[0:chunkEnd]) hasher.Write(uncompressedChunk[0:chunkEnd]) diff --git a/cache/disk/casblob/casblob_test.go b/cache/disk/casblob/casblob_test.go index 32906fec8..8cb668f41 100644 --- a/cache/disk/casblob/casblob_test.go +++ b/cache/disk/casblob/casblob_test.go @@ -82,3 +82,76 @@ func TestZstdFromLegacy(t *testing.T) { t.Fatalf("Unexpected content sha %s, expected %s", hs, hash) } } + +// blobSizeForBenchmark is deliberately several times the 1 MiB chunk size, so +// that WriteAndClose compresses the blob in multiple chunks. Each chunk in the +// zstd write path currently allocates a fresh output buffer +// (zstd.EncodeAll(in, nil)), so the transient garbage produced per Put scales +// with the blob size. This is the allocation churn that drives the OOM under +// concurrent upload bursts (see bazel-remote-oom-findings.md). +const blobSizeForBenchmark = 16 * 1024 * 1024 // 16 MiB => 16 chunks + +// writeBlob runs a single WriteAndClose of the pre-generated blob to a fresh +// temp file, then removes it. It is the unit of work for the benchmarks below. +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 the allocations of the zstd write path +// for a single upload. Run with -benchmem; the "B/op" figure is dominated by +// the per-chunk compressed-output buffers and is the regression metric for the +// output-buffer pooling fix. +func BenchmarkWriteAndCloseZstd(b *testing.B) { + zstd, err := zstdimpl.Get("go") + if err != nil { + b.Fatal(err) + } + + // Random (incompressible) data keeps each chunk's output buffer close to + // the full 1 MiB, i.e. the worst case for the per-chunk allocation. + 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 the concurrent upload burst: +// many Puts compressing at once, each churning per-chunk output buffers with no +// memory backpressure. Run with -benchmem to see the aggregate allocation 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..3ba8e44cb 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(dst, in []byte) []byte { + return gozstd.CompressLevel(dst, in, compressionLevel) } // -- Reader pool diff --git a/cache/disk/zstdimpl/gozstd.go b/cache/disk/zstdimpl/gozstd.go index 720da9ee8..cbceee19a 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(dst, in []byte) []byte { + return encoder.EncodeAll(in, dst) } diff --git a/cache/disk/zstdimpl/zstdimpl.go b/cache/disk/zstdimpl/zstdimpl.go index cc0102031..e794521c6 100644 --- a/cache/disk/zstdimpl/zstdimpl.go +++ b/cache/disk/zstdimpl/zstdimpl.go @@ -39,7 +39,13 @@ 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 in and appends the result to dst, returning the + // updated slice (like the underlying zstd EncodeAll). Passing a dst with + // enough spare capacity lets callers reuse a single output buffer across + // many chunks instead of allocating a fresh one per call. Pass a nil dst + // to let the implementation allocate. + EncodeAll(dst, in []byte) []byte } type zstdEncoder interface { From 1a5e86a64ef310cb2d9564a677cec939a26853f3 Mon Sep 17 00:00:00 2001 From: Mattias Jonsson Date: Tue, 14 Jul 2026 01:51:32 +0200 Subject: [PATCH 2/2] casblob: address review comments - zstdimpl.EncodeAll: rename params to (src, dst), matching the order github.com/klauspost/compress/zstd uses. - WriteAndClose: drop the compressBound() helper; size the reused output buffer inline to ZSTD_COMPRESSBOUND (srcSize + srcSize>>8 for a >= 128 KiB input) and document the bound for both backends. - GetZstdReadCloser: pass a nil dst (restores the pre-change allocate-fresh behaviour on this one-shot, rarely-hit recompress path). - benchmark: fix a stale comment reference (link to the PR instead of a non-existent local file) and trim the comments. --- cache/disk/casblob/casblob.go | 41 ++++++++++-------------------- cache/disk/casblob/casblob_test.go | 28 ++++++++------------ cache/disk/zstdimpl/cgozstd.go | 4 +-- cache/disk/zstdimpl/gozstd.go | 4 +-- cache/disk/zstdimpl/zstdimpl.go | 10 +++----- 5 files changed, 32 insertions(+), 55 deletions(-) diff --git a/cache/disk/casblob/casblob.go b/cache/disk/casblob/casblob.go index 2fb6c4400..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,8 +398,7 @@ func GetZstdReadCloser(zstd zstdimpl.ZstdImpl, f *os.File, expectedSize int64, o } chunkToRecompress := uncompressedFirstChunk[remainder:] - dst := make([]byte, 0, compressBound(len(chunkToRecompress))) - recompressedChunk := zstd.EncodeAll(dst, chunkToRecompress) + recompressedChunk := zstd.EncodeAll(chunkToRecompress, nil) br := bytes.NewReader(recompressedChunk) if chunkNum == int64(len(h.chunkOffsets)-2) { @@ -518,27 +519,6 @@ var chunkBufferPool = &sync.Pool{ }, } -// compressBound returns the guaranteed maximum compressed size of srcSize -// bytes. It is a Go port of the ZSTD_COMPRESSBOUND() C macro from libzstd: -// -// #define ZSTD_COMPRESSBOUND(srcSize) \ -// ((srcSize) + ((srcSize)>>8) + \ -// (((srcSize) < (128<<10)) ? (((128<<10) - (srcSize)) >> 11) : 0)) -// -// zstd's own C.ZSTD_compressBound() is used by the cgo implementation; this -// keeps the pure-Go path from depending on any single backend's internals. -// We use it to size the reusable compressed-output buffer so that EncodeAll -// never has to grow (and thus reallocate) it, even for incompressible chunks -// whose output is slightly larger than the input. -func compressBound(srcSize int) int { - const lowLimit = 128 << 10 // 128 KiB - margin := 0 - if srcSize < lowLimit { - margin = (lowLimit - srcSize) >> 11 - } - return srcSize + (srcSize >> 8) + margin -} - // Read from r and write to f, using CompressionType t. // Return the size on disk or an error if something went wrong. func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t CompressionType, hash string, size int64) (int64, error) { @@ -614,10 +594,15 @@ func WriteAndClose(zstd zstdimpl.ZstdImpl, r io.Reader, f *os.File, t Compressio }() uncompressedChunk := *chunkBufferPtr - // Reusable output buffer for the compressed chunks. Sized to the maximum - // compressed size of a chunk so EncodeAll writes into it without growing, - // and reused for every chunk of this blob. - compressedChunkBuffer := make([]byte, 0, compressBound(int(chunkSize))) + // 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() @@ -636,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(compressedChunkBuffer[:0], 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 8cb668f41..bd17a71cf 100644 --- a/cache/disk/casblob/casblob_test.go +++ b/cache/disk/casblob/casblob_test.go @@ -83,16 +83,13 @@ func TestZstdFromLegacy(t *testing.T) { } } -// blobSizeForBenchmark is deliberately several times the 1 MiB chunk size, so -// that WriteAndClose compresses the blob in multiple chunks. Each chunk in the -// zstd write path currently allocates a fresh output buffer -// (zstd.EncodeAll(in, nil)), so the transient garbage produced per Put scales -// with the blob size. This is the allocation churn that drives the OOM under -// concurrent upload bursts (see bazel-remote-oom-findings.md). +// 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 runs a single WriteAndClose of the pre-generated blob to a fresh -// temp file, then removes it. It is the unit of work for the benchmarks below. +// 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 { @@ -109,18 +106,16 @@ func writeBlob(tb testing.TB, zstd zstdimpl.ZstdImpl, dir string, data []byte, h } } -// BenchmarkWriteAndCloseZstd measures the allocations of the zstd write path -// for a single upload. Run with -benchmem; the "B/op" figure is dominated by -// the per-chunk compressed-output buffers and is the regression metric for the -// output-buffer pooling fix. +// 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) } - // Random (incompressible) data keeps each chunk's output buffer close to - // the full 1 MiB, i.e. the worst case for the per-chunk allocation. + // Incompressible data is the worst case: each chunk's output stays near the + // full 1 MiB. data, hash := testutils.RandomDataAndHash(blobSizeForBenchmark) dir := b.TempDir() @@ -133,9 +128,8 @@ func BenchmarkWriteAndCloseZstd(b *testing.B) { } } -// BenchmarkWriteAndCloseZstdParallel reproduces the concurrent upload burst: -// many Puts compressing at once, each churning per-chunk output buffers with no -// memory backpressure. Run with -benchmem to see the aggregate allocation rate. +// 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 { diff --git a/cache/disk/zstdimpl/cgozstd.go b/cache/disk/zstdimpl/cgozstd.go index 3ba8e44cb..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(dst, in []byte) []byte { - return gozstd.CompressLevel(dst, 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 cbceee19a..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(dst, in []byte) []byte { - return encoder.EncodeAll(in, dst) +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 e794521c6..9b50c3992 100644 --- a/cache/disk/zstdimpl/zstdimpl.go +++ b/cache/disk/zstdimpl/zstdimpl.go @@ -40,12 +40,10 @@ type ZstdImpl interface { GetEncoder(out io.WriteCloser) (zstdEncoder, error) DecodeAll(in []byte) ([]byte, error) - // EncodeAll compresses in and appends the result to dst, returning the - // updated slice (like the underlying zstd EncodeAll). Passing a dst with - // enough spare capacity lets callers reuse a single output buffer across - // many chunks instead of allocating a fresh one per call. Pass a nil dst - // to let the implementation allocate. - EncodeAll(dst, 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 {