From 2f76f2641fc64aaf9584b865b2eea9823c874ab1 Mon Sep 17 00:00:00 2001 From: Jeremy Volkman Date: Fri, 1 May 2026 17:53:18 +0000 Subject: [PATCH] feat: add optional Platform to action cache Command and Actions. Allow specifying static platform information for RBE services that require Command and/or Action resource to have it. --- pkg/configuration/new_asset_store.go | 2 +- .../configuration/bb_remote_asset/BUILD.bazel | 2 + .../bb_remote_asset/bb_remote_asset.pb.go | 34 ++++--- .../bb_remote_asset/bb_remote_asset.proto | 5 ++ pkg/storage/action_cache_asset_store.go | 10 ++- pkg/storage/action_cache_asset_store_test.go | 88 +++++++++++++++++-- 6 files changed, 119 insertions(+), 22 deletions(-) diff --git a/pkg/configuration/new_asset_store.go b/pkg/configuration/new_asset_store.go index 72e2c8d..d1dedf5 100644 --- a/pkg/configuration/new_asset_store.go +++ b/pkg/configuration/new_asset_store.go @@ -52,7 +52,7 @@ func NewAssetStoreFromConfiguration( if err != nil { return nil, err } - assetStore = storage.NewActionCacheAssetStore(actionCache.BlobAccess, contentAddressableStorage.BlobAccess, maximumMessageSizeBytes) + assetStore = storage.NewActionCacheAssetStore(actionCache.BlobAccess, contentAddressableStorage.BlobAccess, maximumMessageSizeBytes, configuration.ActionCachePlatform) default: return nil, status.Errorf(codes.InvalidArgument, "Asset Cache configuration is invalid as no supported Asset Cache is defined.") } diff --git a/pkg/proto/configuration/bb_remote_asset/BUILD.bazel b/pkg/proto/configuration/bb_remote_asset/BUILD.bazel index bed9633..cfbc19c 100644 --- a/pkg/proto/configuration/bb_remote_asset/BUILD.bazel +++ b/pkg/proto/configuration/bb_remote_asset/BUILD.bazel @@ -9,6 +9,7 @@ proto_library( visibility = ["//visibility:public"], deps = [ "//pkg/proto/configuration/bb_remote_asset/fetch:fetch_proto", + "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_proto", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/auth:auth_proto", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/blobstore:blobstore_proto", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/global:global_proto", @@ -24,6 +25,7 @@ go_proto_library( visibility = ["//visibility:public"], deps = [ "//pkg/proto/configuration/bb_remote_asset/fetch", + "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/auth", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/blobstore", "@com_github_buildbarn_bb_storage//pkg/proto/configuration/global", diff --git a/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.pb.go b/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.pb.go index 3cc1166..509612f 100644 --- a/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.pb.go +++ b/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.pb.go @@ -7,6 +7,7 @@ package bb_remote_asset import ( + v2 "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2" fetch "github.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/fetch" auth "github.com/buildbarn/bb-storage/pkg/proto/configuration/auth" blobstore "github.com/buildbarn/bb-storage/pkg/proto/configuration/blobstore" @@ -149,9 +150,10 @@ type AssetCacheConfiguration struct { // // *AssetCacheConfiguration_BlobAccess // *AssetCacheConfiguration_ActionCache - Backend isAssetCacheConfiguration_Backend `protobuf_oneof:"backend"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + Backend isAssetCacheConfiguration_Backend `protobuf_oneof:"backend"` + ActionCachePlatform *v2.Platform `protobuf:"bytes,3,opt,name=action_cache_platform,json=actionCachePlatform,proto3" json:"action_cache_platform,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *AssetCacheConfiguration) Reset() { @@ -209,6 +211,13 @@ func (x *AssetCacheConfiguration) GetActionCache() *blobstore.BlobAccessConfigur return nil } +func (x *AssetCacheConfiguration) GetActionCachePlatform() *v2.Platform { + if x != nil { + return x.ActionCachePlatform + } + return nil +} + type isAssetCacheConfiguration_Backend interface { isAssetCacheConfiguration_Backend() } @@ -229,7 +238,7 @@ var File_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remote_ const file_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remote_asset_bb_remote_asset_proto_rawDesc = "" + "\n" + - "bgithub.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto\x12'buildbarn.configuration.bb_remote_asset\x1a`github.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/fetch/fetcher.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/auth/auth.proto\x1aQgithub.com/buildbarn/bb-storage/pkg/proto/configuration/blobstore/blobstore.proto\x1aKgithub.com/buildbarn/bb-storage/pkg/proto/configuration/global/global.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/grpc/grpc.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/zstd/zstd.proto\"\x8d\a\n" + + "bgithub.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto\x12'buildbarn.configuration.bb_remote_asset\x1a6build/bazel/remote/execution/v2/remote_execution.proto\x1a`github.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/fetch/fetcher.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/auth/auth.proto\x1aQgithub.com/buildbarn/bb-storage/pkg/proto/configuration/blobstore/blobstore.proto\x1aKgithub.com/buildbarn/bb-storage/pkg/proto/configuration/global/global.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/grpc/grpc.proto\x1aGgithub.com/buildbarn/bb-storage/pkg/proto/configuration/zstd/zstd.proto\"\x8d\a\n" + "\x18ApplicationConfiguration\x12T\n" + "\fgrpc_servers\x18\x03 \x03(\v21.buildbarn.configuration.grpc.ServerConfigurationR\vgrpcServers\x12z\n" + "\x1bcontent_addressable_storage\x18\x04 \x01(\v2:.buildbarn.configuration.blobstore.BlobAccessConfigurationR\x19contentAddressableStorage\x12;\n" + @@ -242,11 +251,12 @@ const file_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remot "\x10fetch_authorizer\x18\n" + " \x01(\v25.buildbarn.configuration.auth.AuthorizerConfigurationR\x0ffetchAuthorizer\x12^\n" + "\x0fpush_authorizer\x18\v \x01(\v25.buildbarn.configuration.auth.AuthorizerConfigurationR\x0epushAuthorizer\x12L\n" + - "\tzstd_pool\x18\f \x01(\v2/.buildbarn.configuration.zstd.PoolConfigurationR\bzstdPoolJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03\"\xe4\x01\n" + + "\tzstd_pool\x18\f \x01(\v2/.buildbarn.configuration.zstd.PoolConfigurationR\bzstdPoolJ\x04\b\x01\x10\x02J\x04\b\x02\x10\x03\"\xc3\x02\n" + "\x17AssetCacheConfiguration\x12]\n" + "\vblob_access\x18\x01 \x01(\v2:.buildbarn.configuration.blobstore.BlobAccessConfigurationH\x00R\n" + "blobAccess\x12_\n" + - "\faction_cache\x18\x02 \x01(\v2:.buildbarn.configuration.blobstore.BlobAccessConfigurationH\x00R\vactionCacheB\t\n" + + "\faction_cache\x18\x02 \x01(\v2:.buildbarn.configuration.blobstore.BlobAccessConfigurationH\x00R\vactionCache\x12]\n" + + "\x15action_cache_platform\x18\x03 \x01(\v2).build.bazel.remote.execution.v2.PlatformR\x13actionCachePlatformB\t\n" + "\abackendBNZLgithub.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_assetb\x06proto3" var ( @@ -271,6 +281,7 @@ var file_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remote_ (*fetch.FetcherConfiguration)(nil), // 5: buildbarn.configuration.bb_remote_asset.fetch.FetcherConfiguration (*auth.AuthorizerConfiguration)(nil), // 6: buildbarn.configuration.auth.AuthorizerConfiguration (*zstd.PoolConfiguration)(nil), // 7: buildbarn.configuration.zstd.PoolConfiguration + (*v2.Platform)(nil), // 8: build.bazel.remote.execution.v2.Platform } var file_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remote_asset_bb_remote_asset_proto_depIdxs = []int32{ 2, // 0: buildbarn.configuration.bb_remote_asset.ApplicationConfiguration.grpc_servers:type_name -> buildbarn.configuration.grpc.ServerConfiguration @@ -283,11 +294,12 @@ var file_github_com_buildbarn_bb_remote_asset_pkg_proto_configuration_bb_remote_ 7, // 7: buildbarn.configuration.bb_remote_asset.ApplicationConfiguration.zstd_pool:type_name -> buildbarn.configuration.zstd.PoolConfiguration 3, // 8: buildbarn.configuration.bb_remote_asset.AssetCacheConfiguration.blob_access:type_name -> buildbarn.configuration.blobstore.BlobAccessConfiguration 3, // 9: buildbarn.configuration.bb_remote_asset.AssetCacheConfiguration.action_cache:type_name -> buildbarn.configuration.blobstore.BlobAccessConfiguration - 10, // [10:10] is the sub-list for method output_type - 10, // [10:10] is the sub-list for method input_type - 10, // [10:10] is the sub-list for extension type_name - 10, // [10:10] is the sub-list for extension extendee - 0, // [0:10] is the sub-list for field type_name + 8, // 10: buildbarn.configuration.bb_remote_asset.AssetCacheConfiguration.action_cache_platform:type_name -> build.bazel.remote.execution.v2.Platform + 11, // [11:11] is the sub-list for method output_type + 11, // [11:11] is the sub-list for method input_type + 11, // [11:11] is the sub-list for extension type_name + 11, // [11:11] is the sub-list for extension extendee + 0, // [0:11] is the sub-list for field type_name } func init() { diff --git a/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto b/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto index 996ebea..1374c4f 100644 --- a/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto +++ b/pkg/proto/configuration/bb_remote_asset/bb_remote_asset.proto @@ -2,6 +2,7 @@ syntax = "proto3"; package buildbarn.configuration.bb_remote_asset; +import "build/bazel/remote/execution/v2/remote_execution.proto"; import "github.com/buildbarn/bb-remote-asset/pkg/proto/configuration/bb_remote_asset/fetch/fetcher.proto"; import "github.com/buildbarn/bb-storage/pkg/proto/configuration/auth/auth.proto"; import "github.com/buildbarn/bb-storage/pkg/proto/configuration/blobstore/blobstore.proto"; @@ -66,4 +67,8 @@ message AssetCacheConfiguration { // Cache assets in an existing action cache buildbarn.configuration.blobstore.BlobAccessConfiguration action_cache = 2; } + + // Optional platform properties to attach to actions created for asset + // mappings. + build.bazel.remote.execution.v2.Platform action_cache_platform = 3; } diff --git a/pkg/storage/action_cache_asset_store.go b/pkg/storage/action_cache_asset_store.go index 3501833..38fa195 100644 --- a/pkg/storage/action_cache_asset_store.go +++ b/pkg/storage/action_cache_asset_store.go @@ -36,15 +36,17 @@ type actionCacheAssetStore struct { actionCache blobstore.BlobAccess contentAddressableStorage blobstore.BlobAccess maximumMessageSizeBytes int + platform *remoteexecution.Platform } // NewActionCacheAssetStore creates a new AssetStore which stores it's // references as ActionResults in the Action Cache. -func NewActionCacheAssetStore(actionCache, contentAddressableStorage blobstore.BlobAccess, maximumMessageSizeBytes int) AssetStore { +func NewActionCacheAssetStore(actionCache, contentAddressableStorage blobstore.BlobAccess, maximumMessageSizeBytes int, platform *remoteexecution.Platform) AssetStore { return &actionCacheAssetStore{ actionCache: actionCache, contentAddressableStorage: contentAddressableStorage, maximumMessageSizeBytes: maximumMessageSizeBytes, + platform: platform, } } @@ -88,6 +90,7 @@ func (rs *actionCacheAssetStore) assetReferenceToAction(ref *asset.AssetReferenc Arguments: ref.Uris, OutputPaths: []string{"out"}, OutputDirectoryFormat: remoteexecution.Command_TREE_AND_DIRECTORY, + Platform: rs.platform, } _, commandDigest, err := ProtoSerialise(command, digestFunction) if err != nil { @@ -97,10 +100,14 @@ func (rs *actionCacheAssetStore) assetReferenceToAction(ref *asset.AssetReferenc action = &remoteexecution.Action{ CommandDigest: commandDigest.GetProto(), InputRootDigest: directoryDigest.GetProto(), + Platform: rs.platform, } } else { // Generate a command based on the qualifiers command := commandGenerator(ref.Uris[0]) + if command.Platform == nil { + command.Platform = rs.platform + } _, commandDigest, err := ProtoSerialise(command, digestFunction) if err != nil { return nil, nil, err @@ -109,6 +116,7 @@ func (rs *actionCacheAssetStore) assetReferenceToAction(ref *asset.AssetReferenc action = &remoteexecution.Action{ CommandDigest: commandDigest.GetProto(), InputRootDigest: EmptyDigest(digestFunction).GetProto(), + Platform: command.Platform, } } diff --git a/pkg/storage/action_cache_asset_store_test.go b/pkg/storage/action_cache_asset_store_test.go index 57ae73c..787e68c 100644 --- a/pkg/storage/action_cache_asset_store_test.go +++ b/pkg/storage/action_cache_asset_store_test.go @@ -91,7 +91,77 @@ func TestActionCacheAssetStorePutBlob(t *testing.T) { } return status.Error(codes.Internal, "Blob digest not found") }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) + + err = assetStore.Put(ctx, assetRef, assetData, digestFunction) + require.NoError(t, err) +} + +func TestActionCacheAssetStorePutBlobWithPlatform(t *testing.T) { + ctrl, ctx := gomock.WithContext(context.Background(), t) + + instanceName := util.Must(digest.NewInstanceName("")) + digestFunction, err := instanceName.GetDigestFunction(remoteexecution.DigestFunction_SHA256, 0) + require.NoError(t, err) + + blobDigest := &remoteexecution.Digest{ + Hash: "58de0f27ce0f781e5c109f18b0ee6905bdf64f2b1009e225ac67a27f656a0643", + SizeBytes: 111, + } + uri := "https://example.com/example.txt" + assetRef := storage.NewAssetReference([]string{uri}, + []*remoteasset.Qualifier{{Name: "test", Value: "test"}}) + assetData := storage.NewBlobAsset(blobDigest, timestamppb.Now()) + + platform := &remoteexecution.Platform{ + Properties: []*remoteexecution.Platform_Property{ + {Name: "OSFamily", Value: "Linux"}, + }, + } + + // Compute expected digests + qr := storage.NewAssetReference(nil, assetRef.Qualifiers) + _, qrDigest, err := storage.ProtoSerialise(qr, digestFunction) + require.NoError(t, err) + + directory := &remoteexecution.Directory{ + Files: []*remoteexecution.FileNode{{ + Name: "AssetReference", + Digest: qrDigest.GetProto(), + }}, + } + _, directoryDigest, err := storage.ProtoSerialise(directory, digestFunction) + require.NoError(t, err) + + command := &remoteexecution.Command{ + Arguments: assetRef.Uris, + OutputPaths: []string{"out"}, + OutputDirectoryFormat: remoteexecution.Command_TREE_AND_DIRECTORY, + Platform: platform, + } + _, commandDigest, err := storage.ProtoSerialise(command, digestFunction) + require.NoError(t, err) + + action := &remoteexecution.Action{ + CommandDigest: commandDigest.GetProto(), + InputRootDigest: directoryDigest.GetProto(), + Platform: platform, + } + _, actionDigest, err := storage.ProtoSerialise(action, digestFunction) + require.NoError(t, err) + + ac := mock.NewMockBlobAccess(ctrl) + cas := mock.NewMockBlobAccess(ctrl) + + // Expect uploads to CAS with specific digests! + cas.EXPECT().Put(ctx, qrDigest, gomock.Any()).Return(nil) + cas.EXPECT().Put(ctx, directoryDigest, gomock.Any()).Return(nil) + cas.EXPECT().Put(ctx, commandDigest, gomock.Any()).Return(nil) + cas.EXPECT().Put(ctx, actionDigest, gomock.Any()).Return(nil) + + ac.EXPECT().Put(ctx, actionDigest, gomock.Any()).Return(nil) + + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, platform) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NoError(t, err) @@ -184,7 +254,7 @@ func TestActionCacheAssetStorePutDirectory(t *testing.T) { } return status.Error(codes.Internal, "Directory digest not found") }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NoError(t, err) @@ -230,7 +300,7 @@ func TestActionCacheAssetStorePutMalformedDirectory(t *testing.T) { }, buffer.UserProvided)) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NotNil(t, err) @@ -337,7 +407,7 @@ func TestActionCacheAssetStorePutRecursiveDirectory(t *testing.T) { return status.Error(codes.Internal, "Directory digest not found") }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NoError(t, err) @@ -386,7 +456,7 @@ func TestActionCacheAssetStorePutMalformedDirectoryAsBlob(t *testing.T) { } return status.Error(codes.Internal, "Directory digest not found") }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NoError(t, err) @@ -424,7 +494,7 @@ func roundTripTest(t *testing.T, assetRef *asset.AssetReference, assetData *asse return nil }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) err = assetStore.Put(ctx, assetRef, assetData, digestFunction) require.NoError(t, err) @@ -443,7 +513,7 @@ func roundTripTest(t *testing.T, assetRef *asset.AssetReference, assetData *asse return buffer.NewBufferFromError(fmt.Errorf("not in AC")) }) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) asset, err := assetStore.Get(ctx, assetRef, digestFunction) require.NoError(t, err) @@ -534,7 +604,7 @@ func TestActionCacheAssetStoreGetBlob(t *testing.T) { ac := mock.NewMockBlobAccess(ctrl) cas := mock.NewMockBlobAccess(ctrl) ac.EXPECT().Get(ctx, actionDigest).Return(buf) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) _, err = assetStore.Get(ctx, assetRef, digestFunction) require.NoError(t, err) @@ -576,7 +646,7 @@ func TestActionCacheAssetStoreGetDirectory(t *testing.T) { ac := mock.NewMockBlobAccess(ctrl) cas := mock.NewMockBlobAccess(ctrl) ac.EXPECT().Get(ctx, actionDigest).Return(buf) - assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024) + assetStore := storage.NewActionCacheAssetStore(ac, cas, 16*1024*1024, nil) _, err = assetStore.Get(ctx, assetRef, digestFunction) require.NoError(t, err)