From 61ff6ccec77169c23a78a1ffa52ae0266a77dbe9 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 9 Jul 2026 13:27:44 -0700 Subject: [PATCH 1/2] feat(orchestrator): add queue-wide prioritize stage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Selection is per batch and blind to other batches, so it cannot ration a shared build budget: if every batch promoted generously, their combined demand could swamp CI. The speculation design closes that gap with a prioritization step that sees every candidate path across a queue's in-flight batches and admits only what fits the queue's concurrent-build budget. That reconcile is queue-scoped, not batch-scoped, so it gets its own pipeline stage rather than piggybacking on the per-batch speculate flow. ### What? New `prioritize` topic (payload: `entity.QueueID`, partitioned by queue name) and controller. `entity.QueueID` is introduced here with ToBytes/QueueIDFromBytes, mirroring the BatchID payload pattern for queue-scoped stages. Each invocation is a full budget round for one queue: load every Speculating batch's speculation tree (skipping batches not yet speculated), flatten the queue-wide candidates (Selected / Prioritized / Building paths), hand them to the queue's `prioritizer.Prioritizer`, and apply the sparse decisions as captured intent — Promote flips Selected→Prioritized; Cancel on a Building path flips it to Cancelling (this controller never talks to the build system: the build stage owns runner interaction and enacts the persisted intent, retried on every build message until the build terminates); Cancel on a not-yet-building Prioritized path drops it straight to Cancelled; illegal decisions are logged and skipped so a policy bug cannot corrupt tree state. Path lookup for decision application lives on the entity (`SpeculationTree.PathIndex`), keyed by the path ID decisions carry; the controller maps each ID back to its tree via the candidates it loaded this round (never by parsing the ID), and treats a duplicate decision for the same path as a policy bug — logged and skipped like any other illegal decision. Each affected tree is persisted under its own optimistic lock (version arithmetic in the controller); any conflict nacks the round, which is safe to replay because decisions are recomputed from freshly read state — a redelivered round carries no memory of what a previous attempt picked and needs none, since already-promoted paths are counted as slot holders rather than re-promoted. After applying, the controller republishes to build for every batch whose tree carries work the build stage still has to enact — a Prioritized path with no build, or a Cancelling intent — healing dropped build messages idempotently. Wiring: topic registered with an automatic DLQ pair. The DLQ reconciler re-arms the queue rather than terminalizing an entity: it logs the failure and republishes a fresh prioritize round under a distinct, deterministic message ID, so a queue whose only pending work is waiting on prioritization is not stranded when a round dead-letters. The requeue cannot poison-loop — a round carries only the queue name and recomputes from live state — and a persistently failing round cycles retry-ladder→DLQ→requeue at full-ladder cadence, visible in metrics, converging once the fault heals. The sticky prioritizer (never preempts) backed by a static admit-all parity limit is wired as the default per-queue profile. Tree-entity docs are cleaned to describe the data itself (states, invariants, uniqueness) rather than narrating which stage reads or writes what. Nothing publishes to the topic yet — the speculate rework turns it on. ## Test Plan ✅ `make gazelle && make fmt && make test` and `bazel build //service/submitqueue/orchestrator/...`. Controller unit tests cover: empty queue ack, missing-tree skip, promote transition + versioned update + build republish, illegal-decision skip, cancel-on-building capturing intent (Cancelling persisted, build republished, no runner dependency at all), cancel-on-prioritized dropping straight to Cancelled, version-mismatch and prioritizer errors nacking, and republish for pre-existing prioritized paths with zero new decisions. DLQ reconciler tests cover the requeued round's ID/partition/payload, publish-failure nack, and malformed/empty payload rejection. Entity tests cover `SpeculationTree.PathIndex` (order-sensitive base matching, empty tree). --- .../orchestrator/server/BUILD.bazel | 4 + .../submitqueue/orchestrator/server/main.go | 62 ++- submitqueue/core/topickey/topickey.go | 8 + submitqueue/entity/BUILD.bazel | 2 + submitqueue/entity/queue.go | 36 ++ submitqueue/entity/queue_test.go | 72 ++++ submitqueue/entity/speculation_tree.go | 159 +++---- submitqueue/entity/speculation_tree_test.go | 48 +++ .../orchestrator/controller/dlq/BUILD.bazel | 4 + .../orchestrator/controller/dlq/queue.go | 166 +++++++ .../orchestrator/controller/dlq/queue_test.go | 132 ++++++ .../controller/prioritize/BUILD.bazel | 40 ++ .../controller/prioritize/prioritize.go | 408 ++++++++++++++++++ .../controller/prioritize/prioritize_test.go | 395 +++++++++++++++++ 14 files changed, 1439 insertions(+), 97 deletions(-) create mode 100644 submitqueue/entity/queue.go create mode 100644 submitqueue/entity/queue_test.go create mode 100644 submitqueue/orchestrator/controller/dlq/queue.go create mode 100644 submitqueue/orchestrator/controller/dlq/queue_test.go create mode 100644 submitqueue/orchestrator/controller/prioritize/BUILD.bazel create mode 100644 submitqueue/orchestrator/controller/prioritize/prioritize.go create mode 100644 submitqueue/orchestrator/controller/prioritize/prioritize_test.go 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) +} From e30249ea5051ea9ac544c5c0cf52ca6b459de3db Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 10 Jul 2026 07:58:36 -0700 Subject: [PATCH 2/2] feat(speculate): write the speculation tree in shadow mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? This is the first of four slices reproducing the tree-driven speculate/build/buildsignal rework as an incremental, always-shippable stack instead of one large commit. Each slice keeps the orchestrator working end to end so the e2e suite stays green throughout, and reviewers can follow the pipeline shift (direct publish -> tree -> prioritize -> build) one deliberate step at a time. This slice only adds the speculation tree as bookkeeping: nothing downstream reads it yet, so the batch's forward step (when it builds, when it merges) is untouched — with one deliberate exception called out below: a liveness fix to the dependent wake-up that review of this slice surfaced. ### What? speculate's Controller now takes an enumerator, path scorer, selector, and dependency-limit factory, and every Created/Scored/Speculating pass loads or creates the batch's entity.SpeculationTree, applies the scorer's and selector's outputs, and persists it only if something changed (the apply steps report whether they mutated anything, gating a version+1 conditional Update). The controller stays pure mechanics: which paths exist, how they score, and which are promoted or cancelled are the seams' decisions alone — the controller validates and records them. Tree creation is gated by the queue's dependency limit; the controller wraps the enumerator's structure-only paths into persisted entries itself — stamping Candidate, minting each path's immutable ID as `{batchID}/path/{i}`, and skipping structural duplicates as a contract violation — and a concurrent create race re-reads the winner's tree instead of erroring. Seam outputs are consumed by path ID: scores merged by ID (unknown IDs skipped, out-of-[0,1] values clamped, both logged), selector decisions resolved through the ID-keyed PathIndex with duplicates logged and skipped. The pre-existing forward step is otherwise unchanged: Created/Scored batches CAS to Speculating and publish straight to build, and Speculating batches run the original tryFinalize (merge once every dependency has landed, cascade-fail on a failed dependency). Cancelling is the unmodified pre-existing flow. Liveness fix in the forward step: a batch waiting on dependencies (in tryFinalize, or now in the dependency gate) was only ever woken when a dependency was cancelled — a dependency reaching Succeeded or Failed never re-published its dependents, because mergesignal routes every terminal transition back through speculate under the batch's own ID and the terminal branch only re-published conclude. The terminal branch now fans out to dependents for every terminal state, and failOnDependency wakes its own dependents right after the terminal CAS so failures cascade downstream. The fan-out publishes a wake for every listed dependent; a dangling reverse-index entry (left by an abandoned batch creation) surfaces as a not-found speculate message that dead-letters and is skipped by the DLQ reconciler — eliminating that class at the source, by creating the batch before the reverse-index update, is tracked in #354. Cancel decisions on tree paths are recorded as status only (Cancelling on an in-flight build) — nothing in this controller calls a build runner, which is why it takes no buildrunner.Factory. main.go wires the new seams with parity defaults: chain enumerates a single path per batch, the probability path scorer, the all-selector promotes every candidate, and a static dependency limit that is effectively ungated. ## Test Plan `bazel test //submitqueue/... //service/...` — 56/56 targets pass. New unit coverage for the wake-up fix: dependent fan-out on every terminal state (with publish-order assertions), missing BatchDependent row surfaced as an invariant error, and dependent-publish failure after the terminal CAS nacking for redelivery convergence. `make gazelle && make fmt` — no diffs beyond the intended BUILD file updates. ## Issue Part of the speculation rework. Interim fan-out caveats tracked in #354. --- .../orchestrator/server/BUILD.bazel | 8 + .../submitqueue/orchestrator/server/main.go | 81 +- .../controller/speculate/BUILD.bazel | 12 + .../controller/speculate/speculate.go | 436 ++++++++- .../controller/speculate/speculate_test.go | 909 ++++++++++-------- 5 files changed, 1012 insertions(+), 434 deletions(-) diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 2ccb3493..022f999c 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -41,9 +41,17 @@ 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/dependencylimit:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/static:go_default_library", + "//submitqueue/extension/speculation/enumerator:go_default_library", + "//submitqueue/extension/speculation/enumerator/chain:go_default_library", + "//submitqueue/extension/speculation/pathscorer:go_default_library", + "//submitqueue/extension/speculation/pathscorer/probability: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/speculation/selector:go_default_library", + "//submitqueue/extension/speculation/selector/all:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mysql:go_default_library", "//submitqueue/extension/validator/fake:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 4ddda8cf..8a322d4f 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -61,9 +61,17 @@ 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" + "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit" + dependencylimitstatic "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit/static" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator/chain" + "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer" + pathscorerprobability "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer/probability" 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/speculation/selector" + selectorall "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/all" "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" @@ -238,7 +246,7 @@ func run() error { // back to a baseline profile for queues without an explicit entry. This is // the single place queue topology is known; the extension packages stay // queue-agnostic. - queues, err := newQueueRegistry(logger, scope, changeset.New(store.GetRequestStore(), store.GetChangeStore())) + queues, err := newQueueRegistry(logger, scope, changeset.New(store.GetRequestStore(), store.GetChangeStore()), store.GetBatchStore()) if err != nil { return fmt.Errorf("failed to build queue registry: %w", err) } @@ -249,9 +257,13 @@ func run() error { scf := scorerFactory{queues} cof := analyzerFactory{queues} prf := prioritizerFactory{queues} + enf := enumeratorFactory{queues} + ssf := speculationScorerFactory{queues} + slf := speculationSelectorFactory{queues} + dlf := dependencyLimitFactory{queues} // Register controllers - primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, prf, cnt, store) + primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, prf, enf, ssf, slf, dlf, cnt, store) if err != nil { return err } @@ -502,11 +514,15 @@ func newTopicRegistry(q extqueue.Queue, subscriberName string) (consumer.TopicRe // read as "for this queue, here are its scorer, analyzer, change provider, …", and lets // a queue profile start from a baseline and override only what differs. type queueExtensions struct { - changeProvider changeprovider.ChangeProvider - buildRunner buildrunner.BuildRunner - scorer scorer.Scorer - analyzer conflict.Analyzer - prioritizer prioritizer.Prioritizer + changeProvider changeprovider.ChangeProvider + buildRunner buildrunner.BuildRunner + scorer scorer.Scorer + analyzer conflict.Analyzer + prioritizer prioritizer.Prioritizer + enumerator enumerator.Enumerator + speculationScorer pathscorer.Scorer + speculationSelector selector.Selector + dependencyLimit dependencylimit.DependencyLimit } // queueRegistry maps a queue name to its extensions, falling back to a default @@ -558,7 +574,31 @@ func (f prioritizerFactory) For(cfg prioritizer.Config) (prioritizer.Prioritizer 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) { +type enumeratorFactory struct{ reg queueRegistry } + +func (f enumeratorFactory) For(cfg enumerator.Config) (enumerator.Enumerator, error) { + return f.reg.get(cfg.QueueName).enumerator, nil +} + +type speculationScorerFactory struct{ reg queueRegistry } + +func (f speculationScorerFactory) For(cfg pathscorer.Config) (pathscorer.Scorer, error) { + return f.reg.get(cfg.QueueName).speculationScorer, nil +} + +type speculationSelectorFactory struct{ reg queueRegistry } + +func (f speculationSelectorFactory) For(cfg selector.Config) (selector.Selector, error) { + return f.reg.get(cfg.QueueName).speculationSelector, nil +} + +type dependencyLimitFactory struct{ reg queueRegistry } + +func (f dependencyLimitFactory) For(cfg dependencylimit.Config) (dependencylimit.DependencyLimit, error) { + return f.reg.get(cfg.QueueName).dependencyLimit, 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, enf enumerator.Factory, ssf pathscorer.Factory, slf selector.Factory, dlf dependencylimit.Factory, cnt counter.Counter, store storage.Storage) (int, error) { var count int requestController := start.NewController( logger, @@ -648,6 +688,10 @@ func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, logger, scope, store, + enf, + ssf, + slf, + dlf, registry, topickey.TopicKeySpeculate, "orchestrator-speculate", @@ -899,6 +943,11 @@ func newPhabChangeProvider(logger *zap.Logger, scope tally.Scope) (changeprovide // effectively admit-all — until per-queue budgets are configured. const defaultPrioritizationLimit = 1000 +// defaultDependencyLimit is the baseline cap on active dependencies a batch +// may speculate over. It is a parity default — effectively ungated — until +// per-queue limits are configured. +const defaultDependencyLimit = 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 @@ -906,7 +955,7 @@ const defaultPrioritizationLimit = 1000 // conflict analyzer. Queues without an explicit profile fall back to the // baseline. This is the one place queue topology lives; extension packages stay // queue-agnostic. -func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset.Resolver) (queueRegistry, error) { +func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset.Resolver, batchStore storage.BatchStore) (queueRegistry, error) { cp, err := newChangeProvider(logger, scope) if err != nil { return queueRegistry{}, fmt.Errorf("failed to create change provider: %w", err) @@ -933,6 +982,12 @@ func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset. // 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. + // + // The speculation seams default to the single-chain parity policies: + // chain enumerates one path per batch (built on the full ordered + // dependency chain), probability scores paths from the batches' + // predicted-success probabilities and resolved outcomes, all promotes + // every candidate, and the dependency limit is effectively ungated. base := queueExtensions{ changeProvider: cp, buildRunner: buildfake.New(resolver), @@ -943,8 +998,12 @@ 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), - prioritizer: sticky.New(prioritizationlimitstatic.New(defaultPrioritizationLimit)), + analyzer: conflictfake.New(all.New(), nil), + prioritizer: sticky.New(prioritizationlimitstatic.New(defaultPrioritizationLimit)), + enumerator: chain.New(), + speculationScorer: pathscorerprobability.New(batchStore), + speculationSelector: selectorall.New(), + dependencyLimit: dependencylimitstatic.New(defaultDependencyLimit), } // test-queue: bucketed heuristic scorer; conservative (serialized) conflicts diff --git a/submitqueue/orchestrator/controller/speculate/BUILD.bazel b/submitqueue/orchestrator/controller/speculate/BUILD.bazel index 39e36977..b6f312da 100644 --- a/submitqueue/orchestrator/controller/speculate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/speculate/BUILD.bazel @@ -11,6 +11,10 @@ go_library( "//platform/metrics:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/dependencylimit:go_default_library", + "//submitqueue/extension/speculation/enumerator:go_default_library", + "//submitqueue/extension/speculation/pathscorer:go_default_library", + "//submitqueue/extension/speculation/selector:go_default_library", "//submitqueue/extension/storage:go_default_library", "@com_github_uber_go_tally//:go_default_library", "@org_uber_go_zap//:go_default_library", @@ -28,6 +32,14 @@ go_test( "//platform/extension/messagequeue/mock:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/fake:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/mock:go_default_library", + "//submitqueue/extension/speculation/enumerator/fake:go_default_library", + "//submitqueue/extension/speculation/enumerator/mock:go_default_library", + "//submitqueue/extension/speculation/pathscorer/fake:go_default_library", + "//submitqueue/extension/speculation/pathscorer/mock:go_default_library", + "//submitqueue/extension/speculation/selector/fake:go_default_library", + "//submitqueue/extension/speculation/selector/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", diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index f2d2b21b..50380d2d 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -25,30 +25,48 @@ import ( "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/dependencylimit" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator" + "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/selector" "github.com/uber/submitqueue/submitqueue/extension/storage" "go.uber.org/zap" ) // Controller handles speculate queue messages. // -// Naive happy-path algorithm: assume every in-flight build will pass and -// treat batch.Dependencies + [batch.ID] as the single speculation chain. -// Per invocation, the controller advances the batch one step in the -// state machine: +// Each invocation reconciles the batch's entity.SpeculationTree and advances +// the batch one step in the state machine. Tree reconciliation is pure +// mechanics: the controller runs the queue's configured seams — enumerator +// (path structure), path scorer (path scores), selector (path decisions) — +// validates their outputs, applies them to the persisted tree, and writes it +// back under optimistic concurrency. Which paths exist, how they score, and +// which are promoted or cancelled are the seams' decisions alone; the +// controller never originates one. No downstream stage reads the tree yet, +// so the forward step below is driven by the batch's own state, not the +// tree. // -// - Created or Scored → publish to build, transition to Speculating. -// - Speculating → if all deps are Succeeded, publish to merge and -// transition to Merging; otherwise no-op (or fail-fast if a dep is -// in a non-succeeding terminal state). +// Per invocation, the controller advances the batch one step in the state +// machine: +// +// - Created, Scored, or Speculating → speculateBatch: reconcile the tree +// (created on the first pass the dependency gate admits), then advance +// the batch — publish to build and CAS to Speculating for +// Created/Scored, or tryFinalize for Speculating. // - Cancelling → cancel any in-flight Build entity, respeculate // dependents, CAS to terminal Cancelled, publish to conclude. The // cancel controller hands the batch off in this state and speculate // drives it to terminal. // - Merging → no-op (owned by the merge controller). -// - Terminal → re-fan-out to conclude for self-healing in case a -// prior publish was lost. For terminal Cancelled, also re-fan-out -// dependents so a crash between the terminal CAS and the dependent -// publish does not strand them. +// - Terminal → re-publish the dependent fan-out and the conclude +// event. Every terminal transition is routed back through this controller +// (by mergesignal, or by the cancel flow), so this branch is how waiting +// dependents learn a dependency resolved; redelivery makes the same +// branch the self-heal for a lost publish. +// +// Cancel decisions are recorded as path status only +// (SpeculationPathStatusCancelling on an in-flight build) — nothing in this +// file calls a build runner. // // The controller is re-triggered on every relevant downstream event // (buildsignal, merge), so each call simply re-evaluates the current @@ -57,6 +75,10 @@ type Controller struct { logger *zap.SugaredLogger metricsScope tally.Scope store storage.Storage + enumerators enumerator.Factory + scorers pathscorer.Factory + selectors selector.Factory + depLimits dependencylimit.Factory registry consumer.TopicRegistry topicKey consumer.TopicKey consumerGroup string @@ -73,6 +95,10 @@ func NewController( logger *zap.SugaredLogger, scope tally.Scope, store storage.Storage, + enumerators enumerator.Factory, + scorers pathscorer.Factory, + selectors selector.Factory, + depLimits dependencylimit.Factory, registry consumer.TopicRegistry, topicKey consumer.TopicKey, consumerGroup string, @@ -81,13 +107,18 @@ func NewController( logger: logger.Named("speculate_controller"), metricsScope: scope.SubScope("speculate_controller"), store: store, + enumerators: enumerators, + scorers: scorers, + selectors: selectors, + depLimits: depLimits, registry: registry, topicKey: topicKey, consumerGroup: consumerGroup, } } -// Process advances a batch one step along the naive happy-path. +// Process reconciles the batch's speculation tree and advances the batch one +// step in the state machine. // Returns nil to ack (success), or error to nack (retry). func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (retErr error) { op := metrics.Begin(c.metricsScope, opName) @@ -114,17 +145,17 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return c.cancelBatch(ctx, batch) } - // Terminal state: re-fan-out for self-healing in case a previous publish - // was lost. Always re-publish to conclude (idempotent on the batch ID). - // For Cancelled specifically also re-publish to dependents — a crash - // between the terminal CAS and the dependent publish would otherwise - // leave them stuck waiting on a Cancelled dep. + // Terminal state: wake dependents and re-publish conclude. This branch is + // the dependent wake-up for success and failure — mergesignal routes every + // terminal transition back through speculate under the batch's own ID, and + // re-publishing the dependents here lets each of them re-run its own + // dependency gate / tryFinalize against the new dependency state. On + // redelivery the same branch doubles as the crash self-heal: both the + // dependent fan-out and the conclude publish are idempotent re-sends. if batch.State.IsTerminal() { metrics.NamedCounter(c.metricsScope, opName, "self_heal_terminal", 1) - if batch.State == entity.BatchStateCancelled { - if err := c.respeculateDependents(ctx, batch); err != nil { - return err - } + if err := c.respeculateDependents(ctx, batch); err != nil { + return err } return c.fanout(ctx, batch.ID, batch.Queue) } @@ -135,19 +166,344 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } + switch batch.State { + case entity.BatchStateCreated, entity.BatchStateScored, entity.BatchStateSpeculating: + return c.speculateBatch(ctx, batch) + default: + metrics.NamedCounter(c.metricsScope, opName, "unexpected_state", 1) + return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) + } +} + +// speculateBatch is the unified entry point for Created, Scored, and +// Speculating batches. It loads the batch's speculation tree (creating it on +// the first pass the dependency gate admits), applies the scorer's and +// selector's outputs, persists the tree if anything changed, and then +// advances the batch itself: Created/Scored publish to build and CAS to +// Speculating; Speculating runs tryFinalize. +func (c *Controller) speculateBatch(ctx context.Context, batch entity.Batch) error { + deps, err := c.fetchDependencies(ctx, batch) + if err != nil { + return err + } + + tree, err := c.store.GetSpeculationTreeStore().Get(ctx, batch.ID) + if err != nil { + if !errors.Is(err, storage.ErrNotFound) { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to get speculation tree for batch %s: %w", batch.ID, err) + } + + // Dependency gate: only consulted before a tree exists. Once a batch + // has a tree, later passes just re-score/re-select it. + blocked, gerr := c.dependencyGateBlocks(ctx, batch, deps) + if gerr != nil { + return gerr + } + if blocked { + metrics.NamedCounter(c.metricsScope, opName, "dependency_gate_blocked", 1) + c.logger.Debugw("active dependency count exceeds queue's dependency limit; waiting", + "batch_id", batch.ID, + "dependency_count", len(deps), + ) + return nil + } + + tree, err = c.createTree(ctx, batch, deps) + if err != nil { + return err + } + } + + scr, err := c.scorers.For(pathscorer.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) + return fmt.Errorf("failed to get speculation scorer for queue %s: %w", batch.Queue, err) + } + scores, err := scr.Score(ctx, tree) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) + return fmt.Errorf("failed to score speculation tree for batch %s: %w", batch.ID, err) + } + tree, scoresChanged := c.applyScores(batch, tree, scores) + + sel, err := c.selectors.For(selector.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "selector_errors", 1) + return fmt.Errorf("failed to get selector for queue %s: %w", batch.Queue, err) + } + decisions, err := sel.Select(ctx, tree) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "selector_errors", 1) + return fmt.Errorf("failed to select speculation decisions for batch %s: %w", batch.ID, err) + } + + tree, selectionChanged := c.applySelection(batch, tree, decisions) + + if scoresChanged || selectionChanged { + newVersion := tree.Version + 1 + if err := c.store.GetSpeculationTreeStore().Update(ctx, batch.ID, 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", batch.ID, err) + } + tree.Version = newVersion + } + + // The tree does not influence the forward step below — no downstream + // stage consumes it yet. switch batch.State { case entity.BatchStateCreated, entity.BatchStateScored: return c.startSpeculation(ctx, batch) case entity.BatchStateSpeculating: return c.tryFinalize(ctx, batch) default: - metrics.NamedCounter(c.metricsScope, opName, "unexpected_state", 1) - return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) + return nil + } +} + +// dependencyGateBlocks reports whether batch's count of active dependencies +// (entity.DependencyBatchStates) exceeds the queue's current dependency +// limit. It is consulted only before a batch's speculation tree exists. +// +// The gate applies the dependencylimit extension's value; the application — +// counting active dependencies and expressing "wait" as an acked no-op — is +// admission control on the batch's pipeline progress and deliberately stays +// in the controller, out of the structure-only enumerator seam. A blocked +// batch is woken by the next dependency event: every dependency terminal +// transition re-publishes the dependents of the newly terminal batch (see +// the terminal branch in Process), and the active count only shrinks via +// those same transitions, so no unblocking event can be missed; a raised +// limit takes effect at the next such event. This cannot deadlock: +// dependencies point at strictly earlier batches (the graph is a DAG) and a +// batch with no active dependencies is never blocked, so the head of every +// chain keeps progressing and eventually wakes its dependents. +func (c *Controller) dependencyGateBlocks(ctx context.Context, batch entity.Batch, deps []entity.Batch) (bool, error) { + active := 0 + for _, d := range deps { + if isActiveDependency(d.State) { + active++ + } + } + + limiter, err := c.depLimits.For(dependencylimit.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "dependency_limit_errors", 1) + return false, fmt.Errorf("failed to get dependency limit for queue %s: %w", batch.Queue, err) + } + limit, err := limiter.Limit(ctx) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "dependency_limit_errors", 1) + return false, fmt.Errorf("failed to get dependency limit value for queue %s: %w", batch.Queue, err) + } + + return active > limit, nil +} + +// isActiveDependency reports whether s is one of the states that makes an +// in-flight batch eligible to be a dependency (entity.DependencyBatchStates). +func isActiveDependency(s entity.BatchState) bool { + for _, st := range entity.DependencyBatchStates() { + if st == s { + return true + } + } + return false +} + +// createTree enumerates and persists a batch's speculation tree the first +// time it is seen, with every path stamped Candidate. Concurrent creation +// (two events racing to create the same tree) is resolved by re-reading the +// winner's tree rather than erroring: enumeration is deterministic given the +// same (batchID, deps), so either creator's structure is equivalent and the +// loser simply adopts what won. +func (c *Controller) createTree(ctx context.Context, batch entity.Batch, deps []entity.Batch) (entity.SpeculationTree, error) { + enumFactory, err := c.enumerators.For(enumerator.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "enumerator_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to get enumerator for queue %s: %w", batch.Queue, err) + } + + paths, err := enumFactory.Enumerate(ctx, batch, deps) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "enumerator_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to enumerate speculation paths for batch %s: %w", batch.ID, err) + } + + // The enumerator returns structure only; the controller owns everything + // else about a persisted path. Each entry is stamped Candidate and minted + // its ID here, once, at tree creation — immutable thereafter: it is how + // scores, decisions, and the path->build mapping (PathBuild.PathID) refer + // to the path. A structural duplicate from the enumerator is a contract + // violation; the first occurrence wins and the rest are skipped. + infos := make([]entity.SpeculationPathInfo, 0, len(paths)) + for _, p := range paths { + dup := false + for _, existing := range infos { + if existing.Path.Equal(p) { + dup = true + break + } + } + if dup { + metrics.NamedCounter(c.metricsScope, opName, "duplicate_enumerated_path", 1) + c.logger.Warnw("enumerator returned duplicate path; skipped", + "batch_id", batch.ID, + "path", p, + ) + continue + } + infos = append(infos, entity.SpeculationPathInfo{ + ID: fmt.Sprintf("%s/path/%d", batch.ID, len(infos)), + Path: p, + Status: entity.SpeculationPathStatusCandidate, + }) + } + tree := entity.SpeculationTree{BatchID: batch.ID, Paths: infos, Version: 1} + + if err := c.store.GetSpeculationTreeStore().Create(ctx, tree); err != nil { + if errors.Is(err, storage.ErrAlreadyExists) { + metrics.NamedCounter(c.metricsScope, opName, "tree_create_race_lost", 1) + existing, gerr := c.store.GetSpeculationTreeStore().Get(ctx, batch.ID) + if gerr != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to re-get speculation tree for batch %s after concurrent create: %w", batch.ID, gerr) + } + return existing, nil + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to create speculation tree for batch %s: %w", batch.ID, err) + } + + metrics.NamedCounter(c.metricsScope, opName, "tree_created", 1) + return tree, nil +} + +// applyScores merges the scorer's path-ID-keyed scores into tree, enforcing +// the seam contract on consume: a score naming a path not in the tree is +// logged and skipped, and a score outside [0, 1] is clamped into range with a +// warning — a scorer bug must not corrupt ranking inputs downstream. Paths +// the scorer omitted keep their last persisted score. The returned bool +// reports whether any persisted score actually changed, so the caller can +// skip the conditional write on a no-op pass. +func (c *Controller) applyScores(batch entity.Batch, tree entity.SpeculationTree, scores []entity.PathScore) (entity.SpeculationTree, bool) { + changed := false + for _, ps := range scores { + idx := tree.PathIndex(ps.PathID) + if idx == -1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_score", 1) + c.logger.Warnw("illegal score: path not found in tree", + "batch_id", batch.ID, + "path_id", ps.PathID, + ) + continue + } + score := ps.Score + if score < 0 || score > 1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_score", 1) + c.logger.Warnw("illegal score: outside [0, 1]; clamped", + "batch_id", batch.ID, + "path_id", ps.PathID, + "score", score, + ) + if score < 0 { + score = 0 + } else { + score = 1 + } + } + if tree.Paths[idx].Score != score { + tree.Paths[idx].Score = score + changed = true + } } + return tree, changed } -// startSpeculation kicks off CI for this batch on top of the speculative head -// (batch.Dependencies assumed to all pass), then transitions to Speculating. +// applySelection applies the selector's decisions to tree.Paths, mutating it +// in place. The selector owns the policy (which paths to promote or cancel); +// this function owns only the bookkeeping, mapping each decision onto the +// path's current status. Mirrors prioritize.go's apply loop: Promote clears a +// Candidate to Selected; Cancel is captured as intent only — a Building path +// moves to Cancelling, and any pre-build status (Candidate, Selected, or +// Prioritized) drops straight to Cancelled. Nothing here touches a build +// runner — a Cancelling status records the intent for whichever stage owns +// runner interaction to enact. A decision naming a path not in +// the tree, a duplicate decision for the same path, or an action the path's +// current status does not support is a policy bug in the selector: it is +// logged as a warning and skipped rather than corrupting the tree. The +// returned bool reports whether any path status changed, so the caller can +// skip the conditional write on a no-op pass. +func (c *Controller) applySelection(batch entity.Batch, tree entity.SpeculationTree, decisions []entity.PathDecision) (entity.SpeculationTree, bool) { + if len(decisions) == 0 { + return tree, false + } + + paths := tree.Paths + changed := false + 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", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + seen[d.PathID] = true + + 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", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + + info := paths[idx] + switch { + case d.Action == entity.SpeculationPathActionPromote && info.Status == entity.SpeculationPathStatusCandidate: + info.Status = entity.SpeculationPathStatusSelected + changed = true + metrics.NamedCounter(c.metricsScope, opName, "selected", 1) + + case d.Action == entity.SpeculationPathActionCancel && info.Status == entity.SpeculationPathStatusBuilding: + info.Status = entity.SpeculationPathStatusCancelling + changed = true + metrics.NamedCounter(c.metricsScope, opName, "cancelling", 1) + + case d.Action == entity.SpeculationPathActionCancel && + (info.Status == entity.SpeculationPathStatusCandidate || + info.Status == entity.SpeculationPathStatusSelected || + info.Status == entity.SpeculationPathStatusPrioritized): + info.Status = entity.SpeculationPathStatusCancelled + changed = true + 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", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + "status", string(info.Status), + ) + continue + } + + paths[idx] = info + } + + tree.Paths = paths + return tree, changed +} + +// startSpeculation publishes the batch to the build stage, then transitions +// it to Speculating. func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) error { c.logger.Infow("starting speculation", "batch_id", batch.ID, @@ -176,7 +532,8 @@ func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) e // out-of-the-way: the cancelled batch will never land, so it can no longer // conflict — drop it from the chain and proceed. Failed deps still cascade // via failOnDependency. If some deps are still in flight, the call is a -// no-op and waits for the next event. +// no-op: each dependency's own terminal pass re-publishes this batch to +// speculate (the terminal branch in Process), so waiting never needs a poll. // // TODO: when a dependency fails we currently fail this batch outright. // We will need to respeculate the failed paths — drop the failed dep @@ -235,7 +592,10 @@ func (c *Controller) tryFinalize(ctx context.Context, batch entity.Batch) error // dependencies has reached a non-succeeding terminal state, then publishes to // the conclude queue so the request store and request log get reconciled. // Without this transition the batch would sit in Speculating forever — no -// downstream event ever fires for it again. +// downstream event ever fires for it again. The batch's own dependents are +// woken right after the terminal CAS so the failure cascades downstream +// instead of stranding them mid-wait; a crash between the CAS and the +// fan-out is recovered by the terminal branch in Process on redelivery. func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, dep entity.Batch) error { metrics.NamedCounter(c.metricsScope, opName, "dependency_failed", 1) c.logger.Warnw("dependency in non-succeeding terminal state; failing batch", @@ -250,6 +610,10 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d return fmt.Errorf("failed to update batch %s state to failed: %w", batch.ID, err) } + if err := c.respeculateDependents(ctx, batch); err != nil { + return err + } + if err := c.publish(ctx, topickey.TopicKeyConclude, batch.ID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to conclude: %w", err) @@ -268,9 +632,8 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d // Order matters for correctness: // // 1. Cancel the in-flight Build entity (build.ID == batch.ID; one Get + one -// UpdateStatus covers all builds for this batch). A future external CI -// integration hooks in here. Idempotent: tolerate ErrNotFound (no build -// was scheduled), skip if already terminal. +// UpdateStatus covers all builds for this batch). Idempotent: tolerate +// ErrNotFound (no build was scheduled), skip if already terminal. // // 2. CAS the batch to terminal Cancelled. This must happen BEFORE the // dependent fan-out: tryFinalize only drops a Cancelled dep from the @@ -331,11 +694,6 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error // covers every build scheduled for the batch. Tolerates ErrNotFound (no // build was ever scheduled — the batch was cancelled before speculation // started building) and skips already-terminal builds. -// -// This is the hook point for a future external CI integration: today the -// system has no external runner, so the local state flip is the complete -// cancellation. Once a runner exists, it must be invoked here before the -// local UpdateStatus. func (c *Controller) cancelBuild(ctx context.Context, batch entity.Batch) error { build, err := c.store.GetBuildStore().Get(ctx, batch.ID) if err != nil { @@ -368,8 +726,10 @@ func (c *Controller) cancelBuild(ctx context.Context, batch entity.Batch) error // the message nacks and either an operator or the batch controller's own // crash-recovery can resolve the inconsistency. // -// Called both from the cancelBatch terminal flow and from the terminal -// self-heal branch on redelivery of an already-Cancelled batch. +// Called from the cancelBatch terminal flow, from failOnDependency, and from +// the terminal branch in Process — the latter both on the first pass after +// mergesignal drives a batch terminal (the normal dependent wake-up) and on +// redelivery (self-heal). func (c *Controller) respeculateDependents(ctx context.Context, batch entity.Batch) error { bd, err := c.store.GetBatchDependentStore().Get(ctx, batch.ID) if err != nil { diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 7eff01b5..d9425bec 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate_test.go +++ b/submitqueue/orchestrator/controller/speculate/speculate_test.go @@ -28,12 +28,26 @@ import ( queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" + dependencylimitfake "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit/fake" + dependencylimitmock "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit/mock" + enumeratorfake "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator/fake" + enumeratormock "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator/mock" + scorerfake "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer/fake" + scorermock "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer/mock" + selectorfake "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/fake" + selectormock "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/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" ) +// pubRec records a single publish call for order/content assertions. +type pubRec struct { + topic string + msgID string +} + // batchIDPayload serializes a BatchID to JSON bytes for test message payloads. func batchIDPayload(t *testing.T, id string) []byte { payload, err := entity.BatchID{ID: id}.ToBytes() @@ -52,33 +66,121 @@ func testBatch(state entity.BatchState, deps ...string) entity.Batch { } } -// newTestController wires a controller with a registry covering all topics the -// speculate controller may publish to. The publisher returns publishErr (or nil). -func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock.MockStorage, publishErr error) *Controller { - logger := zaptest.NewLogger(t).Sugar() - scope := tally.NoopScope +// testHarness wires a Controller against mocked storage plus programmable +// fake seam implementations (enumerator, scorer, selector, dependency limit) +// that tests script per case. Every publish is recorded so tests can assert +// both content and order. There is no build-runner seam yet: nothing in this +// package's speculate.go calls one — cancel decisions on the speculation tree +// are captured as intent only. +type testHarness struct { + controller *Controller + batchStore *storagemock.MockBatchStore + treeStore *storagemock.MockSpeculationTreeStore + buildStore *storagemock.MockBuildStore + depStore *storagemock.MockBatchDependentStore + enum *enumeratorfake.Enumerator + scorer *scorerfake.Scorer + selector *selectorfake.Selector + depLimit *dependencylimitfake.DependencyLimit + records *[]pubRec +} + +// newTestHarness builds a harness whose seams default to parity behavior: +// the fake enumerator produces an empty tree unless seeded, the fake scorer +// echoes its input, the fake selector decides nothing, and the fake +// dependency limit is effectively ungated (1000). Every publish succeeds and +// is appended to records. +func newTestHarness(t *testing.T, ctrl *gomock.Controller) *testHarness { + return newHarness(t, ctrl, nil, 1000) +} + +// newFailingPublishHarness is identical to newTestHarness except every +// publish returns publishErr instead of succeeding. +func newFailingPublishHarness(t *testing.T, ctrl *gomock.Controller, publishErr error) *testHarness { + return newHarness(t, ctrl, publishErr, 1000) +} + +// newDependencyLimitHarness is identical to newTestHarness except the +// dependency limit is set to limit instead of the default ungated value. +func newDependencyLimitHarness(t *testing.T, ctrl *gomock.Controller, limit int) *testHarness { + return newHarness(t, ctrl, nil, limit) +} + +func newHarness(t *testing.T, ctrl *gomock.Controller, publishErr error, depLimit int) *testHarness { + batchStore := storagemock.NewMockBatchStore(ctrl) + treeStore := storagemock.NewMockSpeculationTreeStore(ctrl) + buildStore := storagemock.NewMockBuildStore(ctrl) + depStore := storagemock.NewMockBatchDependentStore(ctrl) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetSpeculationTreeStore().Return(treeStore).AnyTimes() + store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() + + enum := enumeratorfake.New() + enumFactory := enumeratormock.NewMockFactory(ctrl) + enumFactory.EXPECT().For(gomock.Any()).Return(enum, nil).AnyTimes() + + scr := scorerfake.New() + scorerFactory := scorermock.NewMockFactory(ctrl) + scorerFactory.EXPECT().For(gomock.Any()).Return(scr, nil).AnyTimes() - mockPub := queuemock.NewMockPublisher(ctrl) - mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( - func(ctx context.Context, topic string, msg entityqueue.Message) error { - return publishErr - }, - ).AnyTimes() + sel := selectorfake.New() + selectorFactory := selectormock.NewMockFactory(ctrl) + selectorFactory.EXPECT().For(gomock.Any()).Return(sel, nil).AnyTimes() + + dl := dependencylimitfake.New(depLimit) + depLimitFactory := dependencylimitmock.NewMockFactory(ctrl) + depLimitFactory.EXPECT().For(gomock.Any()).Return(dl, nil).AnyTimes() + + var records []pubRec + pub := queuemock.NewMockPublisher(ctrl) + pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, topic string, msg entityqueue.Message) error { + if publishErr != nil { + return publishErr + } + records = append(records, pubRec{topic: topic, msgID: msg.ID}) + return nil + }).AnyTimes() mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() - - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyBuild, Name: "build", Queue: mockQ}, - {Key: topickey.TopicKeyMerge, Name: "merge", Queue: mockQ}, - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - {Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ}, - }, - ) + mockQ.EXPECT().Publisher().Return(pub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, + {Key: topickey.TopicKeyBuild, Name: "build", Queue: mockQ}, + {Key: topickey.TopicKeyMerge, Name: "merge", Queue: mockQ}, + {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, + }) require.NoError(t, err) - return NewController(logger, scope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + c := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + store, + enumFactory, + scorerFactory, + selectorFactory, + depLimitFactory, + registry, + topickey.TopicKeySpeculate, + "orchestrator-speculate", + ) + + return &testHarness{ + controller: c, + batchStore: batchStore, + treeStore: treeStore, + buildStore: buildStore, + depStore: depStore, + enum: enum, + scorer: scr, + selector: sel, + depLimit: dl, + records: &records, + } } // runProcess builds a delivery for batchID and invokes Process once. @@ -92,323 +194,469 @@ func runProcess(t *testing.T, ctrl *gomock.Controller, controller *Controller, b func TestNewController(t *testing.T) { ctrl := gomock.NewController(t) - store := storagemock.NewMockStorage(ctrl) - controller := newTestController(t, ctrl, store, nil) + h := newTestHarness(t, ctrl) - require.NotNil(t, controller) - assert.Equal(t, topickey.TopicKeySpeculate, controller.TopicKey()) - assert.Equal(t, "orchestrator-speculate", controller.ConsumerGroup()) - assert.Equal(t, "speculate", controller.Name()) + require.NotNil(t, h.controller) + assert.Equal(t, topickey.TopicKeySpeculate, h.controller.TopicKey()) + assert.Equal(t, "orchestrator-speculate", h.controller.ConsumerGroup()) + assert.Equal(t, "speculate", h.controller.Name()) - var _ consumer.Controller = controller + var _ consumer.Controller = h.controller } -// startSpeculation: Created/Scored should publish to build and CAS to Speculating with newVersion = oldVersion+1. -func TestController_Process_StartSpeculation(t *testing.T) { - tests := []struct { - name string - state entity.BatchState - }{ - {name: "from_created", state: entity.BatchStateCreated}, - {name: "from_scored", state: entity.BatchStateScored}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctrl := gomock.NewController(t) - batch := testBatch(tt.state) - - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateSpeculating).Return(nil) +func TestController_Process_BadPayload(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + msg := entityqueue.NewMessage("anything", []byte("not-json"), "test-queue", nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) - }) - } + require.Error(t, h.controller.Process(context.Background(), delivery)) } -// tryFinalize: Speculating with no deps should publish to merge and CAS to Merging. -func TestController_Process_FinalizeNoDeps(t *testing.T) { +func TestController_Process_StorageFailure(t *testing.T) { ctrl := gomock.NewController(t) - batch := testBatch(entity.BatchStateSpeculating) + h := newTestHarness(t, ctrl) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.Batch{}, fmt.Errorf("db connection lost")) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + err := runProcess(t, ctrl, h.controller, "test-queue/batch/1") + require.Error(t, err) + assert.False(t, errs.IsRetryable(err)) +} + +func TestController_Process_UnrecognizedState(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateUnknown) + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.Error(t, runProcess(t, ctrl, h.controller, batch.ID)) } -// tryFinalize: Speculating with all deps Succeeded should publish to merge and CAS to Merging. -func TestController_Process_FinalizeAllDepsSucceeded(t *testing.T) { +// Merging is owned by the merge controller — speculate is a no-op for it. +func TestController_Process_MergingNoOp(t *testing.T) { ctrl := gomock.NewController(t) - depA := entity.Batch{ID: "test-queue/batch/0a", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 5} - depB := entity.Batch{ID: "test-queue/batch/0b", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 3} - batch := testBatch(entity.BatchStateSpeculating, depA.ID, depB.ID) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateMerging) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().Get(gomock.Any(), depA.ID).Return(depA, nil) - batchStore.EXPECT().Get(gomock.Any(), depB.ID).Return(depB, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // No UpdateState, no tree access, no publish expected. - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) +} - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) +// Terminal states wake dependents (re-publish speculate for each) and then +// re-publish conclude — this is both the normal dependent wake-up (mergesignal +// routes every terminal transition of a batch back through speculate under +// the batch's own ID) and the crash self-heal for a lost publish. State must +// not change (no UpdateState), and neither BuildStore nor +// SpeculationTreeStore is touched. +func TestController_Process_TerminalSelfHeals(t *testing.T) { + for _, state := range []entity.BatchState{ + entity.BatchStateSucceeded, + entity.BatchStateFailed, + entity.BatchStateCancelled, + } { + t.Run(string(state), func(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(state) + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // No UpdateState, no tree access expected. + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{"test-queue/batch/2", "test-queue/batch/3"}, + Version: 1, + }, nil) + + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{ + {topic: "speculate", msgID: "test-queue/batch/2"}, + {topic: "speculate", msgID: "test-queue/batch/3"}, + {topic: "conclude", msgID: batch.ID}, + }, *h.records) + }) + } } -// tryFinalize: Speculating with a dep still in flight is a no-op (no publish, no state change). -func TestController_Process_WaitingOnDep(t *testing.T) { +// An empty dependents list publishes nothing extra beyond conclude. The +// BatchDependent row itself must still exist for the batch — a missing row +// is a storage invariant violation surfaced as an error, not a normal +// "no dependents" case. +func TestController_Process_TerminalSelfHealNoDependents(t *testing.T) { ctrl := gomock.NewController(t) - dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateSpeculating, Version: 1} - batch := testBatch(entity.BatchStateSpeculating, dep.ID) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateSucceeded) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil) - // No UpdateState expected — gomock will fail if it is called. + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + BatchID: batch.ID, + Dependents: nil, + Version: 1, + }, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "conclude", msgID: batch.ID}}, *h.records) +} + +// A missing BatchDependent row is a storage invariant violation (the batch +// controller creates the row before the batch itself), surfaced as an error +// so the message nacks — unlike a stale dependent entry, which is tolerated. +func TestController_Process_TerminalMissingDependentRowErrors(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateSucceeded) - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{}, storage.ErrNotFound) + + require.Error(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) } -// tryFinalize: a failed dep must fail the batch (Speculating → Failed) and -// publish to conclude. Otherwise the batch livelocks. -func TestController_Process_FailedDepFailsBatch(t *testing.T) { +// Created batch: no tree yet -> enumerate, stamp Candidate, mint path IDs, +// Version 1, Create. Score echoes. Select promotes the lone path to Selected +// -> tree changed -> Update(1,2,...). The forward step then runs +// unchanged: batch CASes Created -> Speculating and publishes to build. +func TestController_Process_CreatedEnumeratesScoresSelects(t *testing.T) { ctrl := gomock.NewController(t) - dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateFailed, Version: 1} - batch := testBatch(entity.BatchStateSpeculating, dep.ID) - batch.Contains = []string{"test-queue/req/1", "test-queue/req/2"} + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateCreated) + path := entity.SpeculationPath{Head: batch.ID} + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{}, storage.ErrNotFound) + + h.enum.Set(batch.ID, []entity.SpeculationPath{path}) + h.treeStore.EXPECT().Create(gomock.Any(), gomock.AssignableToTypeOf(entity.SpeculationTree{})). + DoAndReturn(func(_ context.Context, tree entity.SpeculationTree) error { + require.Len(t, tree.Paths, 1) + assert.Equal(t, entity.SpeculationPathStatusCandidate, tree.Paths[0].Status) + assert.Equal(t, fmt.Sprintf("%s/path/0", batch.ID), tree.Paths[0].ID) + assert.Equal(t, int32(1), tree.Version) + return nil + }) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateFailed).Return(nil) + h.selector.SetDecisions(entity.PathDecision{PathID: fmt.Sprintf("%s/path/0", batch.ID), Action: entity.SpeculationPathActionPromote}) + h.treeStore.EXPECT().Update(gomock.Any(), batch.ID, int32(1), int32(2), gomock.Any()). + DoAndReturn(func(_ context.Context, _ string, _, _ int32, paths []entity.SpeculationPathInfo) error { + require.Len(t, paths, 1) + assert.Equal(t, entity.SpeculationPathStatusSelected, paths[0].Status) + return nil + }) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateSpeculating).Return(nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "build", msgID: batch.ID}}, *h.records) +} + +// TestController_Process_CreateTreeMintsPathIDs pins down the exact minted +// path ID format: fmt.Sprintf("%s/path/%d", batch.ID, i), by enumeration +// index. +func TestController_Process_CreateTreeMintsPathIDs(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateCreated) + pathA := entity.SpeculationPath{Head: batch.ID} + pathB := entity.SpeculationPath{Base: []string{"test-queue/batch/0"}, Head: batch.ID} + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{}, storage.ErrNotFound) + + h.enum.Set(batch.ID, []entity.SpeculationPath{pathA, pathB}) + h.treeStore.EXPECT().Create(gomock.Any(), gomock.AssignableToTypeOf(entity.SpeculationTree{})). + DoAndReturn(func(_ context.Context, tree entity.SpeculationTree) error { + require.Len(t, tree.Paths, 2) + for i, p := range tree.Paths { + assert.Equal(t, fmt.Sprintf("%s/path/%d", batch.ID, i), p.ID) + } + return nil + }) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateSpeculating).Return(nil) - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "build", msgID: batch.ID}}, *h.records) } -// tryFinalize: a cancelled dep is treated as out-of-the-way — it will never -// land and can no longer conflict. The dep is dropped from the chain and the -// batch advances to Merging as if the cancelled dep had succeeded. -func TestController_Process_CancelledDepSkipped(t *testing.T) { +// Create racing with a concurrent creator: ErrAlreadyExists must fall back to +// re-reading the winner's tree rather than erroring. The forward step +// still runs afterward. +func TestController_Process_CreateTreeAlreadyExistsRereads(t *testing.T) { ctrl := gomock.NewController(t) - depCancelled := entity.Batch{ID: "test-queue/batch/0a", Queue: "test-queue", State: entity.BatchStateCancelled, Version: 2} - depSucceeded := entity.Batch{ID: "test-queue/batch/0b", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 5} - batch := testBatch(entity.BatchStateSpeculating, depCancelled.ID, depSucceeded.ID) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateCreated) + path := entity.SpeculationPath{Head: batch.ID} - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().Get(gomock.Any(), depCancelled.ID).Return(depCancelled, nil) - batchStore.EXPECT().Get(gomock.Any(), depSucceeded.ID).Return(depSucceeded, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + existing := entity.SpeculationTree{ + BatchID: batch.ID, + Version: 1, + Paths: []entity.SpeculationPathInfo{{ID: "test-queue/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusCandidate}}, + } + gomock.InOrder( + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{}, storage.ErrNotFound), + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(existing, nil), + ) + h.enum.Set(batch.ID, []entity.SpeculationPath{path}) + h.treeStore.EXPECT().Create(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, tree entity.SpeculationTree) error { + require.Len(t, tree.Paths, 1) + assert.NotEmpty(t, tree.Paths[0].ID, "minted path ID must not be empty") + return storage.ErrAlreadyExists + }) + // No changes this pass (selector decides nothing) -> no Update call. + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateSpeculating).Return(nil) - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "build", msgID: batch.ID}}, *h.records) } -// Merging is owned by the merge controller — speculate is a no-op for it. -func TestController_Process_MergingNoOp(t *testing.T) { +// Dependency gate: too many active dependencies for a batch with no tree yet +// blocks enumeration entirely (ack without creating a tree, no forward +// step either). A later dependency-terminal event re-triggers speculate. +func TestController_Process_DependencyGateBlocksTreeCreation(t *testing.T) { ctrl := gomock.NewController(t) - batch := testBatch(entity.BatchStateMerging) + h := newDependencyLimitHarness(t, ctrl, 1) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - // No UpdateState expected. + depA := entity.Batch{ID: "test-queue/batch/0a", Queue: "test-queue", State: entity.BatchStateCreated, Version: 1} + depB := entity.Batch{ID: "test-queue/batch/0b", Queue: "test-queue", State: entity.BatchStateScored, Version: 1} + batch := testBatch(entity.BatchStateCreated, depA.ID, depB.ID) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().Get(gomock.Any(), depA.ID).Return(depA, nil) + h.batchStore.EXPECT().Get(gomock.Any(), depB.ID).Return(depB, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{}, storage.ErrNotFound) + // No Create, no Enumerate call, no UpdateState expected — the mocks fail + // the test if called. - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) } -// Terminal states re-fan-out to conclude for self-healing in case a previous -// publish was lost. State must not change (no UpdateState). The Cancelled -// terminal also re-fans-out dependents and is covered separately in -// TestController_Process_CancelledTerminalSelfHealsDependents. -func TestController_Process_TerminalSelfHeals(t *testing.T) { - for _, state := range []entity.BatchState{ - entity.BatchStateSucceeded, - entity.BatchStateFailed, - } { - t.Run(string(state), func(t *testing.T) { - ctrl := gomock.NewController(t) - batch := testBatch(state) +// A pass over an unchanged tree must not call Update. The forward +// step still runs: here a still-pending dependency makes tryFinalize wait, +// so no publish happens at all. +func TestController_Process_NoChangeSkipsTreeUpdate(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateSpeculating, Version: 1} + batch := testBatch(entity.BatchStateSpeculating, dep.ID) + path := entity.SpeculationPath{Head: batch.ID} + tree := entity.SpeculationTree{ + BatchID: batch.ID, + Version: 4, + Paths: []entity.SpeculationPathInfo{{ID: "test-queue/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusSelected}}, + } - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - // No UpdateState expected. + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil).AnyTimes() + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(tree, nil) + // Fake scorer echoes, fake selector decides nothing -> no change. + // No treeStore.Update expected. tryFinalize waits on the pending dep, so + // no UpdateState/publish either. - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) +} - // Require exactly one publish to the conclude topic for self-healing. - mockPub := queuemock.NewMockPublisher(ctrl) - mockPub.EXPECT().Publish(gomock.Any(), "conclude", gomock.Any()).Return(nil).Times(1) +// A version mismatch on the speculation tree Update must surface as an error +// (nack; the round is recomputed on redelivery) and must not reach the +// forward step. +func TestController_Process_TreeUpdateVersionMismatchErrors(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateSpeculating) + path := entity.SpeculationPath{Head: batch.ID} + tree := entity.SpeculationTree{ + BatchID: batch.ID, + Version: 1, + Paths: []entity.SpeculationPathInfo{{ID: "test-queue/batch/1/path/0", Path: path, Status: entity.SpeculationPathStatusCandidate}}, + } - mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(tree, nil) + h.selector.SetDecisions(entity.PathDecision{PathID: tree.Paths[0].ID, Action: entity.SpeculationPathActionPromote}) + h.treeStore.EXPECT().Update(gomock.Any(), batch.ID, int32(1), int32(2), gomock.Any()).Return(storage.ErrVersionMismatch) + // tryFinalize must never run: no batchStore.UpdateState, no merge publish. - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - }, - ) - require.NoError(t, err) + err := runProcess(t, ctrl, h.controller, batch.ID) + require.Error(t, err) + assert.ErrorIs(t, err, storage.ErrVersionMismatch) + assert.Empty(t, *h.records) +} - logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") +// tryFinalize: Speculating with no deps should publish to merge and CAS to +// Merging. The batch already has an (empty) speculation tree from its +// Created/Scored pass. +func TestController_Process_FinalizeNoDeps(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + batch := testBatch(entity.BatchStateSpeculating) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) - }) - } + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) + + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "merge", msgID: batch.ID}}, *h.records) } -// Cancelled is terminal: redelivery must re-fan-out dependents (so a crash -// between the terminal CAS and the dependent publish does not strand them) -// AND re-publish to conclude. State must not change (no UpdateState; no -// build cancel). The BuildStore must not be touched on this self-heal path. -func TestController_Process_CancelledTerminalSelfHealsDependents(t *testing.T) { +// tryFinalize: Speculating with all deps Succeeded should publish to merge +// and CAS to Merging. +func TestController_Process_FinalizeAllDepsSucceeded(t *testing.T) { ctrl := gomock.NewController(t) - batch := testBatch(entity.BatchStateCancelled) + h := newTestHarness(t, ctrl) + depA := entity.Batch{ID: "test-queue/batch/0a", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 5} + depB := entity.Batch{ID: "test-queue/batch/0b", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 3} + batch := testBatch(entity.BatchStateSpeculating, depA.ID, depB.ID) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - // No UpdateState expected. + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // speculateBatch fetches dependencies once for the pre-tree gate check, + // and tryFinalize fetches them again itself. + h.batchStore.EXPECT().Get(gomock.Any(), depA.ID).Return(depA, nil).Times(2) + h.batchStore.EXPECT().Get(gomock.Any(), depB.ID).Return(depB, nil).Times(2) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) - depStore := storagemock.NewMockBatchDependentStore(ctrl) - depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "merge", msgID: batch.ID}}, *h.records) +} + +// tryFinalize: Speculating with a dep still in flight is a no-op (no +// publish, no state change). +func TestController_Process_WaitingOnDep(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateSpeculating, Version: 1} + batch := testBatch(entity.BatchStateSpeculating, dep.ID) + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // speculateBatch fetches dependencies once for the pre-tree gate check, + // and tryFinalize fetches them again itself. + h.batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil).Times(2) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + // No UpdateState expected — gomock will fail if it is called. + + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) +} + +// tryFinalize: a failed dep must fail the batch (Speculating → Failed), wake +// its dependents (so the failure cascades), and publish to conclude. +// Otherwise the batch livelocks. +func TestController_Process_FailedDepFailsBatch(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateFailed, Version: 1} + batch := testBatch(entity.BatchStateSpeculating, dep.ID) + batch.Contains = []string{"test-queue/req/1", "test-queue/req/2"} + + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // speculateBatch fetches dependencies once for the pre-tree gate check, + // and tryFinalize fetches them again itself. + h.batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil).Times(2) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateFailed).Return(nil) + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ BatchID: batch.ID, - Dependents: []string{"test-queue/batch/2", "test-queue/batch/3"}, + Dependents: []string{"test-queue/batch/2"}, Version: 1, }, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() - // BuildStore must NOT be touched on the terminal self-heal path. + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{ + {topic: "speculate", msgID: "test-queue/batch/2"}, + {topic: "conclude", msgID: batch.ID}, + }, *h.records) +} - type pubRec struct { - topic string - msgID string - } - var records []pubRec - mockPub := queuemock.NewMockPublisher(ctrl) - mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, topic string, msg entityqueue.Message) error { - records = append(records, pubRec{topic: topic, msgID: msg.ID}) - return nil - }).AnyTimes() +// A dependent-wake publish failure after the terminal CAS nacks the message; +// redelivery converges through the terminal branch, which re-runs the +// fan-out and the conclude publish. +func TestController_Process_FailedDepDependentPublishFailure(t *testing.T) { + ctrl := gomock.NewController(t) + h := newFailingPublishHarness(t, ctrl, fmt.Errorf("publish failed")) + dep := entity.Batch{ID: "test-queue/batch/0", Queue: "test-queue", State: entity.BatchStateFailed, Version: 1} + batch := testBatch(entity.BatchStateSpeculating, dep.ID) - mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // speculateBatch fetches dependencies once for the pre-tree gate check, + // and tryFinalize fetches them again itself. + h.batchStore.EXPECT().Get(gomock.Any(), dep.ID).Return(dep, nil).Times(2) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateFailed).Return(nil) + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{"test-queue/batch/2"}, + Version: 1, + }, nil) - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, - }, - ) - require.NoError(t, err) + require.Error(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Empty(t, *h.records) +} - logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") +// tryFinalize: a cancelled dep is treated as out-of-the-way — it will never +// land and can no longer conflict. The dep is dropped from the chain and the +// batch advances to Merging as if the cancelled dep had succeeded. +func TestController_Process_CancelledDepSkipped(t *testing.T) { + ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) + depCancelled := entity.Batch{ID: "test-queue/batch/0a", Queue: "test-queue", State: entity.BatchStateCancelled, Version: 2} + depSucceeded := entity.Batch{ID: "test-queue/batch/0b", Queue: "test-queue", State: entity.BatchStateSucceeded, Version: 5} + batch := testBatch(entity.BatchStateSpeculating, depCancelled.ID, depSucceeded.ID) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + // speculateBatch fetches dependencies once for the pre-tree gate check, + // and tryFinalize fetches them again itself. + h.batchStore.EXPECT().Get(gomock.Any(), depCancelled.ID).Return(depCancelled, nil).Times(2) + h.batchStore.EXPECT().Get(gomock.Any(), depSucceeded.ID).Return(depSucceeded, nil).Times(2) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateMerging).Return(nil) - assert.Equal(t, []pubRec{ - {topic: "speculate", msgID: "test-queue/batch/2"}, - {topic: "speculate", msgID: "test-queue/batch/3"}, - {topic: "conclude", msgID: batch.ID}, - }, records) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "merge", msgID: batch.ID}}, *h.records) } // Cancelling drives the terminal-cancellation flow: cancel any in-flight // build, CAS the batch to Cancelled, fan out dependents, publish to -// conclude. Validates the full happy-path order with a running build and +// conclude. Validates the full order with a running build and // a couple of dependents. Order matters: dependents must publish AFTER the // terminal CAS so the woken dependents observe the dep as Cancelled (and // drop it from their chain) rather than as still-Cancelling (which would // leave them waiting on a state nobody is going to nudge). func TestController_Process_CancellingTerminalFlow(t *testing.T) { ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) batch := testBatch(entity.BatchStateCancelling) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) - buildStore := storagemock.NewMockBuildStore(ctrl) - buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{ + h.buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{ ID: batch.ID, BatchID: batch.ID, Status: entity.BuildStatusRunning, }, nil) - buildStore.EXPECT().UpdateStatus(gomock.Any(), batch.ID, entity.BuildStatusCancelled).Return(nil) + h.buildStore.EXPECT().UpdateStatus(gomock.Any(), batch.ID, entity.BuildStatusCancelled).Return(nil) - depStore := storagemock.NewMockBatchDependentStore(ctrl) - depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ BatchID: batch.ID, Dependents: []string{"test-queue/batch/2", "test-queue/batch/3"}, Version: 1, }, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() - store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() - - type pubRec struct { - topic string - msgID string - } - var records []pubRec - mockPub := queuemock.NewMockPublisher(ctrl) - mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, topic string, msg entityqueue.Message) error { - records = append(records, pubRec{topic: topic, msgID: msg.ID}) - return nil - }).AnyTimes() - - mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() - - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, - }, - ) - require.NoError(t, err) - - logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") - - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) assert.Equal(t, []pubRec{ {topic: "speculate", msgID: "test-queue/batch/2"}, {topic: "speculate", msgID: "test-queue/batch/3"}, {topic: "conclude", msgID: batch.ID}, - }, records) + }, *h.records) } // If the build for the batch has already reached a terminal status (e.g. CI @@ -417,30 +665,22 @@ func TestController_Process_CancellingTerminalFlow(t *testing.T) { // of the flow (terminal batch CAS, dependent fan-out, conclude) still runs. func TestController_Process_CancellingBuildAlreadyTerminal(t *testing.T) { ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) batch := testBatch(entity.BatchStateCancelling) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) - buildStore := storagemock.NewMockBuildStore(ctrl) - buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{ + h.buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{ ID: batch.ID, BatchID: batch.ID, Status: entity.BuildStatusSucceeded, }, nil) // No UpdateStatus expected — the build is already terminal. - depStore := storagemock.NewMockBatchDependentStore(ctrl) - depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ BatchID: batch.ID, Version: 1, }, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() - store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() - - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) } // If no Build entity exists for the batch (e.g. cancel arrived before @@ -448,28 +688,20 @@ func TestController_Process_CancellingBuildAlreadyTerminal(t *testing.T) { // tolerated and the rest of the cancellation flow must continue. func TestController_Process_CancellingNoBuildYet(t *testing.T) { ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) batch := testBatch(entity.BatchStateCancelling) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) - buildStore := storagemock.NewMockBuildStore(ctrl) - buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) + h.buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) // No UpdateStatus expected. - depStore := storagemock.NewMockBatchDependentStore(ctrl) - depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{ BatchID: batch.ID, Version: 1, }, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() - store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() - - controller := newTestController(t, ctrl, store, nil) - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) } // A batch whose BatchDependent row exists with an empty Dependents list must @@ -478,40 +710,17 @@ func TestController_Process_CancellingNoBuildYet(t *testing.T) { // list at batch creation time and it stays empty if no later batch conflicts. func TestController_Process_CancellingNoDependents(t *testing.T) { ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) batch := testBatch(entity.BatchStateCancelling) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) - - buildStore := storagemock.NewMockBuildStore(ctrl) - buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled).Return(nil) - depStore := storagemock.NewMockBatchDependentStore(ctrl) - depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{BatchID: batch.ID, Dependents: []string{}, Version: 1}, nil) + h.buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) + h.depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{BatchID: batch.ID, Dependents: []string{}, Version: 1}, nil) - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() - store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() - - mockPub := queuemock.NewMockPublisher(ctrl) - mockPub.EXPECT().Publish(gomock.Any(), "conclude", gomock.Any()).Return(nil).Times(1) - - mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() - - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - }, - ) - require.NoError(t, err) - - logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") - - require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) + require.NoError(t, runProcess(t, ctrl, h.controller, batch.ID)) + assert.Equal(t, []pubRec{{topic: "conclude", msgID: batch.ID}}, *h.records) } // storage.ErrVersionMismatch on the terminal CAS must surface as an error @@ -521,101 +730,31 @@ func TestController_Process_CancellingNoDependents(t *testing.T) { // will pick up the (now-terminal) state and complete the fan-out. func TestController_Process_CancellingTerminalCASVersionMismatch(t *testing.T) { ctrl := gomock.NewController(t) + h := newTestHarness(t, ctrl) batch := testBatch(entity.BatchStateCancelling) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled). + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelled). Return(storage.ErrVersionMismatch) - buildStore := storagemock.NewMockBuildStore(ctrl) - buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) - - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() + h.buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) // BatchDependentStore must NOT be touched — terminal CAS failed before fan-out. - // No publish expected (terminal CAS failed before fan-out). - mockPub := queuemock.NewMockPublisher(ctrl) - mockQ := queuemock.NewMockQueue(ctrl) - mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() - - registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{ - {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: mockQ}, - {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, - }, - ) - require.NoError(t, err) - - logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") - - err = runProcess(t, ctrl, controller, batch.ID) + err := runProcess(t, ctrl, h.controller, batch.ID) require.Error(t, err) assert.ErrorIs(t, err, storage.ErrVersionMismatch) -} - -// An unrecognized state must surface as an error so the message is nacked -// instead of silently acked — silently acking would drop the event. -func TestController_Process_UnrecognizedState(t *testing.T) { - ctrl := gomock.NewController(t) - batch := testBatch(entity.BatchStateUnknown) - - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) - - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - - controller := newTestController(t, ctrl, store, nil) - require.Error(t, runProcess(t, ctrl, controller, batch.ID)) -} - -// Storage failure on the primary batch fetch surfaces as an error and is not -// retryable per the controller default (plain fmt.Errorf). -func TestController_Process_StorageFailure(t *testing.T) { - ctrl := gomock.NewController(t) - - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.Batch{}, fmt.Errorf("db connection lost")) - - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - - controller := newTestController(t, ctrl, store, nil) - err := runProcess(t, ctrl, controller, "test-queue/batch/1") - require.Error(t, err) - assert.False(t, errs.IsRetryable(err)) + assert.Empty(t, *h.records) } // Publish failure must not advance the batch state. func TestController_Process_PublishFailure(t *testing.T) { ctrl := gomock.NewController(t) + h := newFailingPublishHarness(t, ctrl, fmt.Errorf("publish failed")) batch := testBatch(entity.BatchStateScored) - batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + h.treeStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.SpeculationTree{BatchID: batch.ID, Version: 1, Paths: []entity.SpeculationPathInfo{}}, nil) // No UpdateState expected — publish fails before we get there. - store := storagemock.NewMockStorage(ctrl) - store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() - - controller := newTestController(t, ctrl, store, fmt.Errorf("publish failed")) - require.Error(t, runProcess(t, ctrl, controller, batch.ID)) -} - -// Malformed payload: deserialize error. -func TestController_Process_BadPayload(t *testing.T) { - ctrl := gomock.NewController(t) - store := storagemock.NewMockStorage(ctrl) - controller := newTestController(t, ctrl, store, nil) - - msg := entityqueue.NewMessage("anything", []byte("not-json"), "test-queue", nil) - delivery := queuemock.NewMockDelivery(ctrl) - delivery.EXPECT().Message().Return(msg).AnyTimes() - delivery.EXPECT().Attempt().Return(1).AnyTimes() - - require.Error(t, controller.Process(context.Background(), delivery)) + require.Error(t, runProcess(t, ctrl, h.controller, batch.ID)) }