feat: create grpcZstd blob access

Author: tyler-french

pkg/blobstore/configuration/cas_blob_access_creator.go

@@ -99,6 +99,15 @@ func (bac *casBlobAccessCreator) NewCustomBlobAccess(terminationGroup program.Gr
 			BlobAccess:      grpcclients.NewCASBlobAccess(client, uuid.NewRandom, 65536),
 			DigestKeyFormat: digest.KeyWithInstance,
 		}, "grpc", nil
+	case *pb.BlobAccessConfiguration_Compressed:
+		client, err := bac.grpcClientFactory.NewClientFromConfiguration(backend.Compressed.Grpc, terminationGroup)
+		if err != nil {
+			return BlobAccessInfo{}, "", err
+		}
+		return BlobAccessInfo{
+			BlobAccess:      grpcclients.NewCASWithZstdBlobAccess(client, uuid.NewRandom, 65536, backend.Compressed.CompressionThresholdBytes),
+			DigestKeyFormat: digest.KeyWithInstance,
+		}, "compressed", nil
 	case *pb.BlobAccessConfiguration_ReferenceExpanding:
 		// The backend used by ReferenceExpandingBlobAccess is
 		// an Indirect Content Addressable Storage (ICAS). This

pkg/blobstore/grpcclients/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "ac_blob_access.go",
         "cas_blob_access.go",
+        "cas_zstd_blob_access.go",
         "fsac_blob_access.go",
         "icas_blob_access.go",
         "iscc_blob_access.go",
