From 854d9b2a2a89c05510a3dc0589e3cc7f623d63b2 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Wed, 12 Aug 2026 15:50:30 -0700 Subject: [PATCH] fix(speculate): write the merge state before dispatching the batch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? `applyOutcome` published a batch to the merge topic before writing `BatchStateMerging`, so a lost compare-and-swap could leave Runway acting on an outcome that was never recorded. That ordering existed to avoid a stall, and the stall is real: nothing re-drives a batch stuck in `Merging`. `Process` self-heals only terminal and `Created` batches, `finalize` walks only heads that are still speculating, and the sole production reader of `BatchStateMerging` is the cancel controller — so a batch written `Merging` whose dispatch never went out would sit there forever. Giving that stall a repair path lets the write come first, which is the ordering the rest of the state machine already wants. ### What? `applyOutcome` is restructured into decide-state → recover → write → dispatch. The `terminal` bool falls out: a second switch mirrors the first and dispatches merge or conclude once the state write has landed. `recoverable` is hoisted above the switch so a cascade-decided *merge* gets a recovery message too, not just a cascade-decided failure. It needs one for the same reason: the write drops the batch out of the speculating set, and `Process`'s self-heal only ever names the trigger batch. `Process` gains a `BatchStateMerging` branch that re-sends the dispatch through the new `dispatchMerge` helper. That keeps the stable `IntentID`, the inverse of `fanout`'s `UniqueID` — for conclude a stable ID would suppress the repair, for merge it is what stops Runway merging the batch twice. One side benefit: a lost state CAS now means the dispatch is never sent at all, narrowing the window where a cancelled batch has a live merge request against it. ## Test Plan ✅ `bazel test //submitqueue/... //platform/...` — 68 tests pass New coverage: the dispatch follows the state write; a lost CAS publishes nothing; a cascade-merged batch gets its recovery signal before the write; `Process` on a `Merging` batch re-dispatches; `dispatchMerge` reuses one message ID per batch. `TestProcess_MergingRunsButDoesNotAct` asserted the old behaviour — that a `Merging` batch publishes nothing — and is replaced by `TestProcess_MergingSelfHeals`. --- .../controller/speculate/finalize.go | 77 ++++++----- .../controller/speculate/run_test.go | 120 ++++++++++++++++++ .../controller/speculate/speculate.go | 9 ++ .../controller/speculate/speculate_test.go | 9 +- 4 files changed, 171 insertions(+), 44 deletions(-) diff --git a/submitqueue/orchestrator/controller/speculate/finalize.go b/submitqueue/orchestrator/controller/speculate/finalize.go index 98d210cf..2f105317 100644 --- a/submitqueue/orchestrator/controller/speculate/finalize.go +++ b/submitqueue/orchestrator/controller/speculate/finalize.go @@ -302,45 +302,19 @@ func (c *Controller) recordOutcome(snap *snapshot, batchID string, decision outc // applyOutcome enacts a decided outcome on a batch, reporting whether the // state write landed. // -// The publish order differs per arm, but it is one rule read twice: a publish -// may precede a state write only when the consumer does not read the state -// that write produces. The merge stage correlates on the batch ID alone, so -// telling it before the write is safe — a batch recorded Merging that Runway -// never heard about would merely stall. Conclude does read the state (it -// reconciles requests from it and rejects a non-terminal batch outright), so -// it is published only after the write, or it would race the consumer into -// the dead-letter queue. -// -// Losing the state compare-and-swap is not an error: another writer got -// there, and the next run reads whatever they wrote. It is reported as not -// landed, because the outcome this run reached is not the one that took -// effect. +// Nothing is dispatched until the state it describes is durable, so no +// consumer can act on an outcome a lost compare-and-swap refused to write. The +// cost is a dispatch that fails on a batch finalize no longer walks, which the +// recovery message and Process's self-heal exist to repair. func (c *Controller) applyOutcome(ctx context.Context, store storage.Storage, batch entity.Batch, decision outcome, isTriggerBatch bool) (bool, error) { var state entity.BatchState - terminal := false switch decision { case outcomeMerge: state = entity.BatchStateMerging - // A batch merges once, so the dispatch names only that as its cause: a - // redelivery that re-derives outcomeMerge because the state write was - // lost dedups against the request already sent, instead of asking - // Runway to merge the same batch twice. - if err := c.publishBatchID(ctx, topickey.TopicKeyMerge, publish.IntentID(batch.ID, "merge-dispatch"), batch.ID, batch.Queue, batch.Queue); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) - return false, fmt.Errorf("failed to publish batch %s to merge: %w", batch.ID, err) - } case outcomeFail, outcomeCancel: - state, terminal = decision.terminalState() - // A batch decided by a cascade is not the one on the message, so no - // retry or dead letter would ever come back to it — give it a recovery - // message of its own before it turns terminal. - if !isTriggerBatch { - if err := c.recoverable(ctx, store, batch); err != nil { - return false, err - } - } + state, _ = decision.terminalState() default: // outcomeWait: nothing to enact. Listed explicitly so an unknown or @@ -348,6 +322,12 @@ func (c *Controller) applyOutcome(ctx context.Context, store storage.Storage, ba return false, nil } + if !isTriggerBatch { + if err := c.recoverable(ctx, store, batch); err != nil { + return false, err + } + } + // Through Transition, so the queue's membership record moves with the // state. A raw CAS would leave the batch filed under the bucket it just // left, and since records are only ever added, every later run of this @@ -369,7 +349,13 @@ func (c *Controller) applyOutcome(ctx context.Context, store storage.Storage, ba "state", string(state), ) - if terminal { + switch decision { + case outcomeMerge: + if err := c.dispatchMerge(ctx, batch); err != nil { + return true, err + } + + case outcomeFail, outcomeCancel: // Named for the run that decided it, so a redelivery re-deriving the // same outcome does not conclude the batch twice, and so it stays // distinct from the conclude mergesignal sends for a merged batch. @@ -383,15 +369,26 @@ func (c *Controller) applyOutcome(ctx context.Context, store storage.Storage, ba return true, nil } -// recoverable gives a batch a message of its own before this run makes it -// terminal, so its fan-out cannot be stranded by a failure afterwards. +// dispatchMerge hands a batch to the merge stage under a stable ID, so both a +// redelivery and the Merging self-heal dedupe against the request already sent +// rather than asking Runway to merge the batch twice. +func (c *Controller) dispatchMerge(ctx context.Context, batch entity.Batch) error { + if err := c.publishBatchID(ctx, topickey.TopicKeyMerge, publish.IntentID(batch.ID, "merge-dispatch"), batch.ID, batch.Queue, batch.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) + return fmt.Errorf("failed to publish batch %s to merge: %w", batch.ID, err) + } + return nil +} + +// recoverable gives a batch a message of its own before this run moves it out +// of the speculating set, so what has to follow the write cannot be stranded +// by a failure afterwards. // -// Every other terminal batch is repaired through the message that names it: a -// redelivery finds it terminal and re-publishes from Process's self-heal -// branch, and a persistent failure lands it in the dead-letter queue by name. -// A batch decided by a cascade has neither — it is not the batch on the -// message, and once terminal it is gone from the queue listing — so without -// this its requests would simply stay unreconciled. +// A batch named by a message is repaired through it: a redelivery re-publishes +// from one of Process's self-heal branches, and a persistent failure +// dead-letters by name. A cascade-decided batch has neither, and finalize only +// walks heads still speculating — so a merged one would never reach Runway, +// and a terminal one would leave its requests unreconciled. // // Distinct per publish: the guarantee being bought is that a message exists at // all, and a stable ID would let the queue answer "one already did" with a diff --git a/submitqueue/orchestrator/controller/speculate/run_test.go b/submitqueue/orchestrator/controller/speculate/run_test.go index aa1fe896..914d031a 100644 --- a/submitqueue/orchestrator/controller/speculate/run_test.go +++ b/submitqueue/orchestrator/controller/speculate/run_test.go @@ -700,6 +700,60 @@ func TestRun_MergeableHeadGainsNoNewPath(t *testing.T) { assert.Zero(t, spec.calls) } +// The dispatch is what takes a batch out of this stage's hands, so sending it +// before the write would let Runway act on an outcome a lost compare-and-swap +// refused to record. +func TestRun_MergeableHeadDispatchesAfterTheStateWrite(t *testing.T) { + ctrl := gomock.NewController(t) + passed := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionSucceeds) + + h := newRunHarness(t, ctrl, &scriptedSpeculator{}, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSucceeded}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSucceeded}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(passed, entity.SpeculationPathStatusPassed)}, + Version: 1, + }, nil).AnyTimes() + + var publishedBeforeWrite []string + h.batches.EXPECT(). + Update(gomock.Any(), updateTo{id: head, state: entity.BatchStateMerging}, int32(1), int32(2)). + DoAndReturn(func(context.Context, entity.Batch, int32, int32) error { + publishedBeforeWrite = append([]string(nil), h.published...) + return nil + }) + + require.NoError(t, h.run(head)) + assert.Equal(t, []string{"submitqueue-merge"}, h.published) + assert.Empty(t, publishedBeforeWrite, + "nothing may reach the merge stage before the state it acts on is written") +} + +// The other half: a lost write means another writer owns the batch, so the +// dispatch it would have justified is never sent. +func TestRun_MergeableHeadDispatchesNothingWhenTheStateCASLoses(t *testing.T) { + ctrl := gomock.NewController(t) + passed := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionSucceeds) + + h := newRunHarness(t, ctrl, &scriptedSpeculator{}, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSucceeded}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSucceeded}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(passed, entity.SpeculationPathStatusPassed)}, + Version: 1, + }, nil).AnyTimes() + h.batches.EXPECT(). + Update(gomock.Any(), updateTo{id: head, state: entity.BatchStateMerging}, int32(1), int32(2)). + Return(storage.ErrVersionMismatch) + + require.NoError(t, h.run(head)) + assert.Empty(t, h.published) +} + // A head with no future left fails, and the write order is what keeps conclude // usable: conclude reconciles requests from the batch's state and rejects a // non-terminal one, so it is published only once the terminal write has landed. @@ -1049,6 +1103,40 @@ func TestFanout_MintsADistinctMessageIDPerPublish(t *testing.T) { assert.NotEqual(t, ids[0], ids[1]) } +// The merge dispatch is the mirror image: a batch merges once, so the ID has +// to be stable across the redelivery and the self-heal that both re-derive it. +func TestDispatchMerge_ReusesOneMessageIDPerBatch(t *testing.T) { + ctrl := gomock.NewController(t) + + var ids []string + pub := queuemock.NewMockPublisher(ctrl) + pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + ids = append(ids, msg.ID) + return nil + }, + ).Times(2) + q := queuemock.NewMockQueue(ctrl) + q.EXPECT().Publisher().Return(pub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyMerge, Name: "submitqueue-merge", Queue: q}, + }) + require.NoError(t, err) + + c := NewController( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, staticStorageFactory{store: storagemock.NewMockStorage(ctrl)}, + staticSpeculatorFactory{}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate", + ) + batch := entity.Batch{ID: head, Queue: "q"} + + require.NoError(t, c.dispatchMerge(context.Background(), batch)) + require.NoError(t, c.dispatchMerge(context.Background(), batch)) + + require.Len(t, ids, 2) + assert.Equal(t, ids[0], ids[1]) +} + // cascadePair wires a queue where `prerequisite` reaches a terminal outcome // and `derived` fails only because of it: derived bet that prerequisite would // not succeed, built on that, and its build failed. Until prerequisite is terminal the @@ -1248,6 +1336,38 @@ func TestRun_TriggerBatchNeedsNoRecoverySignal(t *testing.T) { assert.Equal(t, []string{"conclude"}, h.published) } +// Merging needs the same signal for the same reason: the write drops the batch +// out of the speculating set, so nothing else would ever dispatch it. +func TestRun_CascadeDerivedBatchIsGivenARecoverySignalBeforeItMerges(t *testing.T) { + ctrl := gomock.NewController(t) + + merging := entity.Batch{ID: "q/batch/derived", Queue: "q", State: entity.BatchStateSpeculating, Version: 1} + h := newRunHarness(t, ctrl, &scriptedSpeculator{}, []entity.Batch{merging}) + h.noBuildsDispatched() + + // No dependencies, so the passed path has nothing left to settle. + passedPath := entity.SpeculationPath{Head: merging.ID} + h.pathSets.EXPECT().Get(gomock.Any(), merging.ID).Return(entity.SpeculationPathSet{ + Head: merging.ID, + Paths: []entity.SpeculationPathEntry{entryFor(passedPath, entity.SpeculationPathStatusPassed)}, + Version: 1, + }, nil).AnyTimes() + + var publishedBeforeWrite []string + h.batches.EXPECT(). + Update(gomock.Any(), updateTo{id: merging.ID, state: entity.BatchStateMerging}, int32(1), int32(2)). + DoAndReturn(func(context.Context, entity.Batch, int32, int32) error { + publishedBeforeWrite = append([]string(nil), h.published...) + return nil + }) + + // The message names some other batch, so a retry would never come back here. + require.NoError(t, h.run(head)) + + assert.Equal(t, []string{"speculate"}, publishedBeforeWrite) + assert.Equal(t, []string{"speculate", "submitqueue-merge"}, h.published) +} + // A cancelling path whose build is still running finishes only when CI actually // stops. The run writes nothing (the intent is already recorded), publishes // nothing (the poll loop is what keeps asking the runner to stop), and the diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index 9ce4782a..17f29f8c 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -137,6 +137,15 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er } } + // A Merging batch has left the set finalize walks, so a message naming it + // is the only thing that will look at it again. + if batch.State == entity.BatchStateMerging { + metrics.NamedCounter(c.metricsScope, opName, "self_heal_merging", 1) + if err := c.dispatchMerge(ctx, batch); err != nil { + return c.attributed(err, entity.BatchSubject(batch.ID)) + } + } + return c.run(ctx, store, batch) } diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 3d8a1bc3..1fe4ef27 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate_test.go +++ b/submitqueue/orchestrator/controller/speculate/speculate_test.go @@ -276,9 +276,10 @@ func TestProcess_TerminalReplansQueue(t *testing.T) { "the dependent must be re-planned against the terminal outcome, which it can only be weighed against if the terminal batch comes too") } -// A Merging batch is the merge stage's to finish; the run still happens for the -// rest of the queue, but this batch is not an action target. -func TestProcess_MergingRunsButDoesNotAct(t *testing.T) { +// A Merging batch has left the speculating set, so a message naming it is the +// only thing that will look at it again: it re-sends the dispatch to repair +// one lost after the state write. +func TestProcess_MergingSelfHeals(t *testing.T) { ctrl := gomock.NewController(t) h := newProcHarness(t, ctrl, nil) batch := testBatch(entity.BatchStateMerging) @@ -287,7 +288,7 @@ func TestProcess_MergingRunsButDoesNotAct(t *testing.T) { h.listsInFlight() require.NoError(t, h.process(t, ctrl, batch.ID)) - assert.Empty(t, h.published) + assert.Equal(t, []string{"submitqueue-merge"}, h.published) } func TestProcess_Errors(t *testing.T) {