perf: upload local/remote at the same time

Author: tyler-french

enterprise/server/byte_stream_server_proxy/BUILD

@@ -21,7 +21,6 @@ go_library(
         "//server/util/authutil",
         "//server/util/bytebufferpool",
         "//server/util/compression",
-        "//server/util/lib/set",
         "//server/util/log",
         "//server/util/prefix",
         "//server/util/status",

enterprise/server/byte_stream_server_proxy/byte_stream_server_proxy.go

@@ -20,7 +20,6 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/util/authutil"
 	"github.com/buildbuddy-io/buildbuddy/server/util/bytebufferpool"
 	"github.com/buildbuddy-io/buildbuddy/server/util/compression"
-	"github.com/buildbuddy-io/buildbuddy/server/util/lib/set"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/prefix"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
@@ -30,10 +29,12 @@ import (
 
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
 	bspb "google.golang.org/genproto/googleapis/bytestream"
+	"google.golang.org/grpc/codes"
+	gstatus "google.golang.org/grpc/status"
 )
 
 var (
-	chunkUploadConcurrency = flag.Int("cache_proxy.chunk_upload_concurrency", 8, "Maximum number of concurrent chunk uploads when uploading missing chunks to remote cache.")
+	chunkUploadConcurrency = flag.Int("cache_proxy.chunk_upload_concurrency", 2, "Maximum number of concurrent chunk BatchUpdateBlobsRequest uploads in flight when uploading missing chunks to remote cache.")
 )
 
 type ByteStreamServerProxy struct {
@@ -733,33 +734,84 @@ func (s *ByteStreamServerProxy) writeChunked(ctx context.Context, stream bspb.By
 	compressBuf := s.compressBufPool.Get(chunking.MaxChunkSizeBytes())
 	defer s.compressBufPool.Put(compressBuf)
 
-	// chunkWriteFn is called on each new chunk once it's available through the chunker's pipe.
-	// We write chunks to local first, then use FindMissingBlobs + upload for remote.
-	// Chunks are stored and read compressed with ZSTD.
+	uploadGroup, uCtx := errgroup.WithContext(ctx)
+	uploadGroup.SetLimit(*chunkUploadConcurrency)
+	pendingBatch := &repb.BatchUpdateBlobsRequest{
+		InstanceName:   instanceName,
+		DigestFunction: digestFunction,
+	}
+	var pendingBatchSize int64
+	var chunksDeduped int
+	var chunkBytesDeduped int64
+
 	chunkWriteFn := func(chunkData []byte) error {
 		chunkDigest, err := digest.Compute(bytes.NewReader(chunkData), digestFunction)
 		if err != nil {
-			return status.InternalErrorf("computing chunked digest for Write: %s", err)
+			return status.WrapError(err, "computing chunk digest")
 		}
 		chunkDigests = append(chunkDigests, chunkDigest)
 
 		chunkRN := digest.NewCASResourceName(chunkDigest, instanceName, digestFunction)
 		chunkRN.SetCompressor(repb.Compressor_ZSTD)
 		compressedData := compression.CompressZstd(compressBuf, chunkData)
-		stream := &rawWriteStream{
+
+		eg, gCtx := errgroup.WithContext(ctx)
+		eg.Go(func() error {
+			fmbResp, err := s.remoteCAS.FindMissingBlobs(gCtx, &repb.FindMissingBlobsRequest{
+				InstanceName:   instanceName,
+				DigestFunction: digestFunction,
+				BlobDigests:    []*repb.Digest{chunkDigest},
+			})
+			if err != nil {
+				return status.WrapErrorf(err, "finding missing blob on remote for chunk %s", chunkRN.DownloadString())
+			}
+			if len(fmbResp.GetMissingBlobDigests()) == 0 {
+				chunksDeduped++
+				chunkBytesDeduped += chunkDigest.GetSizeBytes()
+				return nil
+			}
+
+			additionalSize := int64(len(compressedData))
+
+			// If this chunk won't fit, fire off the current batch.
+			// uploadGroup.Go blocks when 2 RPCs are in flight.
+			if pendingBatchSize+additionalSize > cachetools.BatchUploadLimitBytes && len(pendingBatch.Requests) > 0 {
+				batch := pendingBatch
+				pendingBatch = &repb.BatchUpdateBlobsRequest{
+					InstanceName:   instanceName,
+					DigestFunction: digestFunction,
+				}
+				pendingBatchSize = 0
+				uploadGroup.Go(func() error {
+					return s.sendBatchUpload(uCtx, batch)
+				})
+			}
+
+			dataCopy := make([]byte, len(compressedData))
+			copy(dataCopy, compressedData)
+			pendingBatch.Requests = append(pendingBatch.Requests, &repb.BatchUpdateBlobsRequest_Request{
+				Digest:     chunkDigest,
+				Data:       dataCopy,
+				Compressor: repb.Compressor_ZSTD,
+			})
+			pendingBatchSize += additionalSize
+			return nil
+		})
+
+		localStream := &rawWriteStream{
 			ctx:          ctx,
 			resourceName: chunkRN.NewUploadString(),
 			data:         compressedData,
 		}
-		if err := s.local.Write(stream); err != nil {
-			return status.InternalErrorf("writing chunk %s to local: %s", chunkRN.DownloadString(), err)
+		if err := s.local.Write(localStream); err != nil {
+			return status.WrapErrorf(err, "writing chunk %s to local", chunkRN.DownloadString())
 		}
-		return nil
+		return eg.Wait()
 	}
 
 	chunker, err := chunking.NewChunker(ctx, int(chunking.AvgChunkSizeBytes()), chunkWriteFn)
 	if err != nil {
-		return writeChunkedResult{}, status.InternalErrorf("creating chunker: %s", err)
+		return writeChunkedResult{}, status.WrapError(err, "creating chunker")
 	}
 	defer chunker.Close()
 
@@ -800,20 +852,30 @@ func (s *ByteStreamServerProxy) writeChunked(ctx context.Context, stream bspb.By
 		}
 	}
 
-	// Close blocks until all chunk writes complete, ensuring chunkDigests is fully populated.
 	if err := chunker.Close(); err != nil {
 		return writeChunkedResult{}, status.InternalErrorf("closing chunker: %s", err)
 	}
 
+	if len(pendingBatch.Requests) > 0 {
+		uploadGroup.Go(func() error {
+			return s.sendBatchUpload(uCtx, pendingBatch)
+		})
+	}
+	if err := uploadGroup.Wait(); err != nil {
+		return writeChunkedResult{}, status.WrapError(err, "uploading chunks to remote")
+	}
+
 	var chunkBytesTotal int64
 	for _, d := range chunkDigests {
 		chunkBytesTotal += d.GetSizeBytes()
 	}
 
 	result := writeChunkedResult{
-		blobBytes:       blobSize,
-		chunksTotal:     len(chunkDigests),
-		chunkBytesTotal: chunkBytesTotal,
+		blobBytes:         blobSize,
+		chunksTotal:       len(chunkDigests),
+		chunksDeduped:     chunksDeduped,
+		chunkBytesTotal:   chunkBytesTotal,
+		chunkBytesDeduped: chunkBytesDeduped,
 	}
 
 	// If there's only 1 chunk, the chunking threshold is misconfigured.
@@ -829,63 +891,22 @@ func (s *ByteStreamServerProxy) writeChunked(ctx context.Context, stream bspb.By
 		DigestFunction: digestFunction,
 	}
 
-	missingBlobs, err := s.remoteCAS.FindMissingBlobs(ctx, manifest.ToFindMissingBlobsRequest())
-	if err != nil {
-		return writeChunkedResult{}, status.InternalErrorf("finding missing blobs on remote: %s", err)
-	}
-	missingDigests := missingBlobs.GetMissingBlobDigests()
-	missingSet := make(set.Set[string], len(missingDigests))
-	for _, d := range missingDigests {
-		missingSet.Add(d.GetHash())
-	}
-
-	// Deduped chunks are ones that already exist on remote.
-	for _, d := range chunkDigests {
-		if !missingSet.Contains(d.GetHash()) {
-			result.chunksDeduped++
-			result.chunkBytesDeduped += d.GetSizeBytes()
-		}
-	}
-	if len(missingDigests) > 0 {
-		if err := s.uploadMissingChunks(ctx, missingDigests, instanceName, digestFunction); err != nil {
-			return writeChunkedResult{}, status.InternalErrorf("uploading missing chunks to remote: %s", err)
-		}
-	}
-
 	if _, err := s.remoteCAS.SpliceBlob(ctx, manifest.ToSpliceBlobRequest()); err != nil {
-		return writeChunkedResult{}, status.InternalErrorf("splice blob on remote: %s", err)
+		return writeChunkedResult{}, status.WrapError(err, "splice blob on remote")
 	}
 
 	return result, stream.SendAndClose(&bspb.WriteResponse{CommittedSize: bytesReceived})
 }
 
-func (s *ByteStreamServerProxy) uploadMissingChunks(ctx context.Context, missingDigests []*repb.Digest, instanceName string, digestFunction repb.DigestFunction_Value) error {
-	g, gCtx := errgroup.WithContext(ctx)
-	g.SetLimit(*chunkUploadConcurrency)
-	for _, chunkDigest := range missingDigests {
-		chunkRN := digest.NewCASResourceName(chunkDigest, instanceName, digestFunction)
-		g.Go(func() error {
-			return s.uploadChunk(gCtx, chunkRN)
-		})
-	}
-	return g.Wait()
-}
-
-func (s *ByteStreamServerProxy) uploadChunk(ctx context.Context, rn *digest.CASResourceName) error {
-	rnCopy, err := digest.CASResourceNameFromProto(rn.ToProto())
+func (s *ByteStreamServerProxy) sendBatchUpload(ctx context.Context, req *repb.BatchUpdateBlobsRequest) error {
+	rsp, err := s.remoteCAS.BatchUpdateBlobs(ctx, req)
 	if err != nil {
 		return err
 	}
-	rnCopy.SetCompressor(repb.Compressor_ZSTD)
-	reader, err := s.localCache.Reader(ctx, rnCopy.ToProto(), 0, 0)
-	if err != nil {
-		return status.InternalErrorf("reading chunk %s from local cache: %s", rn.DownloadString(), err)
-	}
-	defer reader.Close()
-
-	_, _, err = cachetools.UploadFromReaderWithCompression(ctx, s.remote, rnCopy, reader, repb.Compressor_ZSTD)
-	if err != nil {
-		return status.InternalErrorf("uploading chunk %s to remote: %s", rn.DownloadString(), err)
+	for _, r := range rsp.GetResponses() {
+		if c := r.GetStatus().GetCode(); c != 0 {
+			return gstatus.Errorf(codes.Code(c), "batch upload chunk %s: %s", r.GetDigest().GetHash(), r.GetStatus().GetMessage())
+		}
 	}
 	return nil
 }