diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 3bd852c8..022f999c 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -41,6 +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", @@ -54,6 +65,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..8a322d4f 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -61,6 +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" @@ -74,6 +85,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" @@ -234,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) } @@ -244,13 +256,18 @@ func run() error { brf := buildRunnerFactory{queues} 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, 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 } - dlqCount, err := registerDLQControllers(dlqConsumer, logger.Sugar(), scope, store) + dlqCount, err := registerDLQControllers(dlqConsumer, logger.Sugar(), scope, registry, store) if err != nil { return err } @@ -380,6 +397,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 +489,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 @@ -489,10 +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 + 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 @@ -538,7 +568,37 @@ 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 +} + +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, @@ -628,6 +688,10 @@ func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, logger, scope, store, + enf, + ssf, + slf, + dlf, registry, topickey.TopicKeySpeculate, "orchestrator-speculate", @@ -637,6 +701,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 +790,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 +803,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 +938,16 @@ 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 + +// 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 @@ -866,7 +955,7 @@ func newPhabChangeProvider(logger *zap.Logger, scope tally.Scope) (changeprovide // 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) @@ -880,16 +969,25 @@ 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. + // + // 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), @@ -900,7 +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), + 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/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) +} 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)) }