Add a generic message reader interface

cmd/bb_copy/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
     importpath = "github.com/buildbarn/bb-storage/cmd/bb_copy",
     visibility = ["//visibility:private"],
     deps = [
+        "//pkg/blobstore",
         "//pkg/blobstore/configuration",
         "//pkg/blobstore/replication",
         "//pkg/digest",
@@ -15,6 +16,7 @@ go_library(
         "//pkg/proto/configuration/bb_copy",
         "//pkg/util",
         "//pkg/zstd",
+        "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
     ],

cmd/bb_copy/main.go

@@ -2,8 +2,11 @@ package main
 
 import (
 	"context"
+	"io"
 	"os"
 
+	remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
+	"github.com/buildbarn/bb-storage/pkg/blobstore"
 	blobstore_configuration "github.com/buildbarn/bb-storage/pkg/blobstore/configuration"
 	"github.com/buildbarn/bb-storage/pkg/blobstore/replication"
 	"github.com/buildbarn/bb-storage/pkg/digest"
@@ -30,6 +33,14 @@ import (
 // ZIPWritingBlobAccess, this tool can also be used to backup and
 // restore parts of the Content Addressable Storage.
 
+type blobAccessStreamReader struct {
+	blobAccess blobstore.BlobAccess
+}
+
+func (r *blobAccessStreamReader) ReadStream(ctx context.Context, d digest.Digest) (io.ReadCloser, error) {
+	return r.blobAccess.Get(ctx, d).ToReader(), nil
+}
+
 func main() {
 	program.RunMain(func(ctx context.Context, siblingsGroup, dependenciesGroup program.Group) error {
 		if len(os.Args) != 2 {
@@ -75,8 +86,11 @@ func main() {
 		}
 		nestedReplicator := replication.NewNestedBlobReplicator(
 			replicator,
-			sink.DigestKeyFormat,
 			int(configuration.MaximumMessageSizeBytes),
+			blobstore.NewBlobAccessMessageReader(source.BlobAccess, int(configuration.MaximumMessageSizeBytes), func() *remoteexecution.Action { return &remoteexecution.Action{} }),
+			blobstore.NewBlobAccessMessageReader(source.BlobAccess, int(configuration.MaximumMessageSizeBytes), func() *remoteexecution.Directory { return &remoteexecution.Directory{} }),
+			&blobAccessStreamReader{blobAccess: source.BlobAccess},
+			sink.DigestKeyFormat,
 		)
 
 		instanceName, err := digest.NewInstanceName(configuration.InstanceName)

internal/mock/BUILD.bazel

@@ -76,7 +76,10 @@ gomock(
 gomock(
     name = "blobstore_replication",
     out = "blobstore_replication.go",
-    interfaces = ["BlobReplicator"],
+    interfaces = [
+        "BlobReplicator",
+        "BlobStreamReader",
+    ],
     library = "//pkg/blobstore/replication",
     mockgen_model_library = "@org_uber_go_mock//mockgen/model",
     mockgen_tool = "@org_uber_go_mock//mockgen",
@@ -320,6 +323,16 @@ gomock(
     package = "mock",
 )
 
+gomock(
+    name = "cas",
+    out = "cas.go",
+    interfaces = ["MessageReader"],
+    library = "//pkg/cas",
+    mockgen_tool = "@org_uber_go_mock//mockgen",
+    package = "mock",
+    source = "//pkg/cas:message_reader.go",
+)
+
 gomock(
     name = "trace",
     out = "trace.go",
@@ -367,6 +380,7 @@ go_library(
         "buffer.go",
         "builder.go",
         "capabilities.go",
+        "cas.go",
         "clock.go",
         "cloud_aws.go",
         "cloud_gcp.go",
@@ -416,6 +430,7 @@ go_library(
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_protobuf//encoding/protowire",
+        "@org_golang_google_protobuf//proto",
         "@org_uber_go_mock//gomock",
     ],
 )

pkg/blobstore/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
         "action_result_timestamp_injecting_blob_access.go",
         "authorizing_blob_access.go",
         "blob_access.go",
+        "blob_access_message_reader.go",
         "cas_read_buffer_factory.go",
         "deadline_enforcing_blob_access.go",
         "demultiplexing_blob_access.go",
@@ -34,6 +35,7 @@ go_library(
         "//pkg/blobstore/buffer",
         "//pkg/blobstore/slicing",
         "//pkg/capabilities",
+        "//pkg/cas",
         "//pkg/clock",
         "//pkg/cloud/aws",
         "//pkg/cloud/gcp",

pkg/blobstore/blob_access_message_reader.go

@@ -0,0 +1,35 @@
+package blobstore
+
+import (
+	"context"
+
+	"github.com/buildbarn/bb-storage/pkg/cas"
+	"github.com/buildbarn/bb-storage/pkg/digest"
+	"google.golang.org/protobuf/proto"
+)
+
+type blobAccessMessageReader[T proto.Message] struct {
+	blobAccess              BlobAccess
+	maximumMessageSizeBytes int
+	allocator               func() T
+}
+
+// NewBlobAccessMessageReader creates a storage.MessageReader that reads
+// a message from a BlobAccess up to maximumMessageSizeBytes large.
+// Requires an allocator that returns allocated objects of type T.
+func NewBlobAccessMessageReader[T proto.Message](blobAccess BlobAccess, maximumMessageSizeBytes int, allocator func() T) cas.MessageReader[T] {
+	return blobAccessMessageReader[T]{
+		blobAccess:              blobAccess,
+		maximumMessageSizeBytes: maximumMessageSizeBytes,
+		allocator:               allocator,
+	}
+}
+
+func (mr blobAccessMessageReader[T]) ReadMessage(ctx context.Context, d digest.Digest) (T, error) {
+	var zero T
+	ret, err := mr.blobAccess.Get(ctx, d).ToProto(mr.allocator(), mr.maximumMessageSizeBytes)
+	if err != nil {
+		return zero, err
+	}
+	return ret.(T), nil
+}

pkg/blobstore/replication/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//pkg/blobstore",
         "//pkg/blobstore/buffer",
         "//pkg/blobstore/slicing",
+        "//pkg/cas",
         "//pkg/clock",
         "//pkg/digest",
         "//pkg/proto/replicator",

pkg/blobstore/replication/nested_blob_replicator.go

@@ -8,25 +8,40 @@ import (
 	remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
 	"github.com/buildbarn/bb-storage/pkg/blobstore"
 	"github.com/buildbarn/bb-storage/pkg/blobstore/buffer"
+	"github.com/buildbarn/bb-storage/pkg/cas"
 	"github.com/buildbarn/bb-storage/pkg/digest"
 	"github.com/buildbarn/bb-storage/pkg/util"
 
 	"google.golang.org/protobuf/encoding/protowire"
 )
 
+// CASReplicator replicates blobs to the destination Content
+// Addressable Storage.
+type CASReplicator interface {
+	ReplicateMultiple(ctx context.Context, digests digest.Set) error
+}
+
+// BlobStreamReader provides a stream to the raw bytes of a blob.
+type BlobStreamReader interface {
+	ReadStream(ctx context.Context, d digest.Digest) (io.ReadCloser, error)
+}
+
 type blobToReplicate struct {
 	digest       digest.Digest
-	expanderFunc func(ctx context.Context, b buffer.Buffer) error
+	expanderFunc func(ctx context.Context, d digest.Digest) error
 }
 
 // NestedBlobReplicator is a helper type for BlobReplicator that can be
 // used to copy nested hierarchies of objects stored in the Content
 // Addressable Storage (CAS). In the case of the REv2 protocol, these
 // are Action, Directory and Tree messages.
 type NestedBlobReplicator struct {
-	replicator              BlobReplicator
+	replicator              CASReplicator
 	digestKeyFormat         digest.KeyFormat
 	maximumMessageSizeBytes int
+	actionReader            cas.MessageReader[*remoteexecution.Action]
+	directoryReader         cas.MessageReader[*remoteexecution.Directory]
+	treeReader              BlobStreamReader
 
 	lock             sync.Mutex
 	blobsSeen        map[string]struct{}
@@ -37,17 +52,20 @@ type NestedBlobReplicator struct {
 
 // NewNestedBlobReplicator creates a new NestedBlobReplicator that does
 // not have any objects to be replicated queued.
-func NewNestedBlobReplicator(replicator BlobReplicator, digestKeyFormat digest.KeyFormat, maximumMessageSizeBytes int) *NestedBlobReplicator {
+func NewNestedBlobReplicator(replicator CASReplicator, maximumMessageSizeBytes int, actionReader cas.MessageReader[*remoteexecution.Action], directoryReader cas.MessageReader[*remoteexecution.Directory], treeReader BlobStreamReader, digestKeyFormat digest.KeyFormat) *NestedBlobReplicator {
 	return &NestedBlobReplicator{
 		replicator:              replicator,
-		digestKeyFormat:         digestKeyFormat,
 		maximumMessageSizeBytes: maximumMessageSizeBytes,
+		digestKeyFormat:         digestKeyFormat,
+		actionReader:            actionReader,
+		directoryReader:         directoryReader,
+		treeReader:              treeReader,
 
 		blobsSeen: map[string]struct{}{},
 	}
 }
 
-func (nr *NestedBlobReplicator) enqueue(blobDigest digest.Digest, expanderFunc func(ctx context.Context, b buffer.Buffer) error) {
+func (nr *NestedBlobReplicator) enqueue(blobDigest digest.Digest, expanderFunc func(ctx context.Context, d digest.Digest) error) {
 	nr.lock.Lock()
 	defer nr.lock.Unlock()
 
@@ -73,12 +91,11 @@ func (nr *NestedBlobReplicator) maybeWakeUpLocked() {
 // referenced input root and Command message will be replicated as well.
 func (nr *NestedBlobReplicator) EnqueueAction(actionDigest digest.Digest) {
 	digestFunction := actionDigest.GetDigestFunction()
-	nr.enqueue(actionDigest, func(ctx context.Context, b buffer.Buffer) error {
-		actionMessage, err := b.ToProto(&remoteexecution.Action{}, nr.maximumMessageSizeBytes)
+	nr.enqueue(actionDigest, func(ctx context.Context, d digest.Digest) error {
+		action, err := nr.actionReader.ReadMessage(ctx, d)
 		if err != nil {
 			return err
 		}
-		action := actionMessage.(*remoteexecution.Action)
 
 		inputRootDigest, err := digestFunction.NewDigestFromProto(action.InputRootDigest)
 		if err != nil {
@@ -102,12 +119,11 @@ func (nr *NestedBlobReplicator) EnqueueAction(actionDigest digest.Digest) {
 // well, recursively.
 func (nr *NestedBlobReplicator) EnqueueDirectory(directoryDigest digest.Digest) {
 	digestFunction := directoryDigest.GetDigestFunction()
-	nr.enqueue(directoryDigest, func(ctx context.Context, b buffer.Buffer) error {
-		directoryMessage, err := b.ToProto(&remoteexecution.Directory{}, nr.maximumMessageSizeBytes)
+	nr.enqueue(directoryDigest, func(ctx context.Context, d digest.Digest) error {
+		directory, err := nr.directoryReader.ReadMessage(ctx, d)
 		if err != nil {
 			return err
 		}
-		directory := directoryMessage.(*remoteexecution.Directory)
 
 		for i, childDirectory := range directory.Directories {
 			childDigest, err := digestFunction.NewDigestFromProto(childDirectory.Digest)
@@ -136,8 +152,11 @@ func (nr *NestedBlobReplicator) EnqueueDirectory(directoryDigest digest.Digest)
 // file will be replicated as well.
 func (nr *NestedBlobReplicator) EnqueueTree(treeDigest digest.Digest) {
 	digestFunction := treeDigest.GetDigestFunction()
-	nr.enqueue(treeDigest, func(ctx context.Context, b buffer.Buffer) error {
-		r := b.ToReader()
+	nr.enqueue(treeDigest, func(ctx context.Context, d digest.Digest) error {
+		r, err := nr.treeReader.ReadStream(ctx, d)
+		if err != nil {
+			return err
+		}
 		defer r.Close()
 
 		// Gather digests of files contained in the directories.
@@ -214,9 +233,10 @@ func (nr *NestedBlobReplicator) Replicate(ctx context.Context) error {
 		// Replicate a single object.
 		nr.blobsReplicating++
 		nr.lock.Unlock()
+		nr.replicator.ReplicateMultiple(ctx, blobToReplicate.digest.ToSingletonSet())
 		err := blobToReplicate.expanderFunc(
 			ctx,
-			nr.replicator.ReplicateSingle(ctx, blobToReplicate.digest),
+			blobToReplicate.digest,
 		)
 		nr.lock.Lock()
 		nr.blobsReplicating--