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 {