perf: reduce allocations on the chunk reading hot path

pkg/blobstore/buffer/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "cas_reader_buffer.go",
         "cas_validating_chunk_reader.go",
         "cas_validating_reader.go",
+        "chunk_buffer_pool.go",
         "chunk_reader.go",
         "chunk_reader_backed_reader.go",
         "common_conversions.go",
@@ -55,6 +56,8 @@ go_test(
         "new_proto_buffer_from_proto_test.go",
         "new_validated_buffer_from_byte_slice_test.go",
         "new_validated_buffer_from_reader_at_test.go",
+        "reader_backed_chunk_reader_bench_test.go",
+        "reader_backed_chunk_reader_test.go",
         "with_error_handler_test.go",
     ],
     deps = [

pkg/blobstore/buffer/chunk_buffer_pool.go

@@ -0,0 +1,26 @@
+package buffer
+
+import "sync"
+
+// chunkBufferPool caches []byte buffers used by ChunkReader
+// implementations that allocate a fixed-size working buffer per
+// stream. Buffers are stored as *[]byte to avoid the interface-boxing
+// allocation that occurs when storing a non-pointer value in
+// sync.Pool.
+var chunkBufferPool sync.Pool
+
+func getChunkBuffer(size int) *[]byte {
+	if v := chunkBufferPool.Get(); v != nil {
+		b := v.(*[]byte)
+		if cap(*b) >= size {
+			*b = (*b)[:size]
+			return b
+		}
+	}
+	b := make([]byte, size)
+	return &b
+}
+
+func putChunkBuffer(b *[]byte) {
+	chunkBufferPool.Put(b)
+}

pkg/blobstore/buffer/chunk_reader.go

@@ -10,6 +10,12 @@ import (
 // responsible for providing space for the data. This interface is
 // similar to how frame-based transfer protocols work, including the
 // Bytestream protocol used by REv2.
+//
+// The byte slice returned by Read() is only guaranteed to be valid
+// until the next call to Read() or Close() on the same ChunkReader.
+// Implementations are free to reuse the backing storage across Read()
+// calls in order to minimize allocations. Callers that need to retain
+// the contents past the next Read()/Close() must copy them.
 type ChunkReader interface {
 	Read() ([]byte, error)
 	Close()

pkg/blobstore/buffer/multiplexed_chunk_reader.go

@@ -32,6 +32,15 @@ func newMultiplexedChunkReader(r ChunkReader, additionalConsumers int) ChunkRead
 }
 
 func (r *multiplexedChunkReader) readAndShareWithOthers(currentConsumerContinues int) ([]byte, error) {
+	// All waiting consumers receive the exact same slice returned by
+	// the underlying ChunkReader.Read(). This is safe with respect
+	// to the "valid until next Read()/Close()" contract because the
+	// pendingConsumers counter ensures every consumer must call
+	// Read() (or Close()) again before another underlying Read()
+	// can be issued. By the time the next underlying Read()
+	// executes, every consumer has therefore released its reference
+	// to the previously returned slice, so the underlying reader is
+	// free to reuse the backing storage.
 	data, err := r.r.Read()
 	for _, c := range r.waitingConsumers {
 		c <- readResult{data: data, err: err}

pkg/blobstore/buffer/reader_backed_chunk_reader.go

@@ -8,6 +8,7 @@ type readerBackedChunkReader struct {
 	r                     io.ReadCloser
 	maximumChunkSizeBytes int
 
+	buf *[]byte
 	err error
 }
 
@@ -23,8 +24,11 @@ func newReaderBackedChunkReader(r io.ReadCloser, maximumChunkSizeBytes int) Chun
 
 func (r *readerBackedChunkReader) Read() ([]byte, error) {
 	if r.err == nil {
-		b := make([]byte, r.maximumChunkSizeBytes)
-		n, err := io.ReadFull(r.r, b[:])
+		if r.buf == nil {
+			r.buf = getChunkBuffer(r.maximumChunkSizeBytes)
+		}
+		b := *r.buf
+		n, err := io.ReadFull(r.r, b)
 		if err == io.ErrUnexpectedEOF {
 			r.err = io.EOF
 		} else {
@@ -38,5 +42,9 @@ func (r *readerBackedChunkReader) Read() ([]byte, error) {
 }
 
 func (r *readerBackedChunkReader) Close() {
+	if r.buf != nil {
+		putChunkBuffer(r.buf)
+		r.buf = nil
+	}
 	r.r.Close()
 }

pkg/blobstore/buffer/reader_backed_chunk_reader_bench_test.go

@@ -0,0 +1,82 @@
+package buffer_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/buildbarn/bb-storage/pkg/blobstore/buffer"
+)
+
+type readAtCloser struct{ *bytes.Reader }
+
+func (readAtCloser) Close() error { return nil }
+
+const benchChunkSize = 1 << 16
+
+var benchBlobSizes = []struct {
+	name     string
+	blobSize int64
+}{
+	{"BlobSize=4KiB", 4 << 10},
+	{"BlobSize=64KiB", 64 << 10},
+	{"BlobSize=256KiB", 256 << 10},
+	{"BlobSize=4MiB", 4 << 20},
+	{"BlobSize=64MiB", 64 << 20},
+}
+
+// NewValidatedBufferFromReaderAt is the cheapest public entry point
+// that resolves to newReaderBackedChunkReader without wrapping
+// decorators.
+func BenchmarkReaderBackedChunkReader(b *testing.B) {
+	for _, c := range benchBlobSizes {
+		b.Run(c.name, func(b *testing.B) {
+			data := make([]byte, c.blobSize)
+			for i := range data {
+				data[i] = byte(i)
+			}
+
+			b.SetBytes(c.blobSize)
+			b.ReportAllocs()
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				buf := buffer.NewValidatedBufferFromReaderAt(readAtCloser{bytes.NewReader(data)}, c.blobSize)
+				r := buf.ToChunkReader(0, benchChunkSize)
+				for {
+					if _, err := r.Read(); err != nil {
+						break
+					}
+				}
+				r.Close()
+			}
+		})
+	}
+}
+
+func BenchmarkReaderBackedChunkReaderParallel(b *testing.B) {
+	for _, c := range benchBlobSizes {
+		b.Run(c.name, func(b *testing.B) {
+			data := make([]byte, c.blobSize)
+			for i := range data {
+				data[i] = byte(i)
+			}
+
+			b.SetBytes(c.blobSize)
+			b.ReportAllocs()
+			b.ResetTimer()
+
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					buf := buffer.NewValidatedBufferFromReaderAt(readAtCloser{bytes.NewReader(data)}, c.blobSize)
+					r := buf.ToChunkReader(0, benchChunkSize)
+					for {
+						if _, err := r.Read(); err != nil {
+							break
+						}
+					}
+					r.Close()
+				}
+			})
+		})
+	}
+}

pkg/blobstore/buffer/reader_backed_chunk_reader_test.go

@@ -0,0 +1,71 @@
+package buffer_test
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/buildbarn/bb-storage/pkg/blobstore/buffer"
+	"github.com/stretchr/testify/require"
+)
+
+func TestReaderBackedChunkReaderReusesBuffer(t *testing.T) {
+	const (
+		chunkSize = 16
+		blobSize  = int64(chunkSize * 4)
+	)
+	data := make([]byte, blobSize)
+	for i := range data {
+		data[i] = byte(i)
+	}
+
+	b := buffer.NewValidatedBufferFromReaderAt(readAtCloser{bytes.NewReader(data)}, blobSize)
+	r := b.ToChunkReader(0, chunkSize)
+	defer r.Close()
+
+	first, err := r.Read()
+	require.NoError(t, err)
+	require.Equal(t, chunkSize, len(first))
+	firstAddr := &first[0]
+
+	for i := 0; i < 3; i++ {
+		next, err := r.Read()
+		require.NoError(t, err)
+		require.Equal(t, chunkSize, len(next))
+		require.Same(t, firstAddr, &next[0],
+			"successive chunks must share backing storage")
+	}
+
+	_, err = r.Read()
+	require.Equal(t, io.EOF, err)
+}
+
+// TestReaderBackedChunkReaderClosePoolsBuffer retries because
+// sync.Pool may drop entries between calls (e.g. on GC).
+func TestReaderBackedChunkReaderClosePoolsBuffer(t *testing.T) {
+	const chunkSize = 4096
+	data := make([]byte, chunkSize)
+
+	var observedReuse bool
+	for attempt := 0; attempt < 100 && !observedReuse; attempt++ {
+		b1 := buffer.NewValidatedBufferFromReaderAt(readAtCloser{bytes.NewReader(data)}, int64(chunkSize))
+		r1 := b1.ToChunkReader(0, chunkSize)
+		chunk, err := r1.Read()
+		require.NoError(t, err)
+		addr1 := &chunk[0]
+		_, err = r1.Read()
+		require.Equal(t, io.EOF, err)
+		r1.Close()
+
+		b2 := buffer.NewValidatedBufferFromReaderAt(readAtCloser{bytes.NewReader(data)}, int64(chunkSize))
+		r2 := b2.ToChunkReader(0, chunkSize)
+		chunk2, err := r2.Read()
+		require.NoError(t, err)
+		if &chunk2[0] == addr1 {
+			observedReuse = true
+		}
+		r2.Close()
+	}
+	require.True(t, observedReuse,
+		"expected sync.Pool to return a previously-released buffer at least once across 100 attempts")
+}

pkg/blobstore/grpcservers/BUILD.bazel

@@ -23,6 +23,7 @@ go_library(
         "//pkg/zstd",
         "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
+        "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
         "@org_golang_google_protobuf//types/known/emptypb",
@@ -32,14 +33,18 @@ go_library(
 go_test(
     name = "grpcservers_test",
     srcs = [
+        "byte_stream_server_bench_test.go",
         "byte_stream_server_test.go",
         "content_addressable_storage_server_test.go",
         "indirect_content_addressable_storage_server_test.go",
     ],
     deps = [
         ":grpcservers",
         "//internal/mock",
+        "//pkg/blobstore",
         "//pkg/blobstore/buffer",
+        "//pkg/blobstore/slicing",
+        "//pkg/capabilities",
         "//pkg/digest",
         "//pkg/proto/icas",
         "//pkg/testutil",
@@ -51,6 +56,7 @@ go_test(
         "@org_golang_google_genproto_googleapis_rpc//status",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//credentials/insecure",
         "@org_golang_google_grpc//status",
         "@org_golang_google_grpc//test/bufconn",
         "@org_uber_go_mock//gomock",

pkg/blobstore/grpcservers/byte_stream_server.go

@@ -13,6 +13,7 @@ import (
 	bb_zstd "github.com/buildbarn/bb-storage/pkg/zstd"
 
 	"google.golang.org/genproto/googleapis/bytestream"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -56,7 +57,18 @@ func (s *byteStreamServer) Read(in *bytestream.ReadRequest, out bytestream.ByteS
 			if readErr != nil {
 				return readErr
 			}
-			if writeErr := out.Send(&bytestream.ReadResponse{Data: readBuf}); writeErr != nil {
+			// A ChunkReader may reuse the slice returned by
+			// Read() across calls. Encoding into a grpc.PreparedMsg
+			// copies the wire-format bytes synchronously, so the
+			// underlying buffer is no longer referenced by gRPC after
+			// SendMsg returns. This is the only contract-safe way to
+			// hand a reusable buffer to grpc.SendMsg; see
+			// google.golang.org/grpc issues #5857 and #8186.
+			var prepared grpc.PreparedMsg
+			if err := prepared.Encode(out, &bytestream.ReadResponse{Data: readBuf}); err != nil {
+				return err
+			}
+			if writeErr := out.SendMsg(&prepared); writeErr != nil {
 				return writeErr
 			}
 		}
@@ -81,7 +93,14 @@ type readStreamWriter struct {
 }
 
 func (w *readStreamWriter) Write(p []byte) (int, error) {
-	if err := w.out.Send(&bytestream.ReadResponse{Data: p}); err != nil {
+	// Each Write encodes the chunk into a grpc.PreparedMsg before sending,
+	// so the caller's buffer is no longer referenced by gRPC after Write
+	// returns. See the IDENTITY path above for rationale.
+	var prepared grpc.PreparedMsg
+	if err := prepared.Encode(w.out, &bytestream.ReadResponse{Data: p}); err != nil {
+		return 0, err
+	}
+	if err := w.out.SendMsg(&prepared); err != nil {
 		return 0, err
 	}
 	return len(p), nil