@@ -22,6 +23,7 @@ go_library(
         "//pkg/util",
         "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@com_github_google_uuid//:uuid",
+        "@com_github_klauspost_compress//zstd",
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
@@ -32,7 +34,10 @@ go_library(
 
 go_test(
     name = "grpcclients_test",
-    srcs = ["cas_blob_access_test.go"],
+    srcs = [
+        "cas_blob_access_test.go",
+        "cas_zstd_blob_access_test.go",
+    ],
     deps = [
         ":grpcclients",
         "//internal/mock",
@@ -43,6 +48,7 @@ go_test(
         "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto",
         "@bazel_remote_apis//build/bazel/semver:semver_go_proto",
         "@com_github_google_uuid//:uuid",
+        "@com_github_klauspost_compress//zstd",
         "@com_github_stretchr_testify//require",
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
         "@org_golang_google_grpc//:grpc",

pkg/blobstore/grpcclients/cas_blob_access.go

@@ -140,6 +140,10 @@ func (ba *casBlobAccess) Put(ctx context.Context, digest digest.Digest, b buffer
 }
 
 func (ba *casBlobAccess) FindMissing(ctx context.Context, digests digest.Set) (digest.Set, error) {
+	return findMissingBlobsInternal(ctx, digests, ba.contentAddressableStorageClient)
+}
+
+func findMissingBlobsInternal(ctx context.Context, digests digest.Set, cas remoteexecution.ContentAddressableStorageClient) (digest.Set, error) {
 	// Partition all digests by digest function, as the
 	// FindMissingBlobs() RPC can only process digests for a single
 	// instance name and digest function.
@@ -157,7 +161,7 @@ func (ba *casBlobAccess) FindMissing(ctx context.Context, digests digest.Set) (d
 			BlobDigests:    blobDigests,
 			DigestFunction: digestFunction.GetEnumValue(),
 		}
-		response, err := ba.contentAddressableStorageClient.FindMissingBlobs(ctx, &request)
+		response, err := cas.FindMissingBlobs(ctx, &request)
 		if err != nil {
 			return digest.EmptySet, err
 		}

pkg/blobstore/grpcclients/cas_blob_access_test.go

@@ -163,6 +163,107 @@ func TestCASBlobAccessPut(t *testing.T) {
 	})
 }
 
+func TestCASBlobAccessGet(t *testing.T) {
+	ctrl, ctx := gomock.WithContext(context.Background(), t)
+
+	client := mock.NewMockClientConnInterface(ctrl)
+	uuidGenerator := mock.NewMockUUIDGenerator(ctrl)
+	blobAccess := grpcclients.NewCASBlobAccess(client, uuidGenerator.Call, 10)
+
+	t.Run("Success", func(t *testing.T) {
+		blobDigest := digest.MustNewDigest("hello", remoteexecution.DigestFunction_MD5, "8b1a9953c4611296a827abf8c47804d7", 5)
+
+		clientStream := mock.NewMockClientStream(ctrl)
+		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Read").
+			Return(clientStream, nil)
+		clientStream.EXPECT().SendMsg(testutil.EqProto(t, &bytestream.ReadRequest{
+			ResourceName: "hello/blobs/8b1a9953c4611296a827abf8c47804d7/5",
+			ReadOffset:   0,
+			ReadLimit:    0,
+		})).Return(nil)
+		clientStream.EXPECT().RecvMsg(gomock.Any()).DoAndReturn(func(m interface{}) error {
+			resp := m.(*bytestream.ReadResponse)
+			resp.Data = []byte("Hello")
+			return nil
+		})
+		clientStream.EXPECT().RecvMsg(gomock.Any()).Return(io.EOF).AnyTimes()
+		clientStream.EXPECT().CloseSend().Return(nil)
+
+		buffer := blobAccess.Get(ctx, blobDigest)
+		data, err := buffer.ToByteSlice(1000)
+		require.NoError(t, err)
+		require.Equal(t, []byte("Hello"), data)
+	})
+
+	t.Run("SuccessLargeBlob", func(t *testing.T) {
+		// Create large blob data (1000 bytes)
+		expectedData := make([]byte, 1000)
+		for i := range expectedData {
+			expectedData[i] = byte('A' + (i % 26)) // Repeating alphabet pattern
+		}
+		largeDigest := digest.MustNewDigest("hello", remoteexecution.DigestFunction_MD5, "1411ffd5854fa029dc4d231aa89311eb", 1000)
+
+		clientStream := mock.NewMockClientStream(ctrl)
+		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Read").
+			Return(clientStream, nil)
+		clientStream.EXPECT().SendMsg(testutil.EqProto(t, &bytestream.ReadRequest{
+			ResourceName: "hello/blobs/1411ffd5854fa029dc4d231aa89311eb/1000",
+			ReadOffset:   0,
+			ReadLimit:    0,
+		})).Return(nil)
+
+		// Send data in a single chunk (simpler for testing)
+		clientStream.EXPECT().RecvMsg(gomock.Any()).DoAndReturn(func(m interface{}) error {
+			resp := m.(*bytestream.ReadResponse)
+			resp.Data = expectedData
+			return nil
+		})
+		clientStream.EXPECT().RecvMsg(gomock.Any()).Return(io.EOF).AnyTimes()
+		clientStream.EXPECT().CloseSend().Return(nil)
+
+		buffer := blobAccess.Get(ctx, largeDigest)
+		data, err := buffer.ToByteSlice(1500)
+		require.NoError(t, err)
+		require.Equal(t, expectedData, data)
+	})
+
+	t.Run("InitialFailure", func(t *testing.T) {
+		blobDigest := digest.MustNewDigest("hello", remoteexecution.DigestFunction_MD5, "8b1a9953c4611296a827abf8c47804d7", 5)
+
+		// Failure to create the outgoing connection.
+		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Read").
+			Return(nil, status.Error(codes.Internal, "Failed to create outgoing connection"))
+
+		buffer := blobAccess.Get(ctx, blobDigest)
+		_, err := buffer.ToByteSlice(1000)
+		testutil.RequireEqualStatus(t,
+			status.Error(codes.Internal, "Failed to create outgoing connection"),
+			err)
+	})
+
+	t.Run("ReceiveFailure", func(t *testing.T) {
+		blobDigest := digest.MustNewDigest("hello", remoteexecution.DigestFunction_MD5, "8b1a9953c4611296a827abf8c47804d7", 5)
+
+		// Failure to receive a response.
+		clientStream := mock.NewMockClientStream(ctrl)
+		client.EXPECT().NewStream(gomock.Any(), gomock.Any(), "/google.bytestream.ByteStream/Read").
+			Return(clientStream, nil)
+		clientStream.EXPECT().SendMsg(testutil.EqProto(t, &bytestream.ReadRequest{
+			ResourceName: "hello/blobs/8b1a9953c4611296a827abf8c47804d7/5",
+			ReadOffset:   0,
+			ReadLimit:    0,
+		})).Return(nil)
+		clientStream.EXPECT().RecvMsg(gomock.Any()).Return(status.Error(codes.Internal, "Lost connection to server")).AnyTimes()
+		clientStream.EXPECT().CloseSend().Return(nil)
+
+		buffer := blobAccess.Get(ctx, blobDigest)
+		_, err := buffer.ToByteSlice(1000)
+		testutil.RequireEqualStatus(t,
+			status.Error(codes.Internal, "Lost connection to server"),
+			err)
+	})
+}
+
 func TestCASBlobAccessGetCapabilities(t *testing.T) {
 	ctrl, ctx := gomock.WithContext(context.Background(), t)