Land #907

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])
 

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)
+		}
+	})
+}

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

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)
 }

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 {