From 8b56203400d79336cd56ce55192705e86c5cfd56 Mon Sep 17 00:00:00 2001 From: Benjamin Ingberg Date: Sat, 22 Aug 2026 11:57:34 +0200 Subject: [PATCH] Consume the storage.MessageReader interface When fetching messages from the CAS, do it via the storage.MessageReader interface. This completely decouples the scheduler.InMemoryBuildQueue and builder.noopBuildExecutor from the blobstore.BlobAccess interface. --- cmd/bb_noop_worker/BUILD.bazel | 2 + cmd/bb_noop_worker/main.go | 8 +- cmd/bb_scheduler/BUILD.bazel | 1 + cmd/bb_scheduler/main.go | 7 +- cmd/bb_worker/BUILD.bazel | 1 + cmd/bb_worker/main.go | 3 +- internal/mock/BUILD.bazel | 13 ++ pkg/builder/BUILD.bazel | 1 + pkg/builder/local_build_executor.go | 10 +- pkg/builder/local_build_executor_test.go | 71 ++++---- pkg/builder/noop_build_executor.go | 17 +- pkg/builder/noop_build_executor_test.go | 22 ++- pkg/scheduler/BUILD.bazel | 3 +- pkg/scheduler/in_memory_build_queue.go | 13 +- pkg/scheduler/in_memory_build_queue_test.go | 170 +++++++++++--------- 15 files changed, 197 insertions(+), 145 deletions(-) diff --git a/cmd/bb_noop_worker/BUILD.bazel b/cmd/bb_noop_worker/BUILD.bazel index e672978b..22906554 100644 --- a/cmd/bb_noop_worker/BUILD.bazel +++ b/cmd/bb_noop_worker/BUILD.bazel @@ -12,6 +12,8 @@ go_library( "//pkg/filesystem/pool", "//pkg/proto/configuration/bb_noop_worker", "//pkg/proto/remoteworker", + "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", + "@com_github_buildbarn_bb_storage//pkg/blobstore", "@com_github_buildbarn_bb_storage//pkg/blobstore/configuration", "@com_github_buildbarn_bb_storage//pkg/clock", "@com_github_buildbarn_bb_storage//pkg/digest", diff --git a/cmd/bb_noop_worker/main.go b/cmd/bb_noop_worker/main.go index 8ecaba90..d389a208 100644 --- a/cmd/bb_noop_worker/main.go +++ b/cmd/bb_noop_worker/main.go @@ -5,11 +5,13 @@ import ( "net/url" "os" + remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2" re_blobstore "github.com/buildbarn/bb-remote-execution/pkg/blobstore" "github.com/buildbarn/bb-remote-execution/pkg/builder" "github.com/buildbarn/bb-remote-execution/pkg/filesystem/pool" "github.com/buildbarn/bb-remote-execution/pkg/proto/configuration/bb_noop_worker" "github.com/buildbarn/bb-remote-execution/pkg/proto/remoteworker" + "github.com/buildbarn/bb-storage/pkg/blobstore" blobstore_configuration "github.com/buildbarn/bb-storage/pkg/blobstore/configuration" "github.com/buildbarn/bb-storage/pkg/clock" "github.com/buildbarn/bb-storage/pkg/digest" @@ -78,8 +80,10 @@ func main() { buildClient := builder.NewBuildClient( schedulerClient, builder.NewNoopBuildExecutor( - contentAddressableStorage, - int(configuration.MaximumMessageSizeBytes), + blobstore.NewBlobAccessMessageReader[*remoteexecution.Command]( + contentAddressableStorage, + int(configuration.MaximumMessageSizeBytes), + ), browserURL, ), pool.EmptyFilePool, diff --git a/cmd/bb_scheduler/BUILD.bazel b/cmd/bb_scheduler/BUILD.bazel index c312f827..30be9a50 100644 --- a/cmd/bb_scheduler/BUILD.bazel +++ b/cmd/bb_scheduler/BUILD.bazel @@ -39,6 +39,7 @@ go_library( "//pkg/util", "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", "@com_github_buildbarn_bb_storage//pkg/auth/configuration", + "@com_github_buildbarn_bb_storage//pkg/blobstore", "@com_github_buildbarn_bb_storage//pkg/blobstore/configuration", "@com_github_buildbarn_bb_storage//pkg/capabilities", "@com_github_buildbarn_bb_storage//pkg/clock", diff --git a/cmd/bb_scheduler/main.go b/cmd/bb_scheduler/main.go index 6dcac26f..028fabba 100644 --- a/cmd/bb_scheduler/main.go +++ b/cmd/bb_scheduler/main.go @@ -17,6 +17,7 @@ import ( "github.com/buildbarn/bb-remote-execution/pkg/scheduler/initialsizeclass" "github.com/buildbarn/bb-remote-execution/pkg/scheduler/routing" auth_configuration "github.com/buildbarn/bb-storage/pkg/auth/configuration" + "github.com/buildbarn/bb-storage/pkg/blobstore" blobstore_configuration "github.com/buildbarn/bb-storage/pkg/blobstore/configuration" "github.com/buildbarn/bb-storage/pkg/capabilities" "github.com/buildbarn/bb-storage/pkg/clock" @@ -132,7 +133,10 @@ func main() { // TODO: Make timeouts configurable. generator := random.NewFastSingleThreadedGenerator() buildQueue := scheduler.NewInMemoryBuildQueue( - contentAddressableStorage, + blobstore.NewBlobAccessMessageReader[*remoteexecution.Action]( + contentAddressableStorage, + int(configuration.MaximumMessageSizeBytes), + ), clock.SystemClock, uuid.NewRandom, &scheduler.InMemoryBuildQueueConfiguration{ @@ -149,7 +153,6 @@ func main() { WorkerTaskRetryCount: 9, WorkerWithNoSynchronizationsTimeout: time.Minute, }, - int(configuration.MaximumMessageSizeBytes), actionRouter, executeAuthorizer, modifyDrainsAuthorizer, diff --git a/cmd/bb_worker/BUILD.bazel b/cmd/bb_worker/BUILD.bazel index c243ead6..3dddb1e1 100644 --- a/cmd/bb_worker/BUILD.bazel +++ b/cmd/bb_worker/BUILD.bazel @@ -23,6 +23,7 @@ go_library( "//pkg/proto/configuration/bb_worker", "//pkg/proto/remoteworker", "//pkg/proto/runner", + "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", "@com_github_buildbarn_bb_storage//pkg/blobstore", "@com_github_buildbarn_bb_storage//pkg/blobstore/configuration", "@com_github_buildbarn_bb_storage//pkg/clock", diff --git a/cmd/bb_worker/main.go b/cmd/bb_worker/main.go index f9666890..89b8a630 100644 --- a/cmd/bb_worker/main.go +++ b/cmd/bb_worker/main.go @@ -14,6 +14,7 @@ import ( "sync/atomic" "time" + remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2" re_blobstore "github.com/buildbarn/bb-remote-execution/pkg/blobstore" "github.com/buildbarn/bb-remote-execution/pkg/builder" "github.com/buildbarn/bb-remote-execution/pkg/cas" @@ -462,12 +463,12 @@ func main() { buildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorageWriter, + blobstore.NewBlobAccessMessageReader[*remoteexecution.Command](contentAddressableStorageWriter, int(configuration.MaximumMessageSizeBytes)), buildDirectoryCreator, runnerClient, executionTimeoutClock, maximumWritableFileUploadDelay, inputRootCharacterDevices, - int(configuration.MaximumMessageSizeBytes), runnerConfiguration.EnvironmentVariables, configuration.ForceUploadTreesAndDirectories, ) diff --git a/internal/mock/BUILD.bazel b/internal/mock/BUILD.bazel index 3d045bbe..9252a1c9 100644 --- a/internal/mock/BUILD.bazel +++ b/internal/mock/BUILD.bazel @@ -376,6 +376,16 @@ gomock( package = "mock", ) +gomock( + name = "storage", + out = "storage.go", + interfaces = ["MessageReader"], + library = "@com_github_buildbarn_bb_storage//pkg/storage", + mockgen_tool = "@org_uber_go_mock//mockgen", + package = "mock", + source = "@com_github_buildbarn_bb_storage//pkg/storage:message_reader.go", +) + gomock( name = "storage_builder", out = "storage_builder.go", @@ -457,6 +467,7 @@ go_library( ":routing.go", ":runner.go", ":runner_pb.go", + ":storage.go", ":storage_builder.go", ":storage_util.go", ":sync.go", @@ -506,6 +517,7 @@ go_library( "@com_github_buildbarn_bb_storage//pkg/filesystem", "@com_github_buildbarn_bb_storage//pkg/filesystem/path", "@com_github_buildbarn_bb_storage//pkg/proto/iscc", + "@com_github_buildbarn_bb_storage//pkg/storage", "@com_github_buildbarn_bb_storage//pkg/util", "@com_github_google_uuid//:uuid", "@com_google_cloud_go_longrunning//autogen/longrunningpb", @@ -515,6 +527,7 @@ go_library( "@io_opentelemetry_go_otel_trace//embedded", "@org_golang_google_grpc//:grpc", "@org_golang_google_grpc//metadata", + "@org_golang_google_protobuf//proto", "@org_golang_google_protobuf//types/known/anypb:go_default_library", "@org_golang_google_protobuf//types/known/emptypb:go_default_library", "@org_uber_go_mock//gomock", diff --git a/pkg/builder/BUILD.bazel b/pkg/builder/BUILD.bazel index a7de4111..81fce1a8 100644 --- a/pkg/builder/BUILD.bazel +++ b/pkg/builder/BUILD.bazel @@ -56,6 +56,7 @@ go_library( "@com_github_buildbarn_bb_storage//pkg/program", "@com_github_buildbarn_bb_storage//pkg/proto/fsac", "@com_github_buildbarn_bb_storage//pkg/random", + "@com_github_buildbarn_bb_storage//pkg/storage", "@com_github_buildbarn_bb_storage//pkg/util", "@com_github_google_uuid//:uuid", "@com_github_kballard_go_shellquote//:go-shellquote", diff --git a/pkg/builder/local_build_executor.go b/pkg/builder/local_build_executor.go index a1b8ee07..f8ffe03c 100644 --- a/pkg/builder/local_build_executor.go +++ b/pkg/builder/local_build_executor.go @@ -17,6 +17,7 @@ import ( "github.com/buildbarn/bb-storage/pkg/digest" "github.com/buildbarn/bb-storage/pkg/filesystem" "github.com/buildbarn/bb-storage/pkg/filesystem/path" + "github.com/buildbarn/bb-storage/pkg/storage" "github.com/buildbarn/bb-storage/pkg/util" "google.golang.org/grpc/codes" @@ -66,27 +67,27 @@ func (el *capturingErrorLogger) GetError() error { type localBuildExecutor struct { contentAddressableStorage blobstore.BlobAccess + commandReader storage.MessageReader[*remoteexecution.Command] buildDirectoryCreator BuildDirectoryCreator runner runner_pb.RunnerClient clock clock.Clock maximumWritableFileUploadDelay time.Duration inputRootCharacterDevices map[path.Component]filesystem.DeviceNumber - maximumMessageSizeBytes int environmentVariables map[string]string forceUploadTreesAndDirectories bool } // NewLocalBuildExecutor returns a BuildExecutor that executes build // steps on the local system. -func NewLocalBuildExecutor(contentAddressableStorage blobstore.BlobAccess, buildDirectoryCreator BuildDirectoryCreator, runner runner_pb.RunnerClient, clock clock.Clock, maximumWritableFileUploadDelay time.Duration, inputRootCharacterDevices map[path.Component]filesystem.DeviceNumber, maximumMessageSizeBytes int, environmentVariables map[string]string, forceUploadTreesAndDirectories bool) BuildExecutor { +func NewLocalBuildExecutor(contentAddressableStorage blobstore.BlobAccess, commandReader storage.MessageReader[*remoteexecution.Command], buildDirectoryCreator BuildDirectoryCreator, runner runner_pb.RunnerClient, clock clock.Clock, maximumWritableFileUploadDelay time.Duration, inputRootCharacterDevices map[path.Component]filesystem.DeviceNumber, environmentVariables map[string]string, forceUploadTreesAndDirectories bool) BuildExecutor { return &localBuildExecutor{ contentAddressableStorage: contentAddressableStorage, + commandReader: commandReader, buildDirectoryCreator: buildDirectoryCreator, runner: runner, clock: clock, maximumWritableFileUploadDelay: maximumWritableFileUploadDelay, inputRootCharacterDevices: inputRootCharacterDevices, - maximumMessageSizeBytes: maximumMessageSizeBytes, environmentVariables: environmentVariables, forceUploadTreesAndDirectories: forceUploadTreesAndDirectories, } @@ -231,12 +232,11 @@ func (be *localBuildExecutor) Execute(ctx context.Context, filePool pool.FilePoo attachErrorToExecuteResponse(response, util.StatusWrap(err, "Failed to extract digest for command")) return response } - commandMessage, err := be.contentAddressableStorage.Get(ctx, commandDigest).ToProto(&remoteexecution.Command{}, be.maximumMessageSizeBytes) + command, err := be.commandReader.ReadMessage(ctx, commandDigest, &remoteexecution.Command{}) if err != nil { attachErrorToExecuteResponse(response, util.StatusWrap(err, "Failed to obtain command")) return response } - command := commandMessage.(*remoteexecution.Command) outputHierarchy, err := NewOutputHierarchy(command) if err != nil { attachErrorToExecuteResponse(response, err) diff --git a/pkg/builder/local_build_executor_test.go b/pkg/builder/local_build_executor_test.go index 29255ea1..e14f2910 100644 --- a/pkg/builder/local_build_executor_test.go +++ b/pkg/builder/local_build_executor_test.go @@ -35,17 +35,18 @@ func TestLocalBuildExecutorInvalidActionDigest(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) runner := mock.NewMockRunnerClient(ctrl) clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -85,17 +86,18 @@ func TestLocalBuildExecutorMissingAction(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) runner := mock.NewMockRunnerClient(ctrl) clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -128,6 +130,7 @@ func TestLocalBuildExecutorBuildDirectoryCreatorFailedFailed(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) actionDigest := digest.MustNewDigest("netbsd", remoteexecution.DigestFunction_SHA256, "5555555555555555555555555555555555555555555555555555555555555555", 7) buildDirectoryCreator.EXPECT().GetBuildDirectory(ctx, &actionDigest). @@ -136,12 +139,12 @@ func TestLocalBuildExecutorBuildDirectoryCreatorFailedFailed(t *testing.T) { clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -181,6 +184,7 @@ func TestLocalBuildExecutorInputRootPopulationFailed(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) buildDirectory := mock.NewMockBuildDirectory(ctrl) actionDigest := digest.MustNewDigest("netbsd", remoteexecution.DigestFunction_SHA256, "5555555555555555555555555555555555555555555555555555555555555555", 7) @@ -204,12 +208,12 @@ func TestLocalBuildExecutorInputRootPopulationFailed(t *testing.T) { clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -247,16 +251,18 @@ func TestLocalBuildExecutorOutputDirectoryCreationFailure(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) + commandReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("fedora", remoteexecution.DigestFunction_SHA256, "6666666666666666666666666666666666666666666666666666666666666666", 234), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + gomock.Any(), + ).Return(&remoteexecution.Command{ Arguments: []string{"touch", "foo"}, EnvironmentVariables: []*remoteexecution.Command_EnvironmentVariable{ {Name: "PATH", Value: "/bin:/usr/bin"}, }, OutputPaths: []string{"foo/bar/baz"}, - }, buffer.UserProvided)) + }, nil) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) buildDirectory := mock.NewMockBuildDirectory(ctrl) actionDigest := digest.MustNewDigest("fedora", remoteexecution.DigestFunction_SHA256, "5555555555555555555555555555555555555555555555555555555555555555", 7) @@ -281,12 +287,12 @@ func TestLocalBuildExecutorOutputDirectoryCreationFailure(t *testing.T) { clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -328,6 +334,7 @@ func TestLocalBuildExecutorMissingCommand(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) buildDirectory := mock.NewMockBuildDirectory(ctrl) actionDigest := digest.MustNewDigest("netbsd", remoteexecution.DigestFunction_SHA256, "5555555555555555555555555555555555555555555555555555555555555555", 7) @@ -351,12 +358,12 @@ func TestLocalBuildExecutorMissingCommand(t *testing.T) { clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -394,16 +401,18 @@ func TestLocalBuildExecutorOutputSymlinkReadingFailure(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) + commandReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("nintendo64", remoteexecution.DigestFunction_SHA256, "6666666666666666666666666666666666666666666666666666666666666666", 234), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + gomock.Any(), + ).Return(&remoteexecution.Command{ Arguments: []string{"touch", "foo"}, EnvironmentVariables: []*remoteexecution.Command_EnvironmentVariable{ {Name: "PATH", Value: "/bin:/usr/bin"}, }, OutputPaths: []string{"foo"}, - }, buffer.UserProvided)) + }, nil) buildDirectory := mock.NewMockBuildDirectory(ctrl) buildDirectory.EXPECT().UploadFile(ctx, path.MustNewComponent("stdout"), gomock.Any(), gomock.Any()).Return( digest.MustNewDigest("nintendo64", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000005", 567), @@ -481,12 +490,12 @@ func TestLocalBuildExecutorOutputSymlinkReadingFailure(t *testing.T) { }) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -590,10 +599,12 @@ func TestLocalBuildExecutorSuccess(t *testing.T) { // Read operations against the Content Addressable Storage. contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) + commandReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("ubuntu1804", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000002", 234), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + gomock.Any(), + ).Return(&remoteexecution.Command{ Arguments: []string{ "/usr/local/bin/clang", "-MD", @@ -621,7 +632,7 @@ func TestLocalBuildExecutorSuccess(t *testing.T) { }, }, }, - }, buffer.UserProvided)) + }, nil) // Write operations against the Content Addressable Storage. buildDirectory := mock.NewMockBuildDirectory(ctrl) @@ -714,6 +725,7 @@ func TestLocalBuildExecutorSuccess(t *testing.T) { }) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, @@ -721,7 +733,6 @@ func TestLocalBuildExecutorSuccess(t *testing.T) { /* inputRootCharacterDevices = */ map[path.Component]filesystem.DeviceNumber{ path.MustNewComponent("null"): filesystem.NewDeviceNumberFromMajorMinor(1, 3), }, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{ "TEST_VAR": "123", "PWD": "dont-overwrite", @@ -798,17 +809,18 @@ func TestLocalBuildExecutorCachingInvalidTimeout(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) runner := mock.NewMockRunnerClient(ctrl) clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -849,12 +861,14 @@ func TestLocalBuildExecutorInputRootIOFailureDuringExecution(t *testing.T) { // Build directory. buildDirectory := mock.NewMockBuildDirectory(ctrl) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) + commandReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("ubuntu1804", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000002", 234), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + gomock.Any(), + ).Return(&remoteexecution.Command{ Arguments: []string{"clang"}, - }, buffer.UserProvided)) + }, nil) buildDirectory.EXPECT().UploadFile(ctx, path.MustNewComponent("stdout"), gomock.Any(), gomock.Any()).Return( digest.MustNewDigest("ubuntu1804", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000005", 567), nil, @@ -925,12 +939,12 @@ func TestLocalBuildExecutorInputRootIOFailureDuringExecution(t *testing.T) { }) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -984,12 +998,14 @@ func TestLocalBuildExecutorTimeoutDuringExecution(t *testing.T) { // Build directory. buildDirectory := mock.NewMockBuildDirectory(ctrl) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) + commandReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("ubuntu1804", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000002", 234), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + gomock.Any(), + ).Return(&remoteexecution.Command{ Arguments: []string{"clang"}, - }, buffer.UserProvided)) + }, nil) buildDirectory.EXPECT().UploadFile(ctx, path.MustNewComponent("stdout"), gomock.Any(), gomock.Any()).Return( digest.MustNewDigest("ubuntu1804", remoteexecution.DigestFunction_SHA256, "0000000000000000000000000000000000000000000000000000000000000005", 567), nil, @@ -1061,12 +1077,12 @@ func TestLocalBuildExecutorTimeoutDuringExecution(t *testing.T) { }) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, /* maximumWritableFileUploadDelay = */ 10*time.Second, /* inputRootCharacterDevices = */ nil, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) @@ -1126,6 +1142,7 @@ func TestLocalBuildExecutorCharacterDeviceNodeCreationFailed(t *testing.T) { // Build directory. buildDirectory := mock.NewMockBuildDirectory(ctrl) contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) // Build environment. buildDirectoryCreator := mock.NewMockBuildDirectoryCreator(ctrl) @@ -1161,6 +1178,7 @@ func TestLocalBuildExecutorCharacterDeviceNodeCreationFailed(t *testing.T) { clock := mock.NewMockClock(ctrl) localBuildExecutor := builder.NewLocalBuildExecutor( contentAddressableStorage, + commandReader, buildDirectoryCreator, runner, clock, @@ -1168,7 +1186,6 @@ func TestLocalBuildExecutorCharacterDeviceNodeCreationFailed(t *testing.T) { /* inputRootCharacterDevices = */ map[path.Component]filesystem.DeviceNumber{ path.MustNewComponent("null"): filesystem.NewDeviceNumberFromMajorMinor(1, 3), }, - /* maximumMessageSizeBytes = */ 10000, /* environmentVariables = */ map[string]string{}, /* forceUploadTreesAndDirectories = */ false, ) diff --git a/pkg/builder/noop_build_executor.go b/pkg/builder/noop_build_executor.go index 70dd1240..d9d2a9bb 100644 --- a/pkg/builder/noop_build_executor.go +++ b/pkg/builder/noop_build_executor.go @@ -11,8 +11,8 @@ import ( "github.com/buildbarn/bb-remote-execution/pkg/filesystem/pool" "github.com/buildbarn/bb-remote-execution/pkg/proto/remoteworker" re_util "github.com/buildbarn/bb-remote-execution/pkg/util" - "github.com/buildbarn/bb-storage/pkg/blobstore" "github.com/buildbarn/bb-storage/pkg/digest" + "github.com/buildbarn/bb-storage/pkg/storage" "github.com/buildbarn/bb-storage/pkg/util" "google.golang.org/grpc/codes" @@ -20,9 +20,8 @@ import ( ) type noopBuildExecutor struct { - contentAddressableStorage blobstore.BlobAccess - maximumMessageSizeBytes int - browserURL *url.URL + commandReader storage.MessageReader[*remoteexecution.Command] + browserURL *url.URL } // NewNoopBuildExecutor creates a BuildExecutor that always returns an @@ -32,11 +31,10 @@ type noopBuildExecutor struct { // to upload the input root of an action into the Content Addressable // Storage (CAS) without causing it to be executed afterwards. This may // be useful when attempting to debug actions. -func NewNoopBuildExecutor(contentAddressableStorage blobstore.BlobAccess, maximumMessageSizeBytes int, browserURL *url.URL) BuildExecutor { +func NewNoopBuildExecutor(commandReader storage.MessageReader[*remoteexecution.Command], browserURL *url.URL) BuildExecutor { return &noopBuildExecutor{ - contentAddressableStorage: contentAddressableStorage, - maximumMessageSizeBytes: maximumMessageSizeBytes, - browserURL: browserURL, + commandReader: commandReader, + browserURL: browserURL, } } @@ -69,12 +67,11 @@ func (be *noopBuildExecutor) Execute(ctx context.Context, filePool pool.FilePool attachErrorToExecuteResponse(response, util.StatusWrap(err, "Failed to extract digest for command")) return response } - commandMessage, err := be.contentAddressableStorage.Get(ctx, commandDigest).ToProto(&remoteexecution.Command{}, be.maximumMessageSizeBytes) + command, err := be.commandReader.ReadMessage(ctx, commandDigest, &remoteexecution.Command{}) if err != nil { attachErrorToExecuteResponse(response, util.StatusWrap(err, "Failed to obtain command")) return response } - command := commandMessage.(*remoteexecution.Command) errorMessageTemplate := defaultNoopErrorMessageTemplate for _, environmentVariable := range command.EnvironmentVariables { diff --git a/pkg/builder/noop_build_executor_test.go b/pkg/builder/noop_build_executor_test.go index 06500dea..09fa0133 100644 --- a/pkg/builder/noop_build_executor_test.go +++ b/pkg/builder/noop_build_executor_test.go @@ -9,7 +9,6 @@ import ( "github.com/buildbarn/bb-remote-execution/internal/mock" "github.com/buildbarn/bb-remote-execution/pkg/builder" "github.com/buildbarn/bb-remote-execution/pkg/proto/remoteworker" - "github.com/buildbarn/bb-storage/pkg/blobstore/buffer" "github.com/buildbarn/bb-storage/pkg/digest" "github.com/buildbarn/bb-storage/pkg/testutil" @@ -22,10 +21,9 @@ import ( func TestNoopBuildExecutor(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + commandReader := mock.NewMockMessageReader[*remoteexecution.Command](ctrl) buildExecutor := builder.NewNoopBuildExecutor( - contentAddressableStorage, - /* maximumMessageSizeBytes = */ 10000, + commandReader, &url.URL{ Scheme: "http", Host: "example.com", @@ -93,15 +91,15 @@ func TestNoopBuildExecutor(t *testing.T) { t.Run("InvalidTemplate", func(t *testing.T) { // If an invalid template is provided in the // environment, parsing it should fail. - contentAddressableStorage.EXPECT().Get(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "7f53aed4b5489c487be514dd88d3314d966e19b84bc766a972d82246ee6f494f", 150)). - Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + commandReader.EXPECT().ReadMessage(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "7f53aed4b5489c487be514dd88d3314d966e19b84bc766a972d82246ee6f494f", 150), gomock.Any()). + Return(&remoteexecution.Command{ EnvironmentVariables: []*remoteexecution.Command_EnvironmentVariable{ { Name: "NOOP_WORKER_ERROR_MESSAGE_TEMPLATE", Value: "{{ foobarbaz }}", }, }, - }, buffer.UserProvided)) + }, nil) filePool := mock.NewMockFilePool(ctrl) monitor := mock.NewMockUnreadDirectoryMonitor(ctrl) testutil.RequireEqualProto( @@ -137,8 +135,8 @@ func TestNoopBuildExecutor(t *testing.T) { t.Run("SuccessDefaultTemplate", func(t *testing.T) { // If no template is provided in the environment // variables, then a default template should be used. - contentAddressableStorage.EXPECT().Get(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "d134371fd7573f7ef77c90e907c8bfaf95f34b82ac8503dbed5e062fb6fe4702", 200)). - Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{}, buffer.UserProvided)) + commandReader.EXPECT().ReadMessage(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "d134371fd7573f7ef77c90e907c8bfaf95f34b82ac8503dbed5e062fb6fe4702", 200), gomock.Any()). + Return(&remoteexecution.Command{}, nil) filePool := mock.NewMockFilePool(ctrl) monitor := mock.NewMockUnreadDirectoryMonitor(ctrl) testutil.RequireEqualProto( @@ -174,8 +172,8 @@ func TestNoopBuildExecutor(t *testing.T) { t.Run("SuccessCustomTemplate", func(t *testing.T) { // If a custom template is provided in the environment, // it should be preferred over the default template. - contentAddressableStorage.EXPECT().Get(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "9da17cb226048f5bb3e6a20311b551e73ce8ac0d408e69e737d28a8f3179d0ce", 300)). - Return(buffer.NewProtoBufferFromProto(&remoteexecution.Command{ + commandReader.EXPECT().ReadMessage(ctx, digest.MustNewDigest("build", remoteexecution.DigestFunction_SHA256, "9da17cb226048f5bb3e6a20311b551e73ce8ac0d408e69e737d28a8f3179d0ce", 300), gomock.Any()). + Return(&remoteexecution.Command{ EnvironmentVariables: []*remoteexecution.Command_EnvironmentVariable{ { Name: "PATH", @@ -186,7 +184,7 @@ func TestNoopBuildExecutor(t *testing.T) { Value: "Please visit {{ .ActionURL }} to inspect the action", }, }, - }, buffer.UserProvided)) + }, nil) filePool := mock.NewMockFilePool(ctrl) monitor := mock.NewMockUnreadDirectoryMonitor(ctrl) testutil.RequireEqualProto( diff --git a/pkg/scheduler/BUILD.bazel b/pkg/scheduler/BUILD.bazel index 2c1592b0..8c21c21b 100644 --- a/pkg/scheduler/BUILD.bazel +++ b/pkg/scheduler/BUILD.bazel @@ -16,12 +16,12 @@ go_library( "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", "@bazel_remote_apis//build/bazel/semver:semver_go_proto", "@com_github_buildbarn_bb_storage//pkg/auth", - "@com_github_buildbarn_bb_storage//pkg/blobstore", "@com_github_buildbarn_bb_storage//pkg/builder", "@com_github_buildbarn_bb_storage//pkg/capabilities", "@com_github_buildbarn_bb_storage//pkg/clock", "@com_github_buildbarn_bb_storage//pkg/digest", "@com_github_buildbarn_bb_storage//pkg/otel", + "@com_github_buildbarn_bb_storage//pkg/storage", "@com_github_buildbarn_bb_storage//pkg/util", "@com_github_google_uuid//:uuid", "@com_github_prometheus_client_golang//prometheus", @@ -56,7 +56,6 @@ go_test( "//pkg/scheduler/platform", "@bazel_remote_apis//build/bazel/remote/execution/v2:remote_execution_go_proto", "@com_github_buildbarn_bb_storage//pkg/auth", - "@com_github_buildbarn_bb_storage//pkg/blobstore/buffer", "@com_github_buildbarn_bb_storage//pkg/builder", "@com_github_buildbarn_bb_storage//pkg/clock", "@com_github_buildbarn_bb_storage//pkg/digest", diff --git a/pkg/scheduler/in_memory_build_queue.go b/pkg/scheduler/in_memory_build_queue.go index 7240add6..1aee52f1 100644 --- a/pkg/scheduler/in_memory_build_queue.go +++ b/pkg/scheduler/in_memory_build_queue.go @@ -21,12 +21,12 @@ import ( "github.com/buildbarn/bb-remote-execution/pkg/scheduler/platform" "github.com/buildbarn/bb-remote-execution/pkg/scheduler/routing" "github.com/buildbarn/bb-storage/pkg/auth" - "github.com/buildbarn/bb-storage/pkg/blobstore" "github.com/buildbarn/bb-storage/pkg/builder" "github.com/buildbarn/bb-storage/pkg/capabilities" "github.com/buildbarn/bb-storage/pkg/clock" "github.com/buildbarn/bb-storage/pkg/digest" "github.com/buildbarn/bb-storage/pkg/otel" + "github.com/buildbarn/bb-storage/pkg/storage" "github.com/buildbarn/bb-storage/pkg/util" "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" @@ -236,12 +236,11 @@ type InMemoryBuildQueueConfiguration struct { type InMemoryBuildQueue struct { capabilities.Provider - contentAddressableStorage blobstore.BlobAccess + actionReader storage.MessageReader[*remoteexecution.Action] clock clock.Clock uuidGenerator util.UUIDGenerator configuration *InMemoryBuildQueueConfiguration platformQueueAbsenceHardFailureTime time.Time - maximumMessageSizeBytes int actionRouter routing.ActionRouter lock sync.Mutex @@ -316,7 +315,7 @@ var inMemoryBuildQueueCapabilitiesProvider = capabilities.NewStaticProvider(&rem // NewInMemoryBuildQueue creates a new InMemoryBuildQueue that is in the // initial state. It does not have any queues, workers or queued // execution requests. All of these are created by sending it RPCs. -func NewInMemoryBuildQueue(contentAddressableStorage blobstore.BlobAccess, clock clock.Clock, uuidGenerator util.UUIDGenerator, configuration *InMemoryBuildQueueConfiguration, maximumMessageSizeBytes int, actionRouter routing.ActionRouter, executeAuthorizer, modifyDrainsAuthorizer, killOperationsAuthorizer, synchronizeAuthorizer auth.Authorizer) *InMemoryBuildQueue { +func NewInMemoryBuildQueue(actionReader storage.MessageReader[*remoteexecution.Action], clock clock.Clock, uuidGenerator util.UUIDGenerator, configuration *InMemoryBuildQueueConfiguration, actionRouter routing.ActionRouter, executeAuthorizer, modifyDrainsAuthorizer, killOperationsAuthorizer, synchronizeAuthorizer auth.Authorizer) *InMemoryBuildQueue { inMemoryBuildQueuePrometheusMetrics.Do(func() { prometheus.MustRegister(inMemoryBuildQueueInFlightDeduplicationsTotal) @@ -341,12 +340,11 @@ func NewInMemoryBuildQueue(contentAddressableStorage blobstore.BlobAccess, clock return &InMemoryBuildQueue{ Provider: capabilities.NewAuthorizingProvider(inMemoryBuildQueueCapabilitiesProvider, executeAuthorizer), - contentAddressableStorage: contentAddressableStorage, + actionReader: actionReader, clock: clock, uuidGenerator: uuidGenerator, configuration: configuration, platformQueueAbsenceHardFailureTime: clock.Now().Add(configuration.PlatformQueueWithNoWorkersTimeout), - maximumMessageSizeBytes: maximumMessageSizeBytes, actionRouter: actionRouter, platformQueuesTrie: platform.NewTrie(), sizeClassQueues: map[sizeClassKey]*sizeClassQueue{}, @@ -450,11 +448,10 @@ func (bq *InMemoryBuildQueue) Execute(in *remoteexecution.ExecuteRequest, out re if err != nil { return util.StatusWrap(err, "Failed to extract digest for action") } - actionMessage, err := bq.contentAddressableStorage.Get(ctx, actionDigest).ToProto(&remoteexecution.Action{}, bq.maximumMessageSizeBytes) + action, err := bq.actionReader.ReadMessage(ctx, actionDigest, &remoteexecution.Action{}) if err != nil { return util.StatusWrap(err, "Failed to obtain action") } - action := actionMessage.(*remoteexecution.Action) platformKey, err := platform.NewKey(instanceName, action.Platform) if err != nil { return err diff --git a/pkg/scheduler/in_memory_build_queue_test.go b/pkg/scheduler/in_memory_build_queue_test.go index 455c40f5..23754d45 100644 --- a/pkg/scheduler/in_memory_build_queue_test.go +++ b/pkg/scheduler/in_memory_build_queue_test.go @@ -16,7 +16,6 @@ import ( "github.com/buildbarn/bb-remote-execution/pkg/scheduler/invocation" "github.com/buildbarn/bb-remote-execution/pkg/scheduler/platform" "github.com/buildbarn/bb-storage/pkg/auth" - "github.com/buildbarn/bb-storage/pkg/blobstore/buffer" "github.com/buildbarn/bb-storage/pkg/builder" "github.com/buildbarn/bb-storage/pkg/clock" "github.com/buildbarn/bb-storage/pkg/digest" @@ -88,12 +87,12 @@ func getExecutionClient(t *testing.T, buildQueue builder.BuildQueue) remoteexecu func TestInMemoryBuildQueueExecuteBadRequest(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // ExecuteRequest contains an invalid action digest. @@ -111,10 +110,11 @@ func TestInMemoryBuildQueueExecuteBadRequest(t *testing.T) { // Action cannot be found in the Content Addressable Storage (CAS). t.Run("MissingAction", func(t *testing.T) { - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewBufferFromError(status.Error(codes.FailedPrecondition, "Blob not found"))) + gomock.Any(), + ).Return(nil, status.Error(codes.FailedPrecondition, "Blob not found")) stream, err := executionClient.Execute(ctx, &remoteexecution.ExecuteRequest{ InstanceName: "main", @@ -139,10 +139,11 @@ func TestInMemoryBuildQueueExecuteBadRequest(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil).Return(action, platform.MustNewKey("main", platformForTesting), nil, initialSizeClassSelector, nil) initialSizeClassSelector.EXPECT().Abandoned() @@ -170,10 +171,11 @@ func TestInMemoryBuildQueueExecuteBadRequest(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil).Return(action, platform.MustNewKey("main", platformForTesting), nil, initialSizeClassSelector, nil) initialSizeClassSelector.EXPECT().Abandoned() @@ -195,24 +197,25 @@ func TestInMemoryBuildQueueExecuteBadRequest(t *testing.T) { func TestInMemoryBuildQueuePurgeStaleWorkersAndQueues(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) actionRouter := mock.NewMockActionRouter(ctrl) for i := 0; i < 10; i++ { - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(&remoteexecution.Action{ + gomock.Any(), + ).Return(&remoteexecution.Action{ CommandDigest: &remoteexecution.Digest{ Hash: "61c585c297d00409bd477b6b80759c94ec545ab4", SizeBytes: 456, }, DoNotCache: true, - }, buffer.UserProvided)) + }, nil) } clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -498,18 +501,19 @@ func TestInMemoryBuildQueuePurgeStaleOperations(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) for i := 0; i < 2; i++ { - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) } clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -742,22 +746,23 @@ func TestInMemoryBuildQueuePurgeStaleOperations(t *testing.T) { func TestInMemoryBuildQueueCrashLoopingWorker(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) action := &remoteexecution.Action{ CommandDigest: &remoteexecution.Digest{ Hash: "61c585c297d00409bd477b6b80759c94ec545ab4", SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main/suffix", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -968,16 +973,17 @@ func TestInMemoryBuildQueueKillOperationsOperationName(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -1186,16 +1192,17 @@ func TestInMemoryBuildQueueKillOperationsSizeClassQueueWithoutWorkers(t *testing SizeBytes: 456, }, } - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // If the scheduler is in the initial state, the size class @@ -1354,12 +1361,12 @@ func TestInMemoryBuildQueueKillOperationsSizeClassQueueWithoutWorkers(t *testing func TestInMemoryBuildQueueIdleWorkerSynchronizationTimeout(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) // When no work appears, workers should still be woken up // periodically to resynchronize. This ensures that workers that @@ -1409,16 +1416,17 @@ func TestInMemoryBuildQueueDrainedWorker(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) - contentAddressableStorage.EXPECT().Get( + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -1712,12 +1720,12 @@ func TestInMemoryBuildQueueDrainedWorker(t *testing.T) { func TestInMemoryBuildQueueInvocationFairness(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -1797,10 +1805,11 @@ func TestInMemoryBuildQueueInvocationFairness(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_MD5, p.actionHash, 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) requestMetadata := &remoteexecution.RequestMetadata{ ToolInvocationId: p.invocationID, @@ -2109,12 +2118,12 @@ func TestInMemoryBuildQueueInvocationFairness(t *testing.T) { func TestInMemoryBuildQueueInFlightDeduplicationAbandonQueued(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -2178,10 +2187,11 @@ func TestInMemoryBuildQueueInFlightDeduplicationAbandonQueued(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA256, "fc96ea0eee854b45950d3a7448332445730886691b992cb7917da0853664f7c2", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) requestMetadata := &remoteexecution.RequestMetadata{ @@ -2305,12 +2315,12 @@ func TestInMemoryBuildQueueInFlightDeduplicationAbandonQueued(t *testing.T) { func TestInMemoryBuildQueueInFlightDeduplicationAbandonExecuting(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -2375,10 +2385,11 @@ func TestInMemoryBuildQueueInFlightDeduplicationAbandonExecuting(t *testing.T) { }, Platform: platformForTesting, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA256, "fc96ea0eee854b45950d3a7448332445730886691b992cb7917da0853664f7c2", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) requestMetadata := &remoteexecution.RequestMetadata{ @@ -2545,12 +2556,12 @@ func TestInMemoryBuildQueueInFlightDeduplicationAbandonExecuting(t *testing.T) { func TestInMemoryBuildQueuePreferBeingIdle(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Announce a new worker, which creates a queue for operations. @@ -2586,10 +2597,11 @@ func TestInMemoryBuildQueuePreferBeingIdle(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil).Return(action, platform.MustNewKey("main", platformForTesting), nil, initialSizeClassSelector, nil) initialSizeClassLearner := mock.NewMockLearner(ctrl) @@ -2767,12 +2779,12 @@ func TestInMemoryBuildQueuePreferBeingIdle(t *testing.T) { func TestInMemoryBuildQueueMultipleSizeClasses(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Register a platform queue that allows workers up to size @@ -2850,10 +2862,11 @@ func TestInMemoryBuildQueueMultipleSizeClasses(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil).Return(action, platform.MustNewKey("main", platformForTesting), nil, initialSizeClassSelector, nil) initialSizeClassLearner1 := mock.NewMockLearner(ctrl) @@ -3144,12 +3157,12 @@ func TestInMemoryBuildQueueMultipleSizeClasses(t *testing.T) { func TestInMemoryBuildQueueBackgroundRun(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Register a platform queue that allows workers up to size @@ -3206,10 +3219,11 @@ func TestInMemoryBuildQueueBackgroundRun(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("main", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil).Return(action, platform.MustNewKey("main", platformForTesting), nil, initialSizeClassSelector, nil) initialSizeClassLearner1 := mock.NewMockLearner(ctrl) @@ -3468,12 +3482,12 @@ func TestInMemoryBuildQueueBackgroundRun(t *testing.T) { func TestInMemoryBuildQueueIdleSynchronizingWorkers(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) mockClock := mock.NewMockClock(ctrl) mockClock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, mockClock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, mockClock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Common values used by steps below. @@ -3531,10 +3545,11 @@ func TestInMemoryBuildQueueIdleSynchronizingWorkers(t *testing.T) { }) require.NoError(t, err) - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("", remoteexecution.DigestFunction_SHA1, "da39a3ee5e6b4b0d3255bfef95601890afd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)).AnyTimes() + gomock.Any(), + ).Return(action, nil).AnyTimes() // Create a worker that does a blocking Synchronize() call // against the scheduler. @@ -3891,12 +3906,12 @@ func TestInMemoryBuildQueueIdleSynchronizingWorkers(t *testing.T) { func TestInMemoryBuildQueueWorkerInvocationStickinessLimit(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) // Register a platform queue that has a small amount of worker @@ -3943,10 +3958,11 @@ func TestInMemoryBuildQueueWorkerInvocationStickinessLimit(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("", remoteexecution.DigestFunction_SHA1, "0474d2f48968a56da4de20718d8ac23aafd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) requestMetadata := &remoteexecution.RequestMetadata{ ToolInvocationId: p.toolInvocationID, } @@ -4121,13 +4137,13 @@ func TestInMemoryBuildQueueWorkerInvocationStickinessLimit(t *testing.T) { func TestInMemoryBuildQueueAuthorization(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) clock := mock.NewMockClock(ctrl) clock.EXPECT().Now().Return(time.Unix(0, 0)).AnyTimes() uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) authorizer := mock.NewMockAuthorizer(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, authorizer, authorizer, authorizer, authorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, clock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, authorizer, authorizer, authorizer, authorizer) beepboop := util.Must(digest.NewInstanceName("beepboop")) t.Run("GetCapabilities-NotAuthorized", func(t *testing.T) { @@ -4188,10 +4204,11 @@ func TestInMemoryBuildQueueAuthorization(t *testing.T) { Platform: &remoteexecution.Platform{}, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("beepboop", remoteexecution.DigestFunction_SHA1, "61c585c297d00409bd477b6b80759c94ec545ab4", 456), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) initialSizeClassSelector := mock.NewMockSelector(ctrl) actionRouter.EXPECT().RouteAction(gomock.Any(), gomock.Any(), testutil.EqProto(t, action), nil). @@ -4245,12 +4262,12 @@ func TestInMemoryBuildQueueAuthorization(t *testing.T) { func TestInMemoryBuildQueueNestedInvocationsSynchronization(t *testing.T) { ctrl, ctx := gomock.WithContext(context.Background(), t) - contentAddressableStorage := mock.NewMockBlobAccess(ctrl) + actionReader := mock.NewMockMessageReader[*remoteexecution.Action](ctrl) mockClock := mock.NewMockClock(ctrl) mockClock.EXPECT().Now().Return(time.Unix(0, 0)) uuidGenerator := mock.NewMockUUIDGenerator(ctrl) actionRouter := mock.NewMockActionRouter(ctrl) - buildQueue := scheduler.NewInMemoryBuildQueue(contentAddressableStorage, mockClock, uuidGenerator.Call, &buildQueueConfigurationForTesting, 10000, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) + buildQueue := scheduler.NewInMemoryBuildQueue(actionReader, mockClock, uuidGenerator.Call, &buildQueueConfigurationForTesting, actionRouter, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer, allowAllAuthorizer) executionClient := getExecutionClient(t, buildQueue) mockClock.EXPECT().Now().Return(time.Unix(1000, 0)) @@ -4293,10 +4310,11 @@ func TestInMemoryBuildQueueNestedInvocationsSynchronization(t *testing.T) { SizeBytes: 456, }, } - contentAddressableStorage.EXPECT().Get( + actionReader.EXPECT().ReadMessage( gomock.Any(), digest.MustNewDigest("", remoteexecution.DigestFunction_SHA1, "0474d2f48968a56da4de20718d8ac23aafd80709", 123), - ).Return(buffer.NewProtoBufferFromProto(action, buffer.UserProvided)) + gomock.Any(), + ).Return(action, nil) toolInvocationID := &remoteexecution.RequestMetadata{ ToolInvocationId: p.toolInvocationID, }