Set the resource-name header as part of Bytestream requests

Even though I don't think we should do anything on the Buildbarn side to
attempt to interpret these headers, let's at least emit it as part of
outgoing Bytestream requests. That's likely going to make life easier
for people who want to do load balancing this way.

More details: bazelbuild/remote-apis#303

Author: EdSchouten

pkg/blobstore/grpcclients/BUILD.bazel

@@ -25,6 +25,7 @@ go_library(
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
+        "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
     ],
 )

pkg/blobstore/grpcclients/cas_blob_access.go

@@ -14,6 +14,7 @@ import (
 
 	"google.golang.org/genproto/googleapis/bytestream"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 type casBlobAccess struct {
@@ -61,11 +62,17 @@ func (r *byteStreamChunkReader) Close() {
 	}
 }
 
+const resourceNameHeader = "build.bazel.remote.execution.v2.resource-name"
+
 func (ba *casBlobAccess) Get(ctx context.Context, digest digest.Digest) buffer.Buffer {
 	ctxWithCancel, cancel := context.WithCancel(ctx)
-	client, err := ba.byteStreamClient.Read(ctxWithCancel, &bytestream.ReadRequest{
-		ResourceName: digest.GetByteStreamReadPath(remoteexecution.Compressor_IDENTITY),
-	})
+	resourceName := digest.GetByteStreamReadPath(remoteexecution.Compressor_IDENTITY)
+	client, err := ba.byteStreamClient.Read(
+		metadata.AppendToOutgoingContext(ctxWithCancel, resourceNameHeader, resourceName),
+		&bytestream.ReadRequest{
+			ResourceName: resourceName,
+		},
+	)
 	if err != nil {
 		cancel()
 		return buffer.NewBufferFromError(err)
@@ -86,13 +93,15 @@ func (ba *casBlobAccess) Put(ctx context.Context, digest digest.Digest, b buffer
 	defer r.Close()
 
 	ctxWithCancel, cancel := context.WithCancel(ctx)
-	client, err := ba.byteStreamClient.Write(ctxWithCancel)
+	resourceName := digest.GetByteStreamWritePath(uuid.Must(ba.uuidGenerator()), remoteexecution.Compressor_IDENTITY)
+	client, err := ba.byteStreamClient.Write(
+		metadata.AppendToOutgoingContext(ctxWithCancel, resourceNameHeader, resourceName),
+	)
 	if err != nil {
 		cancel()
 		return err
 	}
 
-	resourceName := digest.GetByteStreamWritePath(uuid.Must(ba.uuidGenerator()), remoteexecution.Compressor_IDENTITY)
 	writeOffset := int64(0)
 	for {
 		if data, err := r.Read(); err == nil {

pkg/blobstore/grpcclients/cas_blob_access_test.go

@@ -36,6 +36,7 @@ func TestCASBlobAccessPut(t *testing.T) {
 
 	t.Run("InitialFailure", func(t *testing.T) {
 		// Failure to create the outgoing connection.
+		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Write").
 			Return(nil, status.Error(codes.Internal, "Failed to create outgoing connection"))
 		r := mock.NewMockFileReader(ctrl)
@@ -52,12 +53,12 @@ func TestCASBlobAccessPut(t *testing.T) {
 		// should be returned.
 		clientStream := mock.NewMockClientStream(ctrl)
 		var savedCtx context.Context
+		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Write").
 			DoAndReturn(func(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 				savedCtx = ctx
 				return clientStream, nil
 			})
-		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		r := mock.NewMockFileReader(ctrl)
 		r.EXPECT().ReadAt(gomock.Len(5), int64(0)).Return(0, status.Error(codes.Internal, "Disk on fire"))
 		clientStream.EXPECT().CloseSend().DoAndReturn(func() error {
@@ -79,9 +80,9 @@ func TestCASBlobAccessPut(t *testing.T) {
 		// error message that is returned by
 		// ClientStream.CloseSend().
 		clientStream := mock.NewMockClientStream(ctrl)
+		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Write").
 			Return(clientStream, nil)
-		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		r := mock.NewMockFileReader(ctrl)
 		r.EXPECT().ReadAt(gomock.Len(5), int64(0)).DoAndReturn(func(p []byte, off int64) (int, error) {
 			copy(p, "Hello")
@@ -104,9 +105,9 @@ func TestCASBlobAccessPut(t *testing.T) {
 		// Similar to the previous test, ClientStream.SendMsg()
 		// may fail with io.EOF for the final call.
 		clientStream := mock.NewMockClientStream(ctrl)
+		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Write").
 			Return(clientStream, nil)
-		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		r := mock.NewMockFileReader(ctrl)
 		r.EXPECT().ReadAt(gomock.Len(5), int64(0)).DoAndReturn(func(p []byte, off int64) (int, error) {
 			copy(p, "Hello")
@@ -135,9 +136,9 @@ func TestCASBlobAccessPut(t *testing.T) {
 		// ClientStream.CloseSend() still fails. The error must
 		// still be propagated.
 		clientStream := mock.NewMockClientStream(ctrl)
+		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Write").
 			Return(clientStream, nil)
-		uuidGenerator.EXPECT().Call().Return(uuid, nil)
 		r := mock.NewMockFileReader(ctrl)
 		r.EXPECT().ReadAt(gomock.Len(5), int64(0)).DoAndReturn(func(p []byte, off int64) (int, error) {
 			copy(p, "Hello")