diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 3bd852c8..2ccb3493 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -41,6 +41,9 @@ go_library( "//submitqueue/extension/scorer/composite:go_default_library", "//submitqueue/extension/scorer/fake:go_default_library", "//submitqueue/extension/scorer/heuristic:go_default_library", + "//submitqueue/extension/speculation/prioritizationlimit/static:go_default_library", + "//submitqueue/extension/speculation/prioritizer:go_default_library", + "//submitqueue/extension/speculation/prioritizer/sticky:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mysql:go_default_library", "//submitqueue/extension/validator/fake:go_default_library", @@ -54,6 +57,7 @@ go_library( "//submitqueue/orchestrator/controller/merge:go_default_library", "//submitqueue/orchestrator/controller/mergeconflictsignal:go_default_library", "//submitqueue/orchestrator/controller/mergesignal:go_default_library", + "//submitqueue/orchestrator/controller/prioritize:go_default_library", "//submitqueue/orchestrator/controller/score:go_default_library", "//submitqueue/orchestrator/controller/speculate:go_default_library", "//submitqueue/orchestrator/controller/start:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 2038b2e7..4ddda8cf 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -61,6 +61,9 @@ import ( "github.com/uber/submitqueue/submitqueue/extension/scorer/composite" scorerfake "github.com/uber/submitqueue/submitqueue/extension/scorer/fake" "github.com/uber/submitqueue/submitqueue/extension/scorer/heuristic" + prioritizationlimitstatic "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizationlimit/static" + "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer/sticky" "github.com/uber/submitqueue/submitqueue/extension/storage" mysqlstorage "github.com/uber/submitqueue/submitqueue/extension/storage/mysql" validatorfake "github.com/uber/submitqueue/submitqueue/extension/validator/fake" @@ -74,6 +77,7 @@ import ( "github.com/uber/submitqueue/submitqueue/orchestrator/controller/merge" "github.com/uber/submitqueue/submitqueue/orchestrator/controller/mergeconflictsignal" "github.com/uber/submitqueue/submitqueue/orchestrator/controller/mergesignal" + "github.com/uber/submitqueue/submitqueue/orchestrator/controller/prioritize" "github.com/uber/submitqueue/submitqueue/orchestrator/controller/score" "github.com/uber/submitqueue/submitqueue/orchestrator/controller/speculate" "github.com/uber/submitqueue/submitqueue/orchestrator/controller/start" @@ -244,13 +248,14 @@ func run() error { brf := buildRunnerFactory{queues} scf := scorerFactory{queues} cof := analyzerFactory{queues} + prf := prioritizerFactory{queues} // Register controllers - primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, cnt, store) + primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, prf, cnt, store) if err != nil { return err } - dlqCount, err := registerDLQControllers(dlqConsumer, logger.Sugar(), scope, store) + dlqCount, err := registerDLQControllers(dlqConsumer, logger.Sugar(), scope, registry, store) if err != nil { return err } @@ -380,6 +385,7 @@ func newTopicRegistry(q extqueue.Queue, subscriberName string) (consumer.TopicRe {topickey.TopicKeyBatch, "batch", "orchestrator-batch"}, {topickey.TopicKeyScore, "score", "orchestrator-score"}, {topickey.TopicKeySpeculate, "speculate", "orchestrator-speculate"}, + {topickey.TopicKeyPrioritize, "prioritize", "orchestrator-prioritize"}, {topickey.TopicKeyBuild, "build", "orchestrator-build"}, {topickey.TopicKeyBuildSignal, "buildsignal", "orchestrator-buildsignal"}, {topickey.TopicKeyMerge, "merge", "orchestrator-merge"}, @@ -471,6 +477,13 @@ func newTopicRegistry(q extqueue.Queue, subscriberName string) (consumer.TopicRe // merge-conflict-check queue (⇢); runway performs the merge attempt and // publishes the result to merge-conflict-check-signal, which mergeconflictsignal // consumes before fanning the request out to batch. +// +// prioritize sits alongside this per-batch flow rather than in its line: it +// is queue-wide, not batch-scoped. Its message carries only a queue name; on +// each invocation it loads every Speculating batch's speculation tree for +// that queue, ranks the queue-wide candidate paths against the queue's build +// budget, applies the resulting decisions, and republishes to build for any +// path newly (or still) cleared to run. // TODO(wiring abstraction): queueExtensions + queueRegistry currently live here // as example-local wiring. Evaluate promoting them into a defined abstraction in @@ -493,6 +506,7 @@ type queueExtensions struct { buildRunner buildrunner.BuildRunner scorer scorer.Scorer analyzer conflict.Analyzer + prioritizer prioritizer.Prioritizer } // queueRegistry maps a queue name to its extensions, falling back to a default @@ -538,7 +552,13 @@ func (f analyzerFactory) For(cfg conflict.Config) (conflict.Analyzer, error) { return f.reg.get(cfg.QueueName).analyzer, nil } -func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, registry consumer.TopicRegistry, cpf changeprovider.Factory, brf buildrunner.Factory, scf scorer.Factory, cof conflict.Factory, cnt counter.Counter, store storage.Storage) (int, error) { +type prioritizerFactory struct{ reg queueRegistry } + +func (f prioritizerFactory) For(cfg prioritizer.Config) (prioritizer.Prioritizer, error) { + return f.reg.get(cfg.QueueName).prioritizer, nil +} + +func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, registry consumer.TopicRegistry, cpf changeprovider.Factory, brf buildrunner.Factory, scf scorer.Factory, cof conflict.Factory, prf prioritizer.Factory, cnt counter.Counter, store storage.Storage) (int, error) { var count int requestController := start.NewController( logger, @@ -637,6 +657,20 @@ func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, } count++ + prioritizeController := prioritize.NewController( + logger, + scope, + store, + prf, + registry, + topickey.TopicKeyPrioritize, + "orchestrator-prioritize", + ) + if err := c.Register(prioritizeController); err != nil { + return count, fmt.Errorf("failed to register prioritize controller: %w", err) + } + count++ + buildController := build.NewController( logger, scope, @@ -712,7 +746,7 @@ func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, // registers them with the DLQ consumer. Each reconciler drives the affected // request or batch into a terminal Error/Failed state so the gateway stops // reporting it as stuck-in-progress. -func registerDLQControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, store storage.Storage) (int, error) { +func registerDLQControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, registry consumer.TopicRegistry, store storage.Storage) (int, error) { dlqScope := scope.SubScope("dlq") dlqRegs := []struct { name string @@ -725,6 +759,7 @@ func registerDLQControllers(c consumer.Consumer, logger *zap.SugaredLogger, scop {"batch_dlq", dlq.NewDLQRequestController(logger, dlqScope, store, dlq.DecodeRequestID, dlq.TopicKey(topickey.TopicKeyBatch), "orchestrator-batch-dlq")}, {"score_dlq", dlq.NewDLQBatchController(logger, dlqScope, store, dlq.TopicKey(topickey.TopicKeyScore), "orchestrator-score-dlq")}, {"speculate_dlq", dlq.NewDLQBatchController(logger, dlqScope, store, dlq.TopicKey(topickey.TopicKeySpeculate), "orchestrator-speculate-dlq")}, + {"prioritize_dlq", dlq.NewDLQQueueController(logger, dlqScope, registry, dlq.TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq")}, {"build_dlq", dlq.NewDLQBatchController(logger, dlqScope, store, dlq.TopicKey(topickey.TopicKeyBuild), "orchestrator-build-dlq")}, {"buildsignal_dlq", dlq.NewDLQBuildSignalController(logger, dlqScope, store, dlq.TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq")}, {"merge_dlq", dlq.NewDLQBatchController(logger, dlqScope, store, dlq.TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq")}, @@ -859,6 +894,11 @@ func newPhabChangeProvider(logger *zap.Logger, scope tally.Scope) (changeprovide }), nil } +// defaultPrioritizationLimit is the baseline queue-wide concurrent-build +// budget handed to the sticky prioritizer. It is a parity default — +// effectively admit-all — until per-queue budgets are configured. +const defaultPrioritizationLimit = 1000 + // newQueueRegistry builds the per-queue extension profiles for the example. // Edge integrations (change provider) and the build // runner form a shared baseline; each per-queue profile starts from that @@ -880,16 +920,19 @@ func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset. // Baseline profile: shared edge integrations + a fake build runner (every // build succeeds unless a head URI carries a failure marker), plus permissive - // defaults for scorer and conflict. The build runner instance is shared by - // the build and buildsignal controllers (same profile, same instance) so a - // build's recorded outcome survives across their separate factory lookups. + // defaults for scorer, conflict, and prioritization. The build runner + // instance is shared by the build and buildsignal controllers (same + // profile, same instance) so a build's recorded outcome survives across + // their separate factory lookups. // // The scorer is wrapped by scorerfake so a change URI carrying // "sq-fake=score-error" forces a scoring error end-to-end; it is a pure // passthrough otherwise. The analyzer is wrapped by conflictfake with a nil // predicate (passthrough) — swap the predicate (e.g. conflictfake.FailAlways) // on a queue to exercise the analyzer error path, as e2e-conflict-error-queue - // below does. + // below does. The prioritizer is sticky over a static budget: it never + // preempts a running build and admits Selected candidates by score until + // defaultPrioritizationLimit concurrent builds are in flight. base := queueExtensions{ changeProvider: cp, buildRunner: buildfake.New(resolver), @@ -900,7 +943,8 @@ func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset. )), // TODO: replace the delegate with a real analyzer (e.g. Tango target // analysis). "all" serializes the queue conservatively. - analyzer: conflictfake.New(all.New(), nil), + analyzer: conflictfake.New(all.New(), nil), + prioritizer: sticky.New(prioritizationlimitstatic.New(defaultPrioritizationLimit)), } // test-queue: bucketed heuristic scorer; conservative (serialized) conflicts diff --git a/submitqueue/core/topickey/topickey.go b/submitqueue/core/topickey/topickey.go index 426de3d3..66b23bf3 100644 --- a/submitqueue/core/topickey/topickey.go +++ b/submitqueue/core/topickey/topickey.go @@ -33,6 +33,14 @@ const ( TopicKeyScore TopicKey = "score" // TopicKeySpeculate is the pipeline stage where scored batches are published for speculation. TopicKeySpeculate TopicKey = "speculate" + // TopicKeyPrioritize is the queue-wide reconcile stage that rations the + // build budget across every in-flight batch of a queue. Each message + // carries a QueueID; the consumer loads every Speculating batch's tree, + // runs the queue's Prioritizer over the candidate paths, applies the + // resulting decisions — promoting paths into the build budget, or + // cancelling in-flight paths a preemptive policy evicts — and republishes + // to TopicKeyBuild for any path cleared to run. + TopicKeyPrioritize TopicKey = "prioritize" // TopicKeyBuild is the pipeline stage where speculated batches are published for builds. TopicKeyBuild TopicKey = "build" // TopicKeyBuildSignal is the polling stage for triggered builds. Each diff --git a/submitqueue/entity/BUILD.bazel b/submitqueue/entity/BUILD.bazel index 4e67b83b..60d9fa82 100644 --- a/submitqueue/entity/BUILD.bazel +++ b/submitqueue/entity/BUILD.bazel @@ -14,6 +14,7 @@ go_library( "land_request.go", "merge_result.go", "push_result.go", + "queue.go", "queue_config.go", "request.go", "request_log.go", @@ -36,6 +37,7 @@ go_test( "build_test.go", "cancel_request_test.go", "land_request_test.go", + "queue_test.go", "request_log_test.go", "request_test.go", "speculation_tree_test.go", diff --git a/submitqueue/entity/queue.go b/submitqueue/entity/queue.go new file mode 100644 index 00000000..26409afc --- /dev/null +++ b/submitqueue/entity/queue.go @@ -0,0 +1,36 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package entity + +import "encoding/json" + +// QueueID is the queue-message payload for queue-scoped pipeline stages. It +// carries only the queue name. +type QueueID struct { + // Name is the merge-queue name the message targets. + Name string `json:"name"` +} + +// ToBytes serializes the QueueID to JSON bytes for queue message payload. +func (q QueueID) ToBytes() ([]byte, error) { + return json.Marshal(q) +} + +// QueueIDFromBytes deserializes a QueueID from JSON bytes. +func QueueIDFromBytes(data []byte) (QueueID, error) { + var qid QueueID + err := json.Unmarshal(data, &qid) + return qid, err +} diff --git a/submitqueue/entity/queue_test.go b/submitqueue/entity/queue_test.go new file mode 100644 index 00000000..543ea07b --- /dev/null +++ b/submitqueue/entity/queue_test.go @@ -0,0 +1,72 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package entity + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestQueueID_SerializationRoundTrip(t *testing.T) { + tests := []struct { + name string + queueID QueueID + }{ + { + name: "simple queue name", + queueID: QueueID{Name: "queueA"}, + }, + { + name: "another queue name", + queueID: QueueID{Name: "queueB"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := tt.queueID.ToBytes() + require.NoError(t, err) + + deserialized, err := QueueIDFromBytes(data) + require.NoError(t, err) + + assert.Equal(t, tt.queueID, deserialized) + }) + } +} + +func TestQueueIDFromBytes_InvalidJSON(t *testing.T) { + _, err := QueueIDFromBytes([]byte(`{"invalid": json"}`)) + assert.Error(t, err) +} + +func TestQueueIDFromBytes_EmptyJSON(t *testing.T) { + queueID, err := QueueIDFromBytes([]byte(`{}`)) + require.NoError(t, err) + + assert.Empty(t, queueID.Name) +} + +func TestQueueIDFromBytes_EmptyBytes(t *testing.T) { + _, err := QueueIDFromBytes([]byte{}) + assert.Error(t, err) +} + +func TestQueueIDFromBytes_NilBytes(t *testing.T) { + _, err := QueueIDFromBytes(nil) + assert.Error(t, err) +} diff --git a/submitqueue/entity/speculation_tree.go b/submitqueue/entity/speculation_tree.go index 92c30b56..85f74473 100644 --- a/submitqueue/entity/speculation_tree.go +++ b/submitqueue/entity/speculation_tree.go @@ -19,10 +19,6 @@ import "slices" // SpeculationPath is a single speculation path: an assumed-good prefix of // predecessor batches (Base) on top of which the batch under verification // (Head) is built and validated. -// -// This is the unit the build stage consumes: Base maps to the build runner's -// base changes (an assumed-good prefix to apply) and Head maps to the changes -// being validated. type SpeculationPath struct { // Base is the ordered list of predecessor batch IDs assumed to have passed. // Empty means the path builds the head batch directly on the target branch. @@ -33,18 +29,14 @@ type SpeculationPath struct { // Equal reports whether p and other are structurally the same speculation // path. It is true iff Head matches and Base has the same elements in the same -// order — Base order is the build order and is significant. The controller uses -// it where only structure can identify a path (deduplicating enumerator output, -// carrying entries over across re-enumeration); everything else references a -// persisted path by its assigned ID (SpeculationPathInfo.ID). +// order — Base order is the build order and is significant. Structural +// equality identifies a path that has no assigned ID yet; a persisted path is +// referenced by its ID (SpeculationPathInfo.ID). func (p SpeculationPath) Equal(other SpeculationPath) bool { return p.Head == other.Head && slices.Equal(p.Base, other.Base) } // SpeculationPathStatus is the observed lifecycle state of a speculation path. -// It is written only by the orchestrator's speculate controller (into the -// speculation tree store) and read by the decision seams (selector, prioritizer) -// as input; the seams never write it. type SpeculationPathStatus string const ( @@ -52,30 +44,28 @@ const ( // on init. A persisted path always carries a real status (candidate onward), // so this should never be seen in the store. SpeculationPathStatusUnknown SpeculationPathStatus = "" - // SpeculationPathStatusCandidate is a freshly enumerated path the controller - // has persisted but not yet acted on. + // SpeculationPathStatusCandidate is a freshly enumerated path, persisted but + // not yet acted on. SpeculationPathStatusCandidate SpeculationPathStatus = "candidate" - // SpeculationPathStatusSelected is a path the selector has promoted — a - // per-batch desire to build it — that the queue-wide prioritizer has not yet - // cleared. It is not sent to build while Selected: it waits on the build - // budget until it is prioritized (or dropped). + // SpeculationPathStatusSelected is a path promoted by per-batch selection — + // a desire to build it — that queue-wide prioritization has not yet cleared. + // It is not built while Selected: it waits on the build budget until it is + // prioritized (or dropped). SpeculationPathStatusSelected SpeculationPathStatus = "selected" - // SpeculationPathStatusPrioritized is a path the prioritizer has admitted - // under the queue's build budget; it is cleared to run. The build effector - // triggers only Prioritized paths. + // SpeculationPathStatusPrioritized is a path admitted under the queue's + // build budget: it is cleared to build. SpeculationPathStatusPrioritized SpeculationPathStatus = "prioritized" - // SpeculationPathStatusBuilding is a path a build signal has confirmed is in + // SpeculationPathStatusBuilding is a path whose build is confirmed in // flight; its BuildID is known. SpeculationPathStatusBuilding SpeculationPathStatus = "building" // SpeculationPathStatusPassed is a path whose build succeeded. SpeculationPathStatusPassed SpeculationPathStatus = "passed" // SpeculationPathStatusFailed is a path whose build failed. SpeculationPathStatusFailed SpeculationPathStatus = "failed" - // SpeculationPathStatusCancelling is a path whose in-flight build the - // controller has asked to stop but whose cancellation is not yet confirmed. It - // mirrors the batch-level cancelling intent (BatchStateCancelling): the cancel - // decision is recorded here and the effector drives it to terminal Cancelled. - // A path with no build in flight is dropped straight to Cancelled instead. + // SpeculationPathStatusCancelling is a path whose in-flight build has been + // asked to stop but whose cancellation is not yet confirmed. It mirrors the + // batch-level cancelling intent (BatchStateCancelling). A path with no build + // in flight is dropped straight to Cancelled instead. SpeculationPathStatusCancelling SpeculationPathStatus = "cancelling" // SpeculationPathStatusCancelled is the terminal state for a path that is no // longer pursued — its in-flight build was confirmed stopped, or the path was @@ -83,99 +73,82 @@ const ( SpeculationPathStatusCancelled SpeculationPathStatus = "cancelled" ) -// SpeculationPathAction is the decision a seam (the selector or the prioritizer) -// asks the controller to take for a path. It names the decision, not its effect: -// it is ephemeral (recomputed every time a seam runs) and never persisted. The -// controller maps it to the corresponding SpeculationPathStatus transition — -// applied under the tree's optimistic lock (Version) — and records the result; -// the seams never write status themselves. +// SpeculationPathAction is a requested action for a single speculation path. +// It names the decision, not its effect — applying it yields the corresponding +// SpeculationPathStatus transition — and it is ephemeral: recomputed each time +// decisions are made, never persisted. type SpeculationPathAction string const ( - // SpeculationPathActionUnknown is the unreachable zero value. A seam - // expresses "leave this path as-is" by omitting the path from its decisions, - // not by returning this. + // SpeculationPathActionUnknown is the unreachable zero value. "Leave this + // path as-is" is expressed by omitting the path from the decision set, not + // by this value. SpeculationPathActionUnknown SpeculationPathAction = "" - // SpeculationPathActionPromote asks the controller to advance this path one - // stage toward running. The target status depends on which seam decided it: - // the selector's promote moves a path to Selected; the prioritizer's promote - // moves it to Prioritized (cleared to build). + // SpeculationPathActionPromote advances the path one stage toward building: + // to Selected when decided per batch (selection), to Prioritized — cleared + // to build — when decided queue-wide (prioritization). SpeculationPathActionPromote SpeculationPathAction = "promote" - // SpeculationPathActionCancel asks the controller to stop pursuing this path: - // it moves to Cancelling if a build is in flight (the effector then confirms - // the stop and drives it to terminal Cancelled), or straight to Cancelled if - // no build has started. + // SpeculationPathActionCancel stops pursuing the path: it moves to + // Cancelling if a build is in flight, or straight to Cancelled if no build + // has started. SpeculationPathActionCancel SpeculationPathAction = "cancel" ) // SpeculationPathInfo is the per-path entry in a speculation tree: a path, its -// latest predicted-success score, its controller-owned status, and a link to -// the build dispatched for it (if any). ID and Path are immutable once the -// entry is persisted; Score, Status, and BuildID are updateable, written only -// by the controller under the tree's Version optimistic lock. +// latest predicted-success score, its status, and a link to the build +// dispatched for it (if any). ID and Path are immutable once the entry is +// persisted; Score, Status, and BuildID are updateable under the tree's +// Version optimistic lock. type SpeculationPathInfo struct { - // ID identifies this path. It is assigned by the controller when the path - // entry is first persisted, immutable thereafter, and globally unique — - // not merely unique within its tree, because other entities key rows by it - // alone (SpeculationPathBuild.PathID is a primary key with no extra - // scoping column). Its format is the controller's choice and carries no - // meaning — never parse it. Everything outside the tree names a path by - // this ID: seam outputs (path scores, path decisions) and durable links - // from other entities all refer to it rather than restating the Base/Head - // split. + // ID identifies this path. It is assigned when the path entry is first + // persisted, immutable thereafter, and globally unique — not merely unique + // within its tree, because other entities key rows by it alone + // (SpeculationPathBuild.PathID is a primary key with no extra scoping + // column). Its format carries no meaning — never parse it. Everything + // outside the tree names a path by this ID (PathScore, PathDecision, + // SpeculationPathBuild) rather than restating the Base/Head split. ID string // Path is the Base/Head split this entry covers. Immutable: it identifies // the entry and never changes after the path is first persisted. Path SpeculationPath - // Score is the path's predicted-success score. Updateable: it is computed by - // the scorer and persisted by the controller, not set at enumeration — the - // enumerator produces structure only. It is dynamic: the controller re-runs - // the scorer on every respeculate (as dependencies land, dependency builds - // pass, or sibling paths fail), so the value tracks the latest state rather - // than a figure frozen when the path was first enumerated (~0 until the - // first pass). + // Score is the path's predicted-success score. Updateable: it is recomputed + // as the world changes (dependencies land, dependency builds pass, sibling + // paths fail), so it tracks the latest state rather than a figure frozen at + // enumeration (0 until first scored). Score float32 - // Status is the observed lifecycle state of the path. Updateable: written - // only by the controller; read by the decision seams (scorer, selector, - // prioritizer). + // Status is the observed lifecycle state of the path. Updateable. Status SpeculationPathStatus // BuildID holds the runner-minted build identifier (also the build store's - // primary key) for this path. Updateable: it is empty until the speculate - // controller's reconcile stamps it once a build exists for this path. + // primary key) for this path. Updateable: it is empty until a build exists + // for this path. BuildID string } -// PathScore is the path scorer's verdict for a single path: the -// path's identity and its freshly computed predicted-success score. It is the -// scorer seam's only output — the controller merges scores into the tree by -// path ID and persists them; tree structure and status never pass through the -// scorer. Like PathDecision, it is ephemeral and never persisted. +// PathScore is a freshly computed predicted-success score for a single path, +// named by its ID. Like PathDecision, it is ephemeral and never persisted; the +// score it carries lands in the tree entry (SpeculationPathInfo.Score). type PathScore struct { - // PathID identifies the scored path (SpeculationPathInfo.ID) within the - // tree the scorer was handed. + // PathID identifies the scored path (SpeculationPathInfo.ID). PathID string // Score is the path's predicted-success probability, in [0, 1]. Score float32 } -// PathDecision is a seam's decision for a single path: the action the -// controller should take for it. It is the output of both the selector (per -// batch) and the prioritizer (queue-wide), and is not persisted. A seam returns -// a decision only for the paths it wants to act on; omitted paths are left -// as-is. A seam must return at most one decision per path — the controller -// treats conflicting duplicates as a policy bug, applying the first and -// logging and skipping the rest. +// PathDecision is a requested action for a single speculation path, named by +// its ID. It is ephemeral and never persisted. A decision set covers only the +// paths to act on — omitted paths are left as-is — and carries at most one +// decision per path. type PathDecision struct { // PathID identifies the speculation path the action applies to - // (SpeculationPathInfo.ID), within the tree(s) the seam was handed. + // (SpeculationPathInfo.ID). PathID string - // Action is what the controller should do for the path. + // Action is the requested action for the path. Action SpeculationPathAction } // SpeculationTree is the set of candidate speculation paths for a batch, built -// from its dependency graph. BatchID is immutable; Paths is updateable (the -// controller overwrites it wholesale on every respeculate), guarded by Version. +// from its dependency graph. BatchID is immutable; Paths is updateable — +// overwritten wholesale on re-speculation — guarded by Version. type SpeculationTree struct { // BatchID is the batch for which this speculation tree is constructed. // Immutable: it identifies the tree. @@ -196,7 +169,17 @@ type SpeculationTree struct { // Version is the version of the object. It is used for optimistic locking: // updates are conditional on the persisted version matching the caller's // expected version. Versioning starts at 1 and is incremented for each - // change to the object; version arithmetic is owned by the controller, the - // store performs a pure conditional write. + // change to the object. Version int32 } + +// PathIndex returns the index of the entry in t.Paths whose ID is id, or -1 +// if none is. +func (t SpeculationTree) PathIndex(id string) int { + for i, p := range t.Paths { + if p.ID == id { + return i + } + } + return -1 +} diff --git a/submitqueue/entity/speculation_tree_test.go b/submitqueue/entity/speculation_tree_test.go index 8b6b0428..c50314b7 100644 --- a/submitqueue/entity/speculation_tree_test.go +++ b/submitqueue/entity/speculation_tree_test.go @@ -73,3 +73,51 @@ func TestSpeculationPath_Equal(t *testing.T) { }) } } + +func TestSpeculationTree_PathIndex(t *testing.T) { + tree := SpeculationTree{ + BatchID: "q/batch/3", + Paths: []SpeculationPathInfo{ + {ID: "q/batch/3/path/0", Path: SpeculationPath{Base: []string{"q/batch/1", "q/batch/2"}, Head: "q/batch/3"}}, + {ID: "q/batch/3/path/1", Path: SpeculationPath{Base: []string{"q/batch/1"}, Head: "q/batch/3"}}, + {ID: "q/batch/3/path/2", Path: SpeculationPath{Head: "q/batch/3"}}, + }, + } + + tests := []struct { + name string + id string + want int + }{ + { + name: "first path found", + id: "q/batch/3/path/0", + want: 0, + }, + { + name: "last path found", + id: "q/batch/3/path/2", + want: 2, + }, + { + name: "unknown id not found", + id: "q/batch/3/path/9", + want: -1, + }, + { + name: "empty id not found", + id: "", + want: -1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, tree.PathIndex(tt.id)) + }) + } +} + +func TestSpeculationTree_PathIndex_EmptyTree(t *testing.T) { + assert.Equal(t, -1, SpeculationTree{}.PathIndex("q/batch/1/path/0")) +} diff --git a/submitqueue/orchestrator/controller/dlq/BUILD.bazel b/submitqueue/orchestrator/controller/dlq/BUILD.bazel index 2952dcf1..d6aea342 100644 --- a/submitqueue/orchestrator/controller/dlq/BUILD.bazel +++ b/submitqueue/orchestrator/controller/dlq/BUILD.bazel @@ -9,14 +9,17 @@ go_library( "log.go", "mergeconflictsignal.go", "mergesignal.go", + "queue.go", "request.go", ], importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/dlq", visibility = ["//visibility:public"], deps = [ "//api/runway/messagequeue:go_default_library", + "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", + "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/storage:go_default_library", "@com_github_uber_go_tally//:go_default_library", @@ -33,6 +36,7 @@ go_test( "log_test.go", "mergeconflictsignal_test.go", "mergesignal_test.go", + "queue_test.go", "request_test.go", ], embed = [":go_default_library"], diff --git a/submitqueue/orchestrator/controller/dlq/queue.go b/submitqueue/orchestrator/controller/dlq/queue.go new file mode 100644 index 00000000..2b17cbdb --- /dev/null +++ b/submitqueue/orchestrator/controller/dlq/queue.go @@ -0,0 +1,166 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dlq + +import ( + "context" + "fmt" + + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "github.com/uber/submitqueue/submitqueue/entity" + "go.uber.org/zap" +) + +// queueController is the DLQ reconciler for the prioritize topic. Unlike +// every other primary topic, prioritize's payload carries a QueueID, not a +// request or batch identifier — the stage reconciles a whole queue's build +// budget in one pass, not a single entity. There is therefore no request or +// batch to drive to a terminal state here: the batches the failed round +// would have considered stay exactly where they were (Selected paths stay +// Selected, in-flight builds keep running). +// +// Reconciliation is therefore re-arming, not terminalizing: after logging +// the failure, the controller publishes a fresh prioritize round for the +// queue. Speculate republishes a round on every batch-level event, but a +// queue whose only pending work is waiting on prioritization (nothing +// building, no new requests arriving) may see no further events — without +// the requeue its batches would stay stranded until unrelated activity +// happened to trigger a round. +// +// The requeue cannot poison-loop on payload contents: a round carries only +// the queue name and recomputes entirely from live state, so there is no bad +// datum to replay. A round that fails persistently (e.g. storage down) +// cycles primary retry ladder -> DLQ -> requeue, throttled by the full +// ladder on every cycle and visible in DLQ metrics, and converges the first +// cycle after the fault heals. The requeued message ID is derived from the +// dead-lettered message's ID: distinct from the original round's ID so the +// queue's publish dedup does not swallow it, and deterministic per DLQ +// message so a redelivered DLQ message republishes the same ID and is +// coalesced instead of fanning out. +type queueController struct { + logger *zap.SugaredLogger + metricsScope tally.Scope + registry consumer.TopicRegistry + topicKey consumer.TopicKey + consumerGroup string +} + +// Verify queueController implements consumer.Controller at compile time. +var _ consumer.Controller = (*queueController)(nil) + +// NewDLQQueueController builds a DLQ controller for the prioritize topic. +// registry must resolve topickey.TopicKeyPrioritize — the topic the +// controller re-arms the queue on. +func NewDLQQueueController( + logger *zap.SugaredLogger, + scope tally.Scope, + registry consumer.TopicRegistry, + topicKey consumer.TopicKey, + consumerGroup string, +) consumer.Controller { + name := string(topicKey) + "_controller" + return &queueController{ + logger: logger.Named(name), + metricsScope: scope.SubScope(name), + registry: registry, + topicKey: topicKey, + consumerGroup: consumerGroup, + } +} + +// Process logs a DLQ'd prioritize message for operator visibility, re-arms +// the queue by publishing a fresh prioritize round, and acks. See the +// queueController doc comment for why re-arming is the reconciliation here +// and why it cannot loop unboundedly. +func (c *queueController) Process(ctx context.Context, delivery consumer.Delivery) (retErr error) { + const opName = "process" + + op := metrics.Begin(c.metricsScope, opName) + defer func() { op.Complete(retErr) }() + + msg := delivery.Message() + + qid, err := entity.QueueIDFromBytes(msg.Payload) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) + return fmt.Errorf("failed to decode queue id from dlq payload: %w", err) + } + if qid.Name == "" { + metrics.NamedCounter(c.metricsScope, opName, "empty_id_errors", 1) + return fmt.Errorf("dlq payload decoded to empty queue name") + } + + dmeta := delivery.Metadata() + c.logger.Warnw("dlq message received; requeueing a fresh prioritization round for queue", + "queue", qid.Name, + "attempt", delivery.Attempt(), + "dlq_original_topic", dmeta["dlq.original_topic"], + "dlq_failure_count", dmeta["dlq.failure_count"], + "dlq_last_error", dmeta["dlq.last_error"], + ) + + if err := c.requeueRound(ctx, qid, msg.ID); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) + return fmt.Errorf("failed to requeue prioritize round for queue %s: %w", qid.Name, err) + } + + metrics.NamedCounter(c.metricsScope, opName, "reconciled", 1) + return nil +} + +// requeueRound publishes a fresh prioritize round for the queue. The message +// ID appends a suffix to the dead-lettered message's ID (see the +// queueController doc comment for the dedup reasoning). +func (c *queueController) requeueRound(ctx context.Context, qid entity.QueueID, dlqMsgID string) error { + payload, err := qid.ToBytes() + if err != nil { + return fmt.Errorf("failed to serialize queue ID: %w", err) + } + + msg := entityqueue.NewMessage(dlqMsgID+"/dlq-requeue", payload, qid.Name, nil) + + q, ok := c.registry.Queue(topickey.TopicKeyPrioritize) + if !ok { + return fmt.Errorf("no queue registered for topic key %s", topickey.TopicKeyPrioritize) + } + topicName, ok := c.registry.TopicName(topickey.TopicKeyPrioritize) + if !ok { + return fmt.Errorf("no topic name registered for topic key %s", topickey.TopicKeyPrioritize) + } + + if err := q.Publisher().Publish(ctx, topicName, msg); err != nil { + return fmt.Errorf("failed to publish message: %w", err) + } + return nil +} + +// Name returns the controller name for logging and metrics. +func (c *queueController) Name() string { + return string(c.topicKey) +} + +// TopicKey returns the topic key this controller subscribes to. +func (c *queueController) TopicKey() consumer.TopicKey { + return c.topicKey +} + +// ConsumerGroup returns the consumer group for offset tracking. +func (c *queueController) ConsumerGroup() string { + return c.consumerGroup +} diff --git a/submitqueue/orchestrator/controller/dlq/queue_test.go b/submitqueue/orchestrator/controller/dlq/queue_test.go new file mode 100644 index 00000000..a0f1dccc --- /dev/null +++ b/submitqueue/orchestrator/controller/dlq/queue_test.go @@ -0,0 +1,132 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dlq + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "github.com/uber/submitqueue/submitqueue/entity" + "go.uber.org/mock/gomock" + "go.uber.org/zap/zaptest" +) + +// newPrioritizeRegistry wires a TopicRegistry exposing the prioritize topic +// backed by a mock publisher, returning both so tests can set expectations. +func newPrioritizeRegistry(t *testing.T, ctrl *gomock.Controller) (consumer.TopicRegistry, *queuemock.MockPublisher) { + pub := queuemock.NewMockPublisher(ctrl) + q := queuemock.NewMockQueue(ctrl) + q.EXPECT().Publisher().Return(pub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyPrioritize, Name: "prioritize", Queue: q}, + }) + require.NoError(t, err) + return registry, pub +} + +func TestDLQQueueController_InterfaceAndAccessors(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newPrioritizeRegistry(t, ctrl) + + c := NewDLQQueueController(zaptest.NewLogger(t).Sugar(), testScope(), registry, TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq") + + assert.Equal(t, "prioritize_dlq", c.Name()) + assert.Equal(t, consumer.TopicKey("prioritize_dlq"), c.TopicKey()) + assert.Equal(t, "orchestrator-prioritize-dlq", c.ConsumerGroup()) +} + +// TestDLQQueueController_Process_RequeuesRound verifies the reconciler +// re-arms the queue: a fresh prioritize round is published for the queue, +// with a message ID derived from the dead-lettered message's ID (distinct +// from the original round so publish dedup does not swallow it, stable +// across DLQ redeliveries so they coalesce). +func TestDLQQueueController_Process_RequeuesRound(t *testing.T) { + ctrl := gomock.NewController(t) + registry, pub := newPrioritizeRegistry(t, ctrl) + + c := NewDLQQueueController(zaptest.NewLogger(t).Sugar(), testScope(), registry, TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq") + + payload, err := entity.QueueID{Name: "q"}.ToBytes() + require.NoError(t, err) + + var published entityqueue.Message + pub.EXPECT().Publish(gomock.Any(), "prioritize", gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + published = msg + return nil + }, + ) + + delivery := newMockDelivery(ctrl, payload) + require.NoError(t, c.Process(context.Background(), delivery)) + + assert.Equal(t, "dlq-msg-1/dlq-requeue", published.ID, "requeued ID should derive from the DLQ message's ID") + assert.Equal(t, "q", published.PartitionKey, "rounds are partitioned by queue name") + qid, err := entity.QueueIDFromBytes(published.Payload) + require.NoError(t, err) + assert.Equal(t, "q", qid.Name) +} + +// TestDLQQueueController_Process_PublishFailureNacks verifies a failed +// requeue publish surfaces as an error so the DLQ consumer redelivers and +// the re-arm is eventually made. +func TestDLQQueueController_Process_PublishFailureNacks(t *testing.T) { + ctrl := gomock.NewController(t) + registry, pub := newPrioritizeRegistry(t, ctrl) + + c := NewDLQQueueController(zaptest.NewLogger(t).Sugar(), testScope(), registry, TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq") + + payload, err := entity.QueueID{Name: "q"}.ToBytes() + require.NoError(t, err) + + sentinel := errors.New("publish down") + pub.EXPECT().Publish(gomock.Any(), "prioritize", gomock.Any()).Return(sentinel) + + delivery := newMockDelivery(ctrl, payload) + require.ErrorIs(t, c.Process(context.Background(), delivery), sentinel) +} + +func TestDLQQueueController_Process_MalformedPayloadFails(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newPrioritizeRegistry(t, ctrl) + + c := NewDLQQueueController(zaptest.NewLogger(t).Sugar(), testScope(), registry, TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq") + + delivery := newMockDelivery(ctrl, []byte("garbage")) + err := c.Process(context.Background(), delivery) + require.Error(t, err) +} + +func TestDLQQueueController_Process_EmptyNameFails(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newPrioritizeRegistry(t, ctrl) + + c := NewDLQQueueController(zaptest.NewLogger(t).Sugar(), testScope(), registry, TopicKey(topickey.TopicKeyPrioritize), "orchestrator-prioritize-dlq") + + payload, err := entity.QueueID{Name: ""}.ToBytes() + require.NoError(t, err) + + delivery := newMockDelivery(ctrl, payload) + err = c.Process(context.Background(), delivery) + require.Error(t, err) +} diff --git a/submitqueue/orchestrator/controller/prioritize/BUILD.bazel b/submitqueue/orchestrator/controller/prioritize/BUILD.bazel new file mode 100644 index 00000000..536e0ab5 --- /dev/null +++ b/submitqueue/orchestrator/controller/prioritize/BUILD.bazel @@ -0,0 +1,40 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["prioritize.go"], + importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/prioritize", + visibility = ["//visibility:public"], + deps = [ + "//platform/base/messagequeue:go_default_library", + "//platform/consumer:go_default_library", + "//platform/metrics:go_default_library", + "//submitqueue/core/topickey:go_default_library", + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/prioritizer:go_default_library", + "//submitqueue/extension/storage:go_default_library", + "@com_github_uber_go_tally//:go_default_library", + "@org_uber_go_zap//:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["prioritize_test.go"], + embed = [":go_default_library"], + deps = [ + "//platform/base/messagequeue:go_default_library", + "//platform/consumer:go_default_library", + "//platform/extension/messagequeue/mock:go_default_library", + "//submitqueue/core/topickey:go_default_library", + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/prioritizer/mock:go_default_library", + "//submitqueue/extension/storage:go_default_library", + "//submitqueue/extension/storage/mock:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + "@com_github_uber_go_tally//:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + "@org_uber_go_zap//zaptest:go_default_library", + ], +) diff --git a/submitqueue/orchestrator/controller/prioritize/prioritize.go b/submitqueue/orchestrator/controller/prioritize/prioritize.go new file mode 100644 index 00000000..6bb53293 --- /dev/null +++ b/submitqueue/orchestrator/controller/prioritize/prioritize.go @@ -0,0 +1,408 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package prioritize implements the queue-wide reconcile stage that rations +// a shared build budget across every in-flight batch of a queue. +// +// Unlike every other pipeline stage, prioritize is not batch-scoped: its +// queue message carries only a queue name (entity.QueueID), and each +// invocation re-evaluates every Speculating batch's speculation tree for +// that queue together. It gathers the queue-wide set of candidate paths +// (Selected, Prioritized, or Building), hands them to the queue's +// prioritizer.Prioritizer, and applies the returned decisions: +// +// - Promote on a Selected path clears it to run (-> Prioritized). +// - Cancel on a Building path asks the build runner to stop it and marks +// it Cancelling; the build stage's own signal loop confirms the stop. +// - Cancel on a Prioritized path (no build yet) drops it straight to +// Cancelled. +// +// Each affected tree is persisted under its own optimistic lock, so a +// version conflict only nacks and re-derives that tree's part of the round +// on redelivery — the whole computation is a pure function of freshly read +// state, so recomputing it is always safe. After applying decisions, the +// controller republishes to the build topic for every batch whose tree has +// at least one Prioritized path with no build yet, not just newly promoted +// ones — this heals a build message dropped by a prior crash and is itself +// idempotent, since the build stage dedups on batch ID. +package prioritize + +import ( + "context" + "errors" + "fmt" + + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer" + "github.com/uber/submitqueue/submitqueue/extension/storage" + "go.uber.org/zap" +) + +// opName is the metric operation name shared by every emit in this file. +const opName = "process" + +// Controller consumes queue-wide prioritize messages, ranks every candidate +// speculation path in the queue against its build budget, applies the +// resulting decisions to each affected speculation tree, and republishes to +// build for paths cleared to run. +type Controller struct { + logger *zap.SugaredLogger + metricsScope tally.Scope + store storage.Storage + prioritizers prioritizer.Factory + registry consumer.TopicRegistry + topicKey consumer.TopicKey + consumerGroup string +} + +// Verify Controller implements consumer.Controller interface at compile time. +var _ consumer.Controller = (*Controller)(nil) + +// NewController creates a new prioritize controller for the orchestrator. +func NewController( + logger *zap.SugaredLogger, + scope tally.Scope, + store storage.Storage, + prioritizers prioritizer.Factory, + registry consumer.TopicRegistry, + topicKey consumer.TopicKey, + consumerGroup string, +) *Controller { + return &Controller{ + logger: logger.Named("prioritize_controller"), + metricsScope: scope.SubScope("prioritize_controller"), + store: store, + prioritizers: prioritizers, + registry: registry, + topicKey: topicKey, + consumerGroup: consumerGroup, + } +} + +// Process re-evaluates the build budget for one queue: it loads every +// Speculating batch's speculation tree, ranks the queue-wide candidate paths +// through the queue's prioritizer, applies the returned decisions, persists +// the affected trees, and republishes to build for any path now cleared to +// run. Returns nil to ack (success), or error to nack (retry) — the whole +// round is a pure function of freshly read state, so redelivery simply +// recomputes it. +func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (retErr error) { + op := metrics.Begin(c.metricsScope, opName) + defer func() { op.Complete(retErr) }() + + msg := delivery.Message() + + qid, err := entity.QueueIDFromBytes(msg.Payload) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "deserialize_errors", 1) + return fmt.Errorf("failed to deserialize queue ID: %w", err) + } + + batches, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, qid.Name, []entity.BatchState{entity.BatchStateSpeculating}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to get speculating batches for queue %s: %w", qid.Name, err) + } + if len(batches) == 0 { + metrics.NamedCounter(c.metricsScope, opName, "no_speculating_batches", 1) + return nil + } + + trees, err := c.loadTrees(ctx, batches) + if err != nil { + return err + } + if len(trees) == 0 { + metrics.NamedCounter(c.metricsScope, opName, "no_trees", 1) + return nil + } + + candidates := candidatesOf(trees) + + pf, err := c.prioritizers.For(prioritizer.Config{QueueName: qid.Name}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "prioritizer_errors", 1) + return fmt.Errorf("failed to get prioritizer for queue %s: %w", qid.Name, err) + } + decisions, err := pf.Prioritize(ctx, candidates) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "prioritizer_errors", 1) + return fmt.Errorf("failed to prioritize candidates for queue %s: %w", qid.Name, err) + } + + changed, err := c.applyDecisions(ctx, qid.Name, trees, decisions) + if err != nil { + return err + } + + // Persist before publish. If republishBuilds fails below (or the process + // crashes between the two), the message nacks and the redelivered round + // recomputes everything from the persisted state — it carries no memory + // of what this attempt picked, and needs none: paths promoted here are + // now Prioritized, so the prioritizer counts them as slot holders instead + // of re-promoting them, and republishBuilds derives "who needs a build + // message" from the trees themselves (any Prioritized path without a + // build) rather than from this round's decisions, so a dropped publish is + // healed on the next pass. + if err := c.persistTrees(ctx, trees, changed); err != nil { + return err + } + + if err := c.republishBuilds(ctx, qid.Name, trees); err != nil { + return err + } + + return nil +} + +// loadTrees loads the speculation tree for each batch, keyed by batch ID. +// A batch with no tree yet (storage.ErrNotFound) has not been speculated on +// yet and is skipped; any other error aborts the round. +func (c *Controller) loadTrees(ctx context.Context, batches []entity.Batch) (map[string]entity.SpeculationTree, error) { + trees := make(map[string]entity.SpeculationTree, len(batches)) + for _, b := range batches { + tree, err := c.store.GetSpeculationTreeStore().Get(ctx, b.ID) + if err != nil { + if errors.Is(err, storage.ErrNotFound) { + metrics.NamedCounter(c.metricsScope, opName, "tree_not_found_skipped", 1) + continue + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return nil, fmt.Errorf("failed to get speculation tree for batch %s: %w", b.ID, err) + } + trees[b.ID] = tree + } + return trees, nil +} + +// candidatesOf flattens every path across all loaded trees whose status +// represents a queue-wide interest: Selected (wants a slot), or +// Prioritized/Building (holds one). +func candidatesOf(trees map[string]entity.SpeculationTree) []entity.SpeculationPathInfo { + var candidates []entity.SpeculationPathInfo + for _, tree := range trees { + for _, p := range tree.Paths { + switch p.Status { + case entity.SpeculationPathStatusSelected, + entity.SpeculationPathStatusPrioritized, + entity.SpeculationPathStatusBuilding: + candidates = append(candidates, p) + } + } + } + return candidates +} + +// applyDecisions maps each prioritizer decision to a status transition on its +// tree's matching path, mutating trees in place. It returns the set of batch +// IDs whose tree actually changed, so the caller persists only those. +// +// Decisions are captured as intent, never enacted here: this controller does +// not talk to the build system. A Cancel on a Building path only flips it to +// Cancelling in the tree; the build stage — the sole owner of runner +// interaction — enacts the runner cancel when it processes the batch, retried +// for free on every build message until the build reaches a terminal state +// and the speculate reconcile settles the path to Cancelled. Persisting the +// intent before any side effect is what makes the flow crash-safe: a lost +// build message is healed by republishBuilds, not by remembering this round. +// +// A decision naming a batch or path the round did not load, or applying an +// action the path's current status does not support, is a policy bug in the +// prioritizer: it is logged as a warning and skipped rather than corrupting +// the tree. +func (c *Controller) applyDecisions( + ctx context.Context, + queue string, + trees map[string]entity.SpeculationTree, + decisions []entity.PathDecision, +) (map[string]bool, error) { + changed := make(map[string]bool) + + // Decisions name paths by ID only; recover each path's tree from the + // trees loaded this round rather than parsing anything out of the ID. + pathBatch := make(map[string]string) + for batchID, tree := range trees { + for _, p := range tree.Paths { + pathBatch[p.ID] = batchID + } + } + + seen := make(map[string]bool, len(decisions)) + for _, d := range decisions { + if seen[d.PathID] { + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: duplicate decision for path", + "queue", queue, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + seen[d.PathID] = true + + batchID, ok := pathBatch[d.PathID] + if !ok { + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: path not loaded this round", + "queue", queue, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + tree := trees[batchID] + + // pathBatch was built from these same trees, so the ID must resolve; + // guard anyway so an index bug degrades to a skipped decision. + idx := tree.PathIndex(d.PathID) + if idx == -1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: path not found in tree", + "queue", queue, + "batch_id", batchID, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + + info := tree.Paths[idx] + switch { + case d.Action == entity.SpeculationPathActionPromote && info.Status == entity.SpeculationPathStatusSelected: + info.Status = entity.SpeculationPathStatusPrioritized + metrics.NamedCounter(c.metricsScope, opName, "promoted", 1) + + case d.Action == entity.SpeculationPathActionCancel && info.Status == entity.SpeculationPathStatusBuilding: + info.Status = entity.SpeculationPathStatusCancelling + metrics.NamedCounter(c.metricsScope, opName, "cancelling", 1) + + case d.Action == entity.SpeculationPathActionCancel && info.Status == entity.SpeculationPathStatusPrioritized: + info.Status = entity.SpeculationPathStatusCancelled + metrics.NamedCounter(c.metricsScope, opName, "cancelled", 1) + + default: + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: action not valid for path status", + "queue", queue, + "batch_id", batchID, + "path_id", d.PathID, + "action", d.Action, + "status", string(info.Status), + ) + continue + } + + tree.Paths[idx] = info + trees[batchID] = tree + changed[batchID] = true + } + + return changed, nil +} + +// persistTrees writes each changed tree's paths under its own optimistic +// lock, bumping Version in trees on success. Version arithmetic is owned by +// the controller; the store performs a pure conditional write. +func (c *Controller) persistTrees(ctx context.Context, trees map[string]entity.SpeculationTree, changed map[string]bool) error { + for batchID := range changed { + tree := trees[batchID] + newVersion := tree.Version + 1 + if err := c.store.GetSpeculationTreeStore().Update(ctx, batchID, tree.Version, newVersion, tree.Paths); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to update speculation tree for batch %s: %w", batchID, err) + } + tree.Version = newVersion + trees[batchID] = tree + } + return nil +} + +// republishBuilds re-publishes a build message for every batch whose tree +// carries work the build stage still has to enact — a Prioritized path with +// no build yet (needs a trigger), or a Cancelling path (a persisted cancel +// intent whose runner cancel the build stage owns) — not just paths touched +// this round. This heals a build message a prior crash dropped between the +// tree write and the publish; it is safe to repeat because the build stage +// dedups triggers on the path->build mapping and runner cancels are +// idempotent. +func (c *Controller) republishBuilds(ctx context.Context, queue string, trees map[string]entity.SpeculationTree) error { + for batchID, tree := range trees { + needsBuild := false + for _, p := range tree.Paths { + if (p.Status == entity.SpeculationPathStatusPrioritized && p.BuildID == "") || + p.Status == entity.SpeculationPathStatusCancelling { + needsBuild = true + break + } + } + if !needsBuild { + continue + } + if err := c.publishBuild(ctx, batchID, queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) + return fmt.Errorf("failed to publish batch %s to build: %w", batchID, err) + } + metrics.NamedCounter(c.metricsScope, opName, "build_republished", 1) + } + return nil +} + +// publishBuild publishes a batch ID to the build topic. Only the identifier +// travels on the queue; the build controller reloads the full Batch (and, +// eventually, its speculation tree) from storage. +func (c *Controller) publishBuild(ctx context.Context, batchID, partitionKey string) error { + bid := entity.BatchID{ID: batchID} + payload, err := bid.ToBytes() + if err != nil { + return fmt.Errorf("failed to serialize batch ID: %w", err) + } + + msg := entityqueue.NewMessage(batchID, payload, partitionKey, nil) + + q, ok := c.registry.Queue(topickey.TopicKeyBuild) + if !ok { + return fmt.Errorf("no queue registered for topic key %s", topickey.TopicKeyBuild) + } + + topicName, ok := c.registry.TopicName(topickey.TopicKeyBuild) + if !ok { + return fmt.Errorf("no topic name registered for topic key %s", topickey.TopicKeyBuild) + } + + if err := q.Publisher().Publish(ctx, topicName, msg); err != nil { + return fmt.Errorf("failed to publish message: %w", err) + } + + return nil +} + +// Name returns the controller name for logging and metrics. +func (c *Controller) Name() string { + return "prioritize" +} + +// TopicKey returns the topic key this controller subscribes to. +func (c *Controller) TopicKey() consumer.TopicKey { + return c.topicKey +} + +// ConsumerGroup returns the consumer group for offset tracking. +func (c *Controller) ConsumerGroup() string { + return c.consumerGroup +} diff --git a/submitqueue/orchestrator/controller/prioritize/prioritize_test.go b/submitqueue/orchestrator/controller/prioritize/prioritize_test.go new file mode 100644 index 00000000..44364f57 --- /dev/null +++ b/submitqueue/orchestrator/controller/prioritize/prioritize_test.go @@ -0,0 +1,395 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package prioritize + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "github.com/uber/submitqueue/submitqueue/entity" + prioritizermock "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer/mock" + "github.com/uber/submitqueue/submitqueue/extension/storage" + storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" + "go.uber.org/mock/gomock" + "go.uber.org/zap/zaptest" +) + +// testHarness wires a Controller against mocked storage, prioritizer, build +// runner, and a build topic queue so tests can assert what got persisted and +// published. +type testHarness struct { + controller *Controller + batchStore *storagemock.MockBatchStore + treeStore *storagemock.MockSpeculationTreeStore + prio *prioritizermock.MockPrioritizer + buildPub *queuemock.MockPublisher +} + +func newTestHarness(t *testing.T, ctrl *gomock.Controller) *testHarness { + prio := prioritizermock.NewMockPrioritizer(ctrl) + prioFactory := prioritizermock.NewMockFactory(ctrl) + prioFactory.EXPECT().For(gomock.Any()).Return(prio, nil).AnyTimes() + + buildPub := queuemock.NewMockPublisher(ctrl) + buildQ := queuemock.NewMockQueue(ctrl) + buildQ.EXPECT().Publisher().Return(buildPub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyPrioritize, Name: "prioritize", Queue: buildQ}, + {Key: topickey.TopicKeyBuild, Name: "build", Queue: buildQ}, + }) + require.NoError(t, err) + + batchStore := storagemock.NewMockBatchStore(ctrl) + treeStore := storagemock.NewMockSpeculationTreeStore(ctrl) + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetSpeculationTreeStore().Return(treeStore).AnyTimes() + + c := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + store, + prioFactory, + registry, + topickey.TopicKeyPrioritize, + "orchestrator-prioritize", + ) + return &testHarness{ + controller: c, + batchStore: batchStore, + treeStore: treeStore, + prio: prio, + buildPub: buildPub, + } +} + +// queueDelivery builds a delivery whose payload is a QueueID, matching the +// on-queue contract for this stage. +func queueDelivery(t *testing.T, ctrl *gomock.Controller, queue string) consumer.Delivery { + t.Helper() + payload, err := entity.QueueID{Name: queue}.ToBytes() + require.NoError(t, err) + msg := entityqueue.NewMessage(queue, payload, queue, nil) + d := queuemock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + d.EXPECT().Attempt().Return(1).AnyTimes() + return d +} + +func TestController_Identity(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + assert.Equal(t, "prioritize", h.controller.Name()) + assert.Equal(t, topickey.TopicKeyPrioritize, h.controller.TopicKey()) + assert.Equal(t, "orchestrator-prioritize", h.controller.ConsumerGroup()) + + var _ consumer.Controller = h.controller +} + +// TestController_Process_NoSpeculatingBatches verifies an empty queue just +// acks without touching the prioritizer or publishing anything. +func TestController_Process_NoSpeculatingBatches(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return(nil, nil) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_TreeMissingSkipsBatch verifies a batch with no +// speculation tree yet (not yet speculated on) is skipped rather than +// erroring the whole round. +func TestController_Process_TreeMissingSkipsBatch(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(entity.SpeculationTree{}, storage.ErrNotFound) + + // No prioritizer, no publish expected — the harness mocks fail the test if called unexpectedly. + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_PromoteApplied verifies a Promote decision on a +// Selected path transitions it to Prioritized, persists the tree with the +// correct old/new version pair, and republishes the batch to build. +func TestController_Process_PromoteApplied(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Base: nil, Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 3, + Paths: []entity.SpeculationPathInfo{{ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusSelected, Score: 0.9}}, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return([]entity.PathDecision{ + {PathID: "q/batch/1/path/0", Action: entity.SpeculationPathActionPromote}, + }, nil) + h.treeStore.EXPECT(). + Update(gomock.Any(), "q/batch/1", int32(3), int32(4), gomock.AssignableToTypeOf([]entity.SpeculationPathInfo{})). + DoAndReturn(func(_ context.Context, _ string, _, _ int32, paths []entity.SpeculationPathInfo) error { + require.Len(t, paths, 1) + assert.Equal(t, entity.SpeculationPathStatusPrioritized, paths[0].Status) + return nil + }) + h.buildPub.EXPECT(). + Publish(gomock.Any(), "build", gomock.AssignableToTypeOf(entityqueue.Message{})). + DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { + bid, err := entity.BatchIDFromBytes(msg.Payload) + require.NoError(t, err) + assert.Equal(t, "q/batch/1", bid.ID) + assert.Equal(t, "q", msg.PartitionKey) + return nil + }) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_IllegalDecisionSkipped verifies a decision that does +// not match any loaded path is dropped without an Update call, and nothing +// is published (no path reached Prioritized). +func TestController_Process_IllegalDecisionSkipped(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 1, + Paths: []entity.SpeculationPathInfo{{ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusSelected}}, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + // Cancel on a Selected path is not a legal transition: skip it. + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return([]entity.PathDecision{ + {PathID: "q/batch/1/path/0", Action: entity.SpeculationPathActionCancel}, + }, nil) + // No treeStore.Update, no publish expected. + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_CancelBuildingCapturesIntent verifies Cancel on a +// Building path only flips it to Cancelling in the tree — intent capture, +// never a runner call from this controller — and that a build message is +// republished for the batch so the build stage enacts the cancel. +func TestController_Process_CancelBuildingCapturesIntent(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 2, + Paths: []entity.SpeculationPathInfo{ + {ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusBuilding, BuildID: "build-1"}, + }, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return([]entity.PathDecision{ + {PathID: "q/batch/1/path/0", Action: entity.SpeculationPathActionCancel}, + }, nil) + h.treeStore.EXPECT(). + Update(gomock.Any(), "q/batch/1", int32(2), int32(3), gomock.AssignableToTypeOf([]entity.SpeculationPathInfo{})). + DoAndReturn(func(_ context.Context, _ string, _, _ int32, paths []entity.SpeculationPathInfo) error { + require.Len(t, paths, 1) + assert.Equal(t, entity.SpeculationPathStatusCancelling, paths[0].Status) + return nil + }) + // The Cancelling path is a persisted intent the build stage must enact, + // so a build message is republished for the batch. + h.buildPub.EXPECT().Publish(gomock.Any(), "build", gomock.Any()).Return(nil) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_CancelPrioritizedNoRunnerCall verifies Cancel on a +// Prioritized path (no build started yet) drops straight to Cancelled — +// there is no in-flight work, so no intent to hand to the build stage and no +// build republish. +func TestController_Process_CancelPrioritizedNoRunnerCall(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 5, + Paths: []entity.SpeculationPathInfo{ + {ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusPrioritized}, + }, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return([]entity.PathDecision{ + {PathID: "q/batch/1/path/0", Action: entity.SpeculationPathActionCancel}, + }, nil) + h.treeStore.EXPECT(). + Update(gomock.Any(), "q/batch/1", int32(5), int32(6), gomock.AssignableToTypeOf([]entity.SpeculationPathInfo{})). + DoAndReturn(func(_ context.Context, _ string, _, _ int32, paths []entity.SpeculationPathInfo) error { + require.Len(t, paths, 1) + assert.Equal(t, entity.SpeculationPathStatusCancelled, paths[0].Status) + return nil + }) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +// TestController_Process_VersionMismatchErrors verifies ErrVersionMismatch +// from the tree Update surfaces as an error (nack; the round is recomputed +// on redelivery). +func TestController_Process_VersionMismatchErrors(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 1, + Paths: []entity.SpeculationPathInfo{{ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusSelected}}, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return([]entity.PathDecision{ + {PathID: "q/batch/1/path/0", Action: entity.SpeculationPathActionPromote}, + }, nil) + h.treeStore.EXPECT(). + Update(gomock.Any(), "q/batch/1", int32(1), int32(2), gomock.Any()). + Return(storage.ErrVersionMismatch) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.Error(t, err) + assert.True(t, errors.Is(err, storage.ErrVersionMismatch)) +} + +// TestController_Process_PrioritizerErrors verifies a Prioritize failure +// surfaces as an error without touching storage further. +func TestController_Process_PrioritizerErrors(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 1, + Paths: []entity.SpeculationPathInfo{{ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusSelected}}, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return(nil, errors.New("policy boom")) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.Error(t, err) +} + +// TestController_Process_RepublishesPreExistingPrioritized verifies a batch +// whose tree already has a Prioritized path with no BuildID gets republished +// to build even when the prioritizer returns zero decisions this round — +// self-healing a dropped build message. +func TestController_Process_RepublishesPreExistingPrioritized(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + path := entity.SpeculationPath{Head: "q/batch/1"} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", State: entity.BatchStateSpeculating} + tree := entity.SpeculationTree{ + BatchID: "q/batch/1", + Version: 4, + Paths: []entity.SpeculationPathInfo{ + {ID: "q/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusPrioritized, BuildID: ""}, + }, + } + + h.batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", []entity.BatchState{entity.BatchStateSpeculating}). + Return([]entity.Batch{batch}, nil) + h.treeStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(tree, nil) + h.prio.EXPECT().Prioritize(gomock.Any(), gomock.Any()).Return(nil, nil) + // No treeStore.Update expected: nothing changed this round. + h.buildPub.EXPECT(). + Publish(gomock.Any(), "build", gomock.AssignableToTypeOf(entityqueue.Message{})). + DoAndReturn(func(_ context.Context, _ string, msg entityqueue.Message) error { + bid, err := entity.BatchIDFromBytes(msg.Payload) + require.NoError(t, err) + assert.Equal(t, "q/batch/1", bid.ID) + return nil + }) + + err := h.controller.Process(context.Background(), queueDelivery(t, ctrl, "q")) + require.NoError(t, err) +} + +func TestController_Process_MalformedPayload(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + + msg := entityqueue.NewMessage("bad", []byte(`{"invalid"`), "q", nil) + d := queuemock.NewMockDelivery(ctrl) + d.EXPECT().Message().Return(msg).AnyTimes() + d.EXPECT().Attempt().Return(1).AnyTimes() + + err := h.controller.Process(context.Background(), d) + require.Error(t, err) +}