From 527af81bc6a0e9054beeecc31ec8391fc3cd6bf7 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 10 Jul 2026 07:58:36 -0700 Subject: [PATCH] feat(speculate): write the speculation tree in shadow mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This is the first of four slices reproducing the tree-driven speculate/build/buildsignal rework as an incremental, always-shippable stack instead of one large commit. Each slice keeps the orchestrator working end to end so the e2e suite stays green throughout, and reviewers can follow the pipeline shift (direct publish -> tree -> prioritize -> build) one deliberate step at a time. This slice only adds the speculation tree as bookkeeping: nothing downstream reads it yet, so the batch's forward step (when it builds, when it merges) is untouched — with one deliberate exception called out below: a liveness fix to the dependent wake-up that review of this slice surfaced. speculate's Controller now takes an enumerator, path scorer, selector, and dependency-limit factory, and every Created/Scored/Speculating pass loads or creates the batch's entity.SpeculationTree, applies the scorer's and selector's outputs, and persists it only if something changed (the apply steps report whether they mutated anything, gating a version+1 conditional Update). The controller stays pure mechanics: which paths exist, how they score, and which are promoted or cancelled are the seams' decisions alone — the controller validates and records them. Tree creation is gated by the queue's dependency limit; the controller wraps the enumerator's structure-only paths into persisted entries itself — stamping Candidate, minting each path's immutable ID as `{batchID}/path/{i}`, and skipping structural duplicates as a contract violation — and a concurrent create race re-reads the winner's tree instead of erroring. Seam outputs are consumed by path ID: scores merged by ID (unknown IDs skipped, out-of-[0,1] values clamped, both logged), selector decisions resolved through the ID-keyed PathIndex with duplicates logged and skipped. The pre-existing forward step is otherwise unchanged: Created/Scored batches CAS to Speculating and publish straight to build, and Speculating batches run the original tryFinalize (merge once every dependency has landed, cascade-fail on a failed dependency). Cancelling is the unmodified pre-existing flow. Liveness fix in the forward step: a batch waiting on dependencies (in tryFinalize, or now in the dependency gate) was only ever woken when a dependency was cancelled — a dependency reaching Succeeded or Failed never re-published its dependents, because mergesignal routes every terminal transition back through speculate under the batch's own ID and the terminal branch only re-published conclude. The terminal branch now fans out to dependents for every terminal state, and failOnDependency wakes its own dependents right after the terminal CAS so failures cascade downstream. The fan-out publishes a wake for every listed dependent; a dangling reverse-index entry (left by an abandoned batch creation) surfaces as a not-found speculate message that dead-letters and is skipped by the DLQ reconciler — eliminating that class at the source, by creating the batch before the reverse-index update, is tracked in #354. Cancel decisions on tree paths are recorded as status only (Cancelling on an in-flight build) — nothing in this controller calls a build runner, which is why it takes no buildrunner.Factory. main.go wires the new seams with parity defaults: chain enumerates a single path per batch, the probability path scorer, the all-selector promotes every candidate, and a static dependency limit that is effectively ungated. `bazel test //submitqueue/... //service/...` — 56/56 targets pass. New unit coverage for the wake-up fix: dependent fan-out on every terminal state (with publish-order assertions), missing BatchDependent row surfaced as an invariant error, and dependent-publish failure after the terminal CAS nacking for redelivery convergence. `make gazelle && make fmt` — no diffs beyond the intended BUILD file updates. Part of the speculation rework. Interim fan-out caveats tracked in #354. --- .../orchestrator/server/BUILD.bazel | 8 + .../submitqueue/orchestrator/server/main.go | 81 +- .../controller/speculate/BUILD.bazel | 12 + .../controller/speculate/speculate.go | 436 ++++++++- .../controller/speculate/speculate_test.go | 909 ++++++++++-------- 5 files changed, 1012 insertions(+), 434 deletions(-) diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index 2ccb3493..022f999c 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -41,9 +41,17 @@ go_library( "//submitqueue/extension/scorer/composite:go_default_library", "//submitqueue/extension/scorer/fake:go_default_library", "//submitqueue/extension/scorer/heuristic:go_default_library", + "//submitqueue/extension/speculation/dependencylimit:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/static:go_default_library", + "//submitqueue/extension/speculation/enumerator:go_default_library", + "//submitqueue/extension/speculation/enumerator/chain:go_default_library", + "//submitqueue/extension/speculation/pathscorer:go_default_library", + "//submitqueue/extension/speculation/pathscorer/probability:go_default_library", "//submitqueue/extension/speculation/prioritizationlimit/static:go_default_library", "//submitqueue/extension/speculation/prioritizer:go_default_library", "//submitqueue/extension/speculation/prioritizer/sticky:go_default_library", + "//submitqueue/extension/speculation/selector:go_default_library", + "//submitqueue/extension/speculation/selector/all:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mysql:go_default_library", "//submitqueue/extension/validator/fake:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 6ed707ca..b22fefb9 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -61,9 +61,17 @@ import ( "github.com/uber/submitqueue/submitqueue/extension/scorer/composite" scorerfake "github.com/uber/submitqueue/submitqueue/extension/scorer/fake" "github.com/uber/submitqueue/submitqueue/extension/scorer/heuristic" + "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit" + dependencylimitstatic "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit/static" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator/chain" + "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer" + pathscorerprobability "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer/probability" prioritizationlimitstatic "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizationlimit/static" "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer" "github.com/uber/submitqueue/submitqueue/extension/speculation/prioritizer/sticky" + "github.com/uber/submitqueue/submitqueue/extension/speculation/selector" + selectorall "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/all" "github.com/uber/submitqueue/submitqueue/extension/storage" mysqlstorage "github.com/uber/submitqueue/submitqueue/extension/storage/mysql" validatorfake "github.com/uber/submitqueue/submitqueue/extension/validator/fake" @@ -238,7 +246,7 @@ func run() error { // back to a baseline profile for queues without an explicit entry. This is // the single place queue topology is known; the extension packages stay // queue-agnostic. - queues, err := newQueueRegistry(logger, scope, changeset.New(store.GetRequestStore(), store.GetChangeStore())) + queues, err := newQueueRegistry(logger, scope, changeset.New(store.GetRequestStore(), store.GetChangeStore()), store.GetBatchStore()) if err != nil { return fmt.Errorf("failed to build queue registry: %w", err) } @@ -249,9 +257,13 @@ func run() error { scf := scorerFactory{queues} cof := analyzerFactory{queues} prf := prioritizerFactory{queues} + enf := enumeratorFactory{queues} + ssf := speculationScorerFactory{queues} + slf := speculationSelectorFactory{queues} + dlf := dependencyLimitFactory{queues} // Register controllers - primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, prf, cnt, store) + primaryCount, err := registerPrimaryControllers(primaryConsumer, logger.Sugar(), scope, registry, cpf, brf, scf, cof, prf, enf, ssf, slf, dlf, cnt, store) if err != nil { return err } @@ -502,11 +514,15 @@ func newTopicRegistry(q extqueue.Queue, subscriberName string) (consumer.TopicRe // read as "for this queue, here are its scorer, analyzer, change provider, …", and lets // a queue profile start from a baseline and override only what differs. type queueExtensions struct { - changeProvider changeprovider.ChangeProvider - buildRunner buildrunner.BuildRunner - scorer scorer.Scorer - analyzer conflict.Analyzer - prioritizer prioritizer.Prioritizer + changeProvider changeprovider.ChangeProvider + buildRunner buildrunner.BuildRunner + scorer scorer.Scorer + analyzer conflict.Analyzer + prioritizer prioritizer.Prioritizer + enumerator enumerator.Enumerator + speculationScorer pathscorer.Scorer + speculationSelector selector.Selector + dependencyLimit dependencylimit.DependencyLimit } // queueRegistry maps a queue name to its extensions, falling back to a default @@ -558,7 +574,31 @@ func (f prioritizerFactory) For(cfg prioritizer.Config) (prioritizer.Prioritizer return f.reg.get(cfg.QueueName).prioritizer, nil } -func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, registry consumer.TopicRegistry, cpf changeprovider.Factory, brf buildrunner.Factory, scf scorer.Factory, cof conflict.Factory, prf prioritizer.Factory, cnt counter.Counter, store storage.Storage) (int, error) { +type enumeratorFactory struct{ reg queueRegistry } + +func (f enumeratorFactory) For(cfg enumerator.Config) (enumerator.Enumerator, error) { + return f.reg.get(cfg.QueueName).enumerator, nil +} + +type speculationScorerFactory struct{ reg queueRegistry } + +func (f speculationScorerFactory) For(cfg pathscorer.Config) (pathscorer.Scorer, error) { + return f.reg.get(cfg.QueueName).speculationScorer, nil +} + +type speculationSelectorFactory struct{ reg queueRegistry } + +func (f speculationSelectorFactory) For(cfg selector.Config) (selector.Selector, error) { + return f.reg.get(cfg.QueueName).speculationSelector, nil +} + +type dependencyLimitFactory struct{ reg queueRegistry } + +func (f dependencyLimitFactory) For(cfg dependencylimit.Config) (dependencylimit.DependencyLimit, error) { + return f.reg.get(cfg.QueueName).dependencyLimit, nil +} + +func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, scope tally.Scope, registry consumer.TopicRegistry, cpf changeprovider.Factory, brf buildrunner.Factory, scf scorer.Factory, cof conflict.Factory, prf prioritizer.Factory, enf enumerator.Factory, ssf pathscorer.Factory, slf selector.Factory, dlf dependencylimit.Factory, cnt counter.Counter, store storage.Storage) (int, error) { var count int requestController := start.NewController( logger, @@ -648,6 +688,10 @@ func registerPrimaryControllers(c consumer.Consumer, logger *zap.SugaredLogger, logger, scope, store, + enf, + ssf, + slf, + dlf, registry, topickey.TopicKeySpeculate, "orchestrator-speculate", @@ -899,6 +943,11 @@ func newPhabChangeProvider(logger *zap.Logger, scope tally.Scope) (changeprovide // effectively admit-all — until per-queue budgets are configured. const defaultPrioritizationLimit = 1000 +// defaultDependencyLimit is the baseline cap on active dependencies a batch +// may speculate over. It is a parity default — effectively ungated — until +// per-queue limits are configured. +const defaultDependencyLimit = 1000 + // newQueueRegistry builds the per-queue extension profiles for the example. // Edge integrations (change provider) and the build // runner form a shared baseline; each per-queue profile starts from that @@ -906,7 +955,7 @@ const defaultPrioritizationLimit = 1000 // conflict analyzer. Queues without an explicit profile fall back to the // baseline. This is the one place queue topology lives; extension packages stay // queue-agnostic. -func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset.Resolver) (queueRegistry, error) { +func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset.Resolver, batchStore storage.BatchStore) (queueRegistry, error) { cp, err := newChangeProvider(logger, scope) if err != nil { return queueRegistry{}, fmt.Errorf("failed to create change provider: %w", err) @@ -933,6 +982,12 @@ func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset. // below does. The prioritizer is sticky over a static budget: it never // preempts a running build and admits Selected candidates by score until // defaultPrioritizationLimit concurrent builds are in flight. + // + // The speculation seams default to the single-chain parity policies: + // chain enumerates one path per batch (built on the full ordered + // dependency chain), probability scores paths from the batches' + // predicted-success probabilities and resolved outcomes, all promotes + // every candidate, and the dependency limit is effectively ungated. base := queueExtensions{ changeProvider: cp, buildRunner: buildfake.New(resolver), @@ -943,8 +998,12 @@ func newQueueRegistry(logger *zap.Logger, scope tally.Scope, resolver changeset. )), // TODO: replace the delegate with a real analyzer (e.g. Tango target // analysis). "all" serializes the queue conservatively. - analyzer: conflictfake.New(all.New(), nil), - prioritizer: sticky.New(prioritizationlimitstatic.New(defaultPrioritizationLimit)), + analyzer: conflictfake.New(all.New(), nil), + prioritizer: sticky.New(prioritizationlimitstatic.New(defaultPrioritizationLimit)), + enumerator: chain.New(), + speculationScorer: pathscorerprobability.New(batchStore), + speculationSelector: selectorall.New(), + dependencyLimit: dependencylimitstatic.New(defaultDependencyLimit), } // test-queue: bucketed heuristic scorer; conservative (serialized) conflicts diff --git a/submitqueue/orchestrator/controller/speculate/BUILD.bazel b/submitqueue/orchestrator/controller/speculate/BUILD.bazel index 39e36977..b6f312da 100644 --- a/submitqueue/orchestrator/controller/speculate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/speculate/BUILD.bazel @@ -11,6 +11,10 @@ go_library( "//platform/metrics:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/dependencylimit:go_default_library", + "//submitqueue/extension/speculation/enumerator:go_default_library", + "//submitqueue/extension/speculation/pathscorer:go_default_library", + "//submitqueue/extension/speculation/selector:go_default_library", "//submitqueue/extension/storage:go_default_library", "@com_github_uber_go_tally//:go_default_library", "@org_uber_go_zap//:go_default_library", @@ -28,6 +32,14 @@ go_test( "//platform/extension/messagequeue/mock:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/fake:go_default_library", + "//submitqueue/extension/speculation/dependencylimit/mock:go_default_library", + "//submitqueue/extension/speculation/enumerator/fake:go_default_library", + "//submitqueue/extension/speculation/enumerator/mock:go_default_library", + "//submitqueue/extension/speculation/pathscorer/fake:go_default_library", + "//submitqueue/extension/speculation/pathscorer/mock:go_default_library", + "//submitqueue/extension/speculation/selector/fake:go_default_library", + "//submitqueue/extension/speculation/selector/mock:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mock:go_default_library", "@com_github_stretchr_testify//assert:go_default_library", diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index f2d2b21b..50380d2d 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -25,30 +25,48 @@ import ( "github.com/uber/submitqueue/platform/metrics" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/dependencylimit" + "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator" + "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/selector" "github.com/uber/submitqueue/submitqueue/extension/storage" "go.uber.org/zap" ) // Controller handles speculate queue messages. // -// Naive happy-path algorithm: assume every in-flight build will pass and -// treat batch.Dependencies + [batch.ID] as the single speculation chain. -// Per invocation, the controller advances the batch one step in the -// state machine: +// Each invocation reconciles the batch's entity.SpeculationTree and advances +// the batch one step in the state machine. Tree reconciliation is pure +// mechanics: the controller runs the queue's configured seams — enumerator +// (path structure), path scorer (path scores), selector (path decisions) — +// validates their outputs, applies them to the persisted tree, and writes it +// back under optimistic concurrency. Which paths exist, how they score, and +// which are promoted or cancelled are the seams' decisions alone; the +// controller never originates one. No downstream stage reads the tree yet, +// so the forward step below is driven by the batch's own state, not the +// tree. // -// - Created or Scored → publish to build, transition to Speculating. -// - Speculating → if all deps are Succeeded, publish to merge and -// transition to Merging; otherwise no-op (or fail-fast if a dep is -// in a non-succeeding terminal state). +// Per invocation, the controller advances the batch one step in the state +// machine: +// +// - Created, Scored, or Speculating → speculateBatch: reconcile the tree +// (created on the first pass the dependency gate admits), then advance +// the batch — publish to build and CAS to Speculating for +// Created/Scored, or tryFinalize for Speculating. // - Cancelling → cancel any in-flight Build entity, respeculate // dependents, CAS to terminal Cancelled, publish to conclude. The // cancel controller hands the batch off in this state and speculate // drives it to terminal. // - Merging → no-op (owned by the merge controller). -// - Terminal → re-fan-out to conclude for self-healing in case a -// prior publish was lost. For terminal Cancelled, also re-fan-out -// dependents so a crash between the terminal CAS and the dependent -// publish does not strand them. +// - Terminal → re-publish the dependent fan-out and the conclude +// event. Every terminal transition is routed back through this controller +// (by mergesignal, or by the cancel flow), so this branch is how waiting +// dependents learn a dependency resolved; redelivery makes the same +// branch the self-heal for a lost publish. +// +// Cancel decisions are recorded as path status only +// (SpeculationPathStatusCancelling on an in-flight build) — nothing in this +// file calls a build runner. // // The controller is re-triggered on every relevant downstream event // (buildsignal, merge), so each call simply re-evaluates the current @@ -57,6 +75,10 @@ type Controller struct { logger *zap.SugaredLogger metricsScope tally.Scope store storage.Storage + enumerators enumerator.Factory + scorers pathscorer.Factory + selectors selector.Factory + depLimits dependencylimit.Factory registry consumer.TopicRegistry topicKey consumer.TopicKey consumerGroup string @@ -73,6 +95,10 @@ func NewController( logger *zap.SugaredLogger, scope tally.Scope, store storage.Storage, + enumerators enumerator.Factory, + scorers pathscorer.Factory, + selectors selector.Factory, + depLimits dependencylimit.Factory, registry consumer.TopicRegistry, topicKey consumer.TopicKey, consumerGroup string, @@ -81,13 +107,18 @@ func NewController( logger: logger.Named("speculate_controller"), metricsScope: scope.SubScope("speculate_controller"), store: store, + enumerators: enumerators, + scorers: scorers, + selectors: selectors, + depLimits: depLimits, registry: registry, topicKey: topicKey, consumerGroup: consumerGroup, } } -// Process advances a batch one step along the naive happy-path. +// Process reconciles the batch's speculation tree and advances the batch one +// step in the state machine. // Returns nil to ack (success), or error to nack (retry). func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (retErr error) { op := metrics.Begin(c.metricsScope, opName) @@ -114,17 +145,17 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return c.cancelBatch(ctx, batch) } - // Terminal state: re-fan-out for self-healing in case a previous publish - // was lost. Always re-publish to conclude (idempotent on the batch ID). - // For Cancelled specifically also re-publish to dependents — a crash - // between the terminal CAS and the dependent publish would otherwise - // leave them stuck waiting on a Cancelled dep. + // Terminal state: wake dependents and re-publish conclude. This branch is + // the dependent wake-up for success and failure — mergesignal routes every + // terminal transition back through speculate under the batch's own ID, and + // re-publishing the dependents here lets each of them re-run its own + // dependency gate / tryFinalize against the new dependency state. On + // redelivery the same branch doubles as the crash self-heal: both the + // dependent fan-out and the conclude publish are idempotent re-sends. if batch.State.IsTerminal() { metrics.NamedCounter(c.metricsScope, opName, "self_heal_terminal", 1) - if batch.State == entity.BatchStateCancelled { - if err := c.respeculateDependents(ctx, batch); err != nil { - return err - } + if err := c.respeculateDependents(ctx, batch); err != nil { + return err } return c.fanout(ctx, batch.ID, batch.Queue) } @@ -135,19 +166,344 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } + switch batch.State { + case entity.BatchStateCreated, entity.BatchStateScored, entity.BatchStateSpeculating: + return c.speculateBatch(ctx, batch) + default: + metrics.NamedCounter(c.metricsScope, opName, "unexpected_state", 1) + return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) + } +} + +// speculateBatch is the unified entry point for Created, Scored, and +// Speculating batches. It loads the batch's speculation tree (creating it on +// the first pass the dependency gate admits), applies the scorer's and +// selector's outputs, persists the tree if anything changed, and then +// advances the batch itself: Created/Scored publish to build and CAS to +// Speculating; Speculating runs tryFinalize. +func (c *Controller) speculateBatch(ctx context.Context, batch entity.Batch) error { + deps, err := c.fetchDependencies(ctx, batch) + if err != nil { + return err + } + + tree, err := c.store.GetSpeculationTreeStore().Get(ctx, batch.ID) + if err != nil { + if !errors.Is(err, storage.ErrNotFound) { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to get speculation tree for batch %s: %w", batch.ID, err) + } + + // Dependency gate: only consulted before a tree exists. Once a batch + // has a tree, later passes just re-score/re-select it. + blocked, gerr := c.dependencyGateBlocks(ctx, batch, deps) + if gerr != nil { + return gerr + } + if blocked { + metrics.NamedCounter(c.metricsScope, opName, "dependency_gate_blocked", 1) + c.logger.Debugw("active dependency count exceeds queue's dependency limit; waiting", + "batch_id", batch.ID, + "dependency_count", len(deps), + ) + return nil + } + + tree, err = c.createTree(ctx, batch, deps) + if err != nil { + return err + } + } + + scr, err := c.scorers.For(pathscorer.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) + return fmt.Errorf("failed to get speculation scorer for queue %s: %w", batch.Queue, err) + } + scores, err := scr.Score(ctx, tree) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) + return fmt.Errorf("failed to score speculation tree for batch %s: %w", batch.ID, err) + } + tree, scoresChanged := c.applyScores(batch, tree, scores) + + sel, err := c.selectors.For(selector.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "selector_errors", 1) + return fmt.Errorf("failed to get selector for queue %s: %w", batch.Queue, err) + } + decisions, err := sel.Select(ctx, tree) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "selector_errors", 1) + return fmt.Errorf("failed to select speculation decisions for batch %s: %w", batch.ID, err) + } + + tree, selectionChanged := c.applySelection(batch, tree, decisions) + + if scoresChanged || selectionChanged { + newVersion := tree.Version + 1 + if err := c.store.GetSpeculationTreeStore().Update(ctx, batch.ID, tree.Version, newVersion, tree.Paths); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to update speculation tree for batch %s: %w", batch.ID, err) + } + tree.Version = newVersion + } + + // The tree does not influence the forward step below — no downstream + // stage consumes it yet. switch batch.State { case entity.BatchStateCreated, entity.BatchStateScored: return c.startSpeculation(ctx, batch) case entity.BatchStateSpeculating: return c.tryFinalize(ctx, batch) default: - metrics.NamedCounter(c.metricsScope, opName, "unexpected_state", 1) - return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) + return nil + } +} + +// dependencyGateBlocks reports whether batch's count of active dependencies +// (entity.DependencyBatchStates) exceeds the queue's current dependency +// limit. It is consulted only before a batch's speculation tree exists. +// +// The gate applies the dependencylimit extension's value; the application — +// counting active dependencies and expressing "wait" as an acked no-op — is +// admission control on the batch's pipeline progress and deliberately stays +// in the controller, out of the structure-only enumerator seam. A blocked +// batch is woken by the next dependency event: every dependency terminal +// transition re-publishes the dependents of the newly terminal batch (see +// the terminal branch in Process), and the active count only shrinks via +// those same transitions, so no unblocking event can be missed; a raised +// limit takes effect at the next such event. This cannot deadlock: +// dependencies point at strictly earlier batches (the graph is a DAG) and a +// batch with no active dependencies is never blocked, so the head of every +// chain keeps progressing and eventually wakes its dependents. +func (c *Controller) dependencyGateBlocks(ctx context.Context, batch entity.Batch, deps []entity.Batch) (bool, error) { + active := 0 + for _, d := range deps { + if isActiveDependency(d.State) { + active++ + } + } + + limiter, err := c.depLimits.For(dependencylimit.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "dependency_limit_errors", 1) + return false, fmt.Errorf("failed to get dependency limit for queue %s: %w", batch.Queue, err) + } + limit, err := limiter.Limit(ctx) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "dependency_limit_errors", 1) + return false, fmt.Errorf("failed to get dependency limit value for queue %s: %w", batch.Queue, err) + } + + return active > limit, nil +} + +// isActiveDependency reports whether s is one of the states that makes an +// in-flight batch eligible to be a dependency (entity.DependencyBatchStates). +func isActiveDependency(s entity.BatchState) bool { + for _, st := range entity.DependencyBatchStates() { + if st == s { + return true + } + } + return false +} + +// createTree enumerates and persists a batch's speculation tree the first +// time it is seen, with every path stamped Candidate. Concurrent creation +// (two events racing to create the same tree) is resolved by re-reading the +// winner's tree rather than erroring: enumeration is deterministic given the +// same (batchID, deps), so either creator's structure is equivalent and the +// loser simply adopts what won. +func (c *Controller) createTree(ctx context.Context, batch entity.Batch, deps []entity.Batch) (entity.SpeculationTree, error) { + enumFactory, err := c.enumerators.For(enumerator.Config{QueueName: batch.Queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "enumerator_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to get enumerator for queue %s: %w", batch.Queue, err) + } + + paths, err := enumFactory.Enumerate(ctx, batch, deps) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "enumerator_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to enumerate speculation paths for batch %s: %w", batch.ID, err) + } + + // The enumerator returns structure only; the controller owns everything + // else about a persisted path. Each entry is stamped Candidate and minted + // its ID here, once, at tree creation — immutable thereafter: it is how + // scores, decisions, and the path->build mapping (PathBuild.PathID) refer + // to the path. A structural duplicate from the enumerator is a contract + // violation; the first occurrence wins and the rest are skipped. + infos := make([]entity.SpeculationPathInfo, 0, len(paths)) + for _, p := range paths { + dup := false + for _, existing := range infos { + if existing.Path.Equal(p) { + dup = true + break + } + } + if dup { + metrics.NamedCounter(c.metricsScope, opName, "duplicate_enumerated_path", 1) + c.logger.Warnw("enumerator returned duplicate path; skipped", + "batch_id", batch.ID, + "path", p, + ) + continue + } + infos = append(infos, entity.SpeculationPathInfo{ + ID: fmt.Sprintf("%s/path/%d", batch.ID, len(infos)), + Path: p, + Status: entity.SpeculationPathStatusCandidate, + }) + } + tree := entity.SpeculationTree{BatchID: batch.ID, Paths: infos, Version: 1} + + if err := c.store.GetSpeculationTreeStore().Create(ctx, tree); err != nil { + if errors.Is(err, storage.ErrAlreadyExists) { + metrics.NamedCounter(c.metricsScope, opName, "tree_create_race_lost", 1) + existing, gerr := c.store.GetSpeculationTreeStore().Get(ctx, batch.ID) + if gerr != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to re-get speculation tree for batch %s after concurrent create: %w", batch.ID, gerr) + } + return existing, nil + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return entity.SpeculationTree{}, fmt.Errorf("failed to create speculation tree for batch %s: %w", batch.ID, err) + } + + metrics.NamedCounter(c.metricsScope, opName, "tree_created", 1) + return tree, nil +} + +// applyScores merges the scorer's path-ID-keyed scores into tree, enforcing +// the seam contract on consume: a score naming a path not in the tree is +// logged and skipped, and a score outside [0, 1] is clamped into range with a +// warning — a scorer bug must not corrupt ranking inputs downstream. Paths +// the scorer omitted keep their last persisted score. The returned bool +// reports whether any persisted score actually changed, so the caller can +// skip the conditional write on a no-op pass. +func (c *Controller) applyScores(batch entity.Batch, tree entity.SpeculationTree, scores []entity.PathScore) (entity.SpeculationTree, bool) { + changed := false + for _, ps := range scores { + idx := tree.PathIndex(ps.PathID) + if idx == -1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_score", 1) + c.logger.Warnw("illegal score: path not found in tree", + "batch_id", batch.ID, + "path_id", ps.PathID, + ) + continue + } + score := ps.Score + if score < 0 || score > 1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_score", 1) + c.logger.Warnw("illegal score: outside [0, 1]; clamped", + "batch_id", batch.ID, + "path_id", ps.PathID, + "score", score, + ) + if score < 0 { + score = 0 + } else { + score = 1 + } + } + if tree.Paths[idx].Score != score { + tree.Paths[idx].Score = score + changed = true + } } + return tree, changed } -// startSpeculation kicks off CI for this batch on top of the speculative head -// (batch.Dependencies assumed to all pass), then transitions to Speculating. +// applySelection applies the selector's decisions to tree.Paths, mutating it +// in place. The selector owns the policy (which paths to promote or cancel); +// this function owns only the bookkeeping, mapping each decision onto the +// path's current status. Mirrors prioritize.go's apply loop: Promote clears a +// Candidate to Selected; Cancel is captured as intent only — a Building path +// moves to Cancelling, and any pre-build status (Candidate, Selected, or +// Prioritized) drops straight to Cancelled. Nothing here touches a build +// runner — a Cancelling status records the intent for whichever stage owns +// runner interaction to enact. A decision naming a path not in +// the tree, a duplicate decision for the same path, or an action the path's +// current status does not support is a policy bug in the selector: it is +// logged as a warning and skipped rather than corrupting the tree. The +// returned bool reports whether any path status changed, so the caller can +// skip the conditional write on a no-op pass. +func (c *Controller) applySelection(batch entity.Batch, tree entity.SpeculationTree, decisions []entity.PathDecision) (entity.SpeculationTree, bool) { + if len(decisions) == 0 { + return tree, false + } + + paths := tree.Paths + changed := false + seen := make(map[string]bool, len(decisions)) + for _, d := range decisions { + if seen[d.PathID] { + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: duplicate decision for path", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + seen[d.PathID] = true + + idx := tree.PathIndex(d.PathID) + if idx == -1 { + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: path not found in tree", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + ) + continue + } + + info := paths[idx] + switch { + case d.Action == entity.SpeculationPathActionPromote && info.Status == entity.SpeculationPathStatusCandidate: + info.Status = entity.SpeculationPathStatusSelected + changed = true + metrics.NamedCounter(c.metricsScope, opName, "selected", 1) + + case d.Action == entity.SpeculationPathActionCancel && info.Status == entity.SpeculationPathStatusBuilding: + info.Status = entity.SpeculationPathStatusCancelling + changed = true + metrics.NamedCounter(c.metricsScope, opName, "cancelling", 1) + + case d.Action == entity.SpeculationPathActionCancel && + (info.Status == entity.SpeculationPathStatusCandidate || + info.Status == entity.SpeculationPathStatusSelected || + info.Status == entity.SpeculationPathStatusPrioritized): + info.Status = entity.SpeculationPathStatusCancelled + changed = true + metrics.NamedCounter(c.metricsScope, opName, "cancelled", 1) + + default: + metrics.NamedCounter(c.metricsScope, opName, "illegal_decision", 1) + c.logger.Warnw("illegal decision: action not valid for path status", + "batch_id", batch.ID, + "path_id", d.PathID, + "action", d.Action, + "status", string(info.Status), + ) + continue + } + + paths[idx] = info + } + + tree.Paths = paths + return tree, changed +} + +// startSpeculation publishes the batch to the build stage, then transitions +// it to Speculating. func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) error { c.logger.Infow("starting speculation", "batch_id", batch.ID, @@ -176,7 +532,8 @@ func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) e // out-of-the-way: the cancelled batch will never land, so it can no longer // conflict — drop it from the chain and proceed. Failed deps still cascade // via failOnDependency. If some deps are still in flight, the call is a -// no-op and waits for the next event. +// no-op: each dependency's own terminal pass re-publishes this batch to +// speculate (the terminal branch in Process), so waiting never needs a poll. // // TODO: when a dependency fails we currently fail this batch outright. // We will need to respeculate the failed paths — drop the failed dep @@ -235,7 +592,10 @@ func (c *Controller) tryFinalize(ctx context.Context, batch entity.Batch) error // dependencies has reached a non-succeeding terminal state, then publishes to // the conclude queue so the request store and request log get reconciled. // Without this transition the batch would sit in Speculating forever — no -// downstream event ever fires for it again. +// downstream event ever fires for it again. The batch's own dependents are +// woken right after the terminal CAS so the failure cascades downstream +// instead of stranding them mid-wait; a crash between the CAS and the +// fan-out is recovered by the terminal branch in Process on redelivery. func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, dep entity.Batch) error { metrics.NamedCounter(c.metricsScope, opName, "dependency_failed", 1) c.logger.Warnw("dependency in non-succeeding terminal state; failing batch", @@ -250,6 +610,10 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d return fmt.Errorf("failed to update batch %s state to failed: %w", batch.ID, err) } + if err := c.respeculateDependents(ctx, batch); err != nil { + return err + } + if err := c.publish(ctx, topickey.TopicKeyConclude, batch.ID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to conclude: %w", err) @@ -268,9 +632,8 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d // Order matters for correctness: // // 1. Cancel the in-flight Build entity (build.ID == batch.ID; one Get + one -// UpdateStatus covers all builds for this batch). A future external CI -// integration hooks in here. Idempotent: tolerate ErrNotFound (no build -// was scheduled), skip if already terminal. +// UpdateStatus covers all builds for this batch). Idempotent: tolerate +// ErrNotFound (no build was scheduled), skip if already terminal. // // 2. CAS the batch to terminal Cancelled. This must happen BEFORE the // dependent fan-out: tryFinalize only drops a Cancelled dep from the @@ -331,11 +694,6 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error // covers every build scheduled for the batch. Tolerates ErrNotFound (no // build was ever scheduled — the batch was cancelled before speculation // started building) and skips already-terminal builds. -// -// This is the hook point for a future external CI integration: today the -// system has no external runner, so the local state flip is the complete -// cancellation. Once a runner exists, it must be invoked here before the -// local UpdateStatus. func (c *Controller) cancelBuild(ctx context.Context, batch entity.Batch) error { build, err := c.store.GetBuildStore().Get(ctx, batch.ID) if err != nil { @@ -368,8 +726,10 @@ func (c *Controller) cancelBuild(ctx context.Context, batch entity.Batch) error // the message nacks and either an operator or the batch controller's own // crash-recovery can resolve the inconsistency. // -// Called both from the cancelBatch terminal flow and from the terminal -// self-heal branch on redelivery of an already-Cancelled batch. +// Called from the cancelBatch terminal flow, from failOnDependency, and from +// the terminal branch in Process — the latter both on the first pass after +// mergesignal drives a batch terminal (the normal dependent wake-up) and on +// redelivery (self-heal). func (c *Controller) respeculateDependents(ctx context.Context, batch entity.Batch) error { bd, err := c.store.GetBatchDependentStore().Get(ctx, batch.ID) if err != nil { diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 3cdbd7ce..03d00b4f 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: "submitqueue-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: "submitqueue-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: "submitqueue-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: "submitqueue-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: "submitqueue-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)) }