From 55134c7542d261caccdc289ae50e1bda28507025 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 30 Jul 2026 10:41:42 -0700 Subject: [PATCH] feat(orchestrator): re-plan the queue from the Speculator each run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? The speculate controller had no speculation in it: it advanced one batch at a time along a single hard-coded chain — every dependency assumed to pass — and nothing ever called the Speculator, the Generator, or the Allocator that had been built for it. This wires that machinery into the pipeline: a queue's paths are chosen by a swappable extension within a build budget, and each is built against only the dependencies it assumes will succeed. ### What? Every message is a dirty signal naming a batch; the controller re-plans that batch's whole queue from a single read: read the state, cancel paths whose assumptions a finished dependency has proven wrong, ask the Speculator, filter its proposals, dispatch what survives. Nothing carries over between runs, so duplicated or reordered signals are harmless and a later run repairs whatever an earlier one left half-done. `doc.go` explains the model in plain terms — no vocabulary section, every term is defined where it is used. The path set keeps exactly one writer — this run. The build stages record what CI did on per-build records; the run folds those into the set and alone decides each path's status. Pending paths are re-dispatched every run until their build is seen running, dispatches partition by batch so heads proceed in parallel, and cancelling paths need no dispatch at all: the poll loop reads the stop off the set and enacts it. Speculation lands inert. The wiring layer passes a placeholder Speculator that proposes nothing, so the run executes end to end but funds no paths; composing real per-queue speculators and turning the feature on is the wiring change at the top of this stack. Batch outcomes still come from the legacy per-batch finalizer, which waits on every dependency — strictly stricter than path-aware finalization, so the system stays correct until the next commit replaces it. ## Test Plan ✅ `bazel test //submitqueue/orchestrator/controller/speculate/...` — assumption checks and proposal filtering are table-driven; run tests cover funding a first path, re-dispatching pending paths, broken-path cancellation, build results recorded onto paths, lost CAS races skipped rather than failed, and Speculator errors abandoning the run. ✅ `make fmt`, `make gazelle` --- .../orchestrator/server/BUILD.bazel | 2 + .../submitqueue/orchestrator/server/main.go | 26 +- submitqueue/orchestrator/BUILD.bazel | 1 + .../controller/speculate/BUILD.bazel | 20 +- .../controller/speculate/check.go | 169 ++++++ .../controller/speculate/check_test.go | 259 ++++++++++ .../controller/speculate/dispatch.go | 213 ++++++++ .../orchestrator/controller/speculate/doc.go | 98 ++++ .../orchestrator/controller/speculate/run.go | 294 +++++++++++ .../controller/speculate/run_test.go | 489 ++++++++++++++++++ .../controller/speculate/snapshot.go | 86 +++ .../controller/speculate/snapshot_test.go | 101 ++++ .../controller/speculate/speculate.go | 157 +++--- .../controller/speculate/speculate_test.go | 59 ++- submitqueue/orchestrator/pipeline.go | 6 +- 15 files changed, 1895 insertions(+), 85 deletions(-) create mode 100644 submitqueue/orchestrator/controller/speculate/check.go create mode 100644 submitqueue/orchestrator/controller/speculate/check_test.go create mode 100644 submitqueue/orchestrator/controller/speculate/dispatch.go create mode 100644 submitqueue/orchestrator/controller/speculate/doc.go create mode 100644 submitqueue/orchestrator/controller/speculate/run.go create mode 100644 submitqueue/orchestrator/controller/speculate/run_test.go create mode 100644 submitqueue/orchestrator/controller/speculate/snapshot.go create mode 100644 submitqueue/orchestrator/controller/speculate/snapshot_test.go diff --git a/service/submitqueue/orchestrator/server/BUILD.bazel b/service/submitqueue/orchestrator/server/BUILD.bazel index e197655a..5860b070 100644 --- a/service/submitqueue/orchestrator/server/BUILD.bazel +++ b/service/submitqueue/orchestrator/server/BUILD.bazel @@ -25,6 +25,7 @@ go_library( "//platform/http:go_default_library", "//platform/pipeline:go_default_library", "//submitqueue/core/changeset:go_default_library", + "//submitqueue/entity:go_default_library", "//submitqueue/extension/buildrunner:go_default_library", "//submitqueue/extension/buildrunner/fake:go_default_library", "//submitqueue/extension/changeprovider:go_default_library", @@ -37,6 +38,7 @@ go_library( "//submitqueue/extension/conflict/fake:go_default_library", "//submitqueue/extension/conflict/fileoverlap:go_default_library", "//submitqueue/extension/conflict/none:go_default_library", + "//submitqueue/extension/speculation/speculator:go_default_library", "//submitqueue/extension/storage/mysql:go_default_library", "//submitqueue/extension/validator/fake:go_default_library", "//submitqueue/orchestrator:go_default_library", diff --git a/service/submitqueue/orchestrator/server/main.go b/service/submitqueue/orchestrator/server/main.go index 0c5bbc61..2d05b957 100644 --- a/service/submitqueue/orchestrator/server/main.go +++ b/service/submitqueue/orchestrator/server/main.go @@ -42,11 +42,13 @@ import ( "github.com/uber/submitqueue/platform/http" "github.com/uber/submitqueue/platform/pipeline" "github.com/uber/submitqueue/submitqueue/core/changeset" + "github.com/uber/submitqueue/submitqueue/entity" "github.com/uber/submitqueue/submitqueue/extension/changeprovider" cpfake "github.com/uber/submitqueue/submitqueue/extension/changeprovider/fake" githubprovider "github.com/uber/submitqueue/submitqueue/extension/changeprovider/github" phabprovider "github.com/uber/submitqueue/submitqueue/extension/changeprovider/phabricator" routingprovider "github.com/uber/submitqueue/submitqueue/extension/changeprovider/routing" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" mysqlstorage "github.com/uber/submitqueue/submitqueue/extension/storage/mysql" validatorfake "github.com/uber/submitqueue/submitqueue/extension/validator/fake" "github.com/uber/submitqueue/submitqueue/orchestrator" @@ -196,7 +198,12 @@ func run() error { BuildRunner: profiles.BuildRunnerFactory(), ChangeProvider: profiles.ChangeProviderFactory(), Analyzer: profiles.AnalyzerFactory(), - Validator: validatorfake.NewFactory(), + // Speculation is wired but inert: the placeholder below proposes + // nothing, so no path is ever funded and no speculative build starts. + // The wiring change at the top of this stack replaces it with real + // per-queue speculators composed from each profile's scorer. + Speculator: noopSpeculators{}, + Validator: validatorfake.NewFactory(), } // Assemble the pipeline: one call builds the topic registry, creates @@ -431,3 +438,20 @@ func parseTimeout(envVal string, defaultVal time.Duration) time.Duration { } return defaultVal } + +// noopSpeculators resolves every queue to a speculator that proposes nothing. +// It keeps the speculate stage inert — no path funded, no build started — +// until per-queue speculators are composed in the profiles. +type noopSpeculators struct{} + +// For returns the propose-nothing speculator for any queue. +func (noopSpeculators) For(speculator.Config) (speculator.Speculator, error) { + return noopSpeculator{}, nil +} + +type noopSpeculator struct{} + +// Speculate proposes no actions, whatever the queue looks like. +func (noopSpeculator) Speculate(context.Context, []entity.Batch, []entity.SpeculationPathSet) ([]entity.Speculation, error) { + return nil, nil +} diff --git a/submitqueue/orchestrator/BUILD.bazel b/submitqueue/orchestrator/BUILD.bazel index 975406ed..82e60a9d 100644 --- a/submitqueue/orchestrator/BUILD.bazel +++ b/submitqueue/orchestrator/BUILD.bazel @@ -14,6 +14,7 @@ go_library( "//submitqueue/extension/buildrunner:go_default_library", "//submitqueue/extension/changeprovider:go_default_library", "//submitqueue/extension/conflict:go_default_library", + "//submitqueue/extension/speculation/speculator:go_default_library", "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/validator:go_default_library", "//submitqueue/orchestrator/controller:go_default_library", diff --git a/submitqueue/orchestrator/controller/speculate/BUILD.bazel b/submitqueue/orchestrator/controller/speculate/BUILD.bazel index 39e36977..c4085b42 100644 --- a/submitqueue/orchestrator/controller/speculate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/speculate/BUILD.bazel @@ -2,15 +2,23 @@ load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "go_default_library", - srcs = ["speculate.go"], + srcs = [ + "check.go", + "dispatch.go", + "doc.go", + "run.go", + "snapshot.go", + "speculate.go", + ], importpath = "github.com/uber/submitqueue/submitqueue/orchestrator/controller/speculate", visibility = ["//visibility:public"], deps = [ - "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", + "//submitqueue/core/publish:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/speculator:go_default_library", "//submitqueue/extension/storage:go_default_library", "@com_github_uber_go_tally//:go_default_library", "@org_uber_go_zap//:go_default_library", @@ -19,7 +27,12 @@ go_library( go_test( name = "go_default_test", - srcs = ["speculate_test.go"], + srcs = [ + "check_test.go", + "run_test.go", + "snapshot_test.go", + "speculate_test.go", + ], embed = [":go_default_library"], deps = [ "//platform/base/messagequeue:go_default_library", @@ -28,6 +41,7 @@ go_test( "//platform/extension/messagequeue/mock:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/speculator: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/check.go b/submitqueue/orchestrator/controller/speculate/check.go new file mode 100644 index 00000000..175127fc --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/check.go @@ -0,0 +1,169 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import "github.com/uber/submitqueue/submitqueue/entity" + +// rejection is why one proposed action was dropped. The reasons are the metric +// dimension for a misbehaving Speculator, so each names a distinct fault. +type rejection string + +const ( + // rejectUnknownAction is a zero-value or unrecognized action. + rejectUnknownAction rejection = "unknown_action" + // rejectUnknownHead names a batch this run never read. + rejectUnknownHead rejection = "unknown_head" + // rejectHeadNotSpeculating targets a batch that is not open to new work. + rejectHeadNotSpeculating rejection = "head_not_speculating" + // rejectMalformedPath is a path whose assumptions do not line up with its + // head's dependency list: one missing or extra, a duplicate, a wrong head, + // or a made-up assumption value. + rejectMalformedPath rejection = "malformed_path" + // rejectBrokenAssumption is a path with an assumption a finished + // dependency has already proven wrong. + rejectBrokenAssumption rejection = "broken_assumption" + // rejectPathTerminal would rebuild a path whose build already finished. + rejectPathTerminal rejection = "path_terminal" + // rejectCancelNotInFlight cancels a path that is not running. + rejectCancelNotInFlight rejection = "cancel_not_in_flight" + // rejectCancelPassed would throw away a build that already passed. + rejectCancelPassed rejection = "cancel_passed" +) + +// filterProposals narrows a Speculator's proposals down to the ones the +// controller is willing to enact, returning the survivors and a reason for +// each drop. +// +// The Speculator is an extension, so its output is untrusted input: it decides +// which paths run, never whether a batch merges or fails. Every rule here +// protects an invariant the extension could otherwise break — acting on a batch +// that is finalizing, resurrecting a path a resolved dependency has ruled out, +// or discarding a passed build the queue is about to merge on. A proposal that +// trips one of these is a bug in the Speculator, not a normal outcome, which is +// why the caller counts them. +func filterProposals(proposals []entity.Speculation, snap snapshot) ([]entity.Speculation, []rejection) { + var kept []entity.Speculation + var rejected []rejection + + for _, proposal := range proposals { + if reason, ok := rejectionReason(proposal, snap); ok { + rejected = append(rejected, reason) + continue + } + kept = append(kept, proposal) + } + + return kept, rejected +} + +// rejectionReason reports why a proposal cannot be enacted, or ok=false if +// it can. +func rejectionReason(proposal entity.Speculation, snap snapshot) (rejection, bool) { + switch proposal.Action { + case entity.PathActionBuild, entity.PathActionCancel: + default: + return rejectUnknownAction, true + } + + head, ok := snap.batches[proposal.Path.Head] + if !ok { + return rejectUnknownHead, true + } + // Only a speculating head is open to new work. Batches in every other + // state are facts the Speculator may reason from, never action targets. + if head.State != entity.BatchStateSpeculating { + return rejectHeadNotSpeculating, true + } + if !isWellFormed(proposal.Path, head) { + return rejectMalformedPath, true + } + if assumptionBroken(proposal.Path, snap) { + return rejectBrokenAssumption, true + } + + entry, stored := findPath(snap.pathSets[head.ID], proposal.Path.ID()) + + if proposal.Action == entity.PathActionCancel { + // Cancelling is only meaningful for a path that is actually running, + // and never for one that passed: that build is the head's way out of + // the queue, and the budget it holds is already spent. + if !stored { + return rejectCancelNotInFlight, true + } + if entry.Status == entity.SpeculationPathStatusPassed { + return rejectCancelPassed, true + } + if entry.Status.IsTerminal() { + return rejectCancelNotInFlight, true + } + return "", false + } + + // A build proposal for a path whose build already finished would discard a + // recorded result and start the same work again. + if stored && entry.Status.IsTerminal() { + return rejectPathTerminal, true + } + + return "", false +} + +// isWellFormed reports whether a path is a proper guess about its head: +// exactly one assumption for each of the head's dependencies, no more and no +// fewer, and every assumption a real value. +// +// A malformed path is not merely suboptimal, it is unmergeable — the merge +// preconditions are read off the path's assumptions, so a path missing a +// dependency would let its head merge without waiting for it. +func isWellFormed(path entity.SpeculationPath, head entity.Batch) bool { + if path.Head != head.ID { + return false + } + if len(path.Dependencies) != len(head.Dependencies) { + return false + } + + required := make(map[string]struct{}, len(head.Dependencies)) + for _, dep := range head.Dependencies { + required[dep] = struct{}{} + } + + for _, dep := range path.Dependencies { + if _, ok := required[dep.Batch]; !ok { + return false + } + delete(required, dep.Batch) + + switch dep.Assumption { + case entity.DependencyAssumptionSucceeds, + entity.DependencyAssumptionFails, + entity.DependencyAssumptionIgnored: + default: + return false + } + } + + return len(required) == 0 +} + +// findPath returns the entry for a path ID in the set. +func findPath(set entity.SpeculationPathSet, pathID string) (entity.SpeculationPathEntry, bool) { + for _, entry := range set.Paths { + if entry.ID == pathID { + return entry, true + } + } + return entity.SpeculationPathEntry{}, false +} diff --git a/submitqueue/orchestrator/controller/speculate/check_test.go b/submitqueue/orchestrator/controller/speculate/check_test.go new file mode 100644 index 00000000..e092a3b6 --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/check_test.go @@ -0,0 +1,259 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/submitqueue/entity" +) + +// checkSnapshot builds a snapshot with a speculating head over dep1 and dep2, +// both unresolved, and the head's path set seeded with the given entries. +func checkSnapshot(headState entity.BatchState, entries ...entity.SpeculationPathEntry) snapshot { + snap := snapshot{ + batches: map[string]entity.Batch{ + head: {ID: head, State: headState, Dependencies: []string{dep1, dep2}}, + dep1: {ID: dep1, State: entity.BatchStateSpeculating}, + dep2: {ID: dep2, State: entity.BatchStateSpeculating}, + }, + pathSets: map[string]entity.SpeculationPathSet{}, + } + if len(entries) > 0 { + snap.pathSets[head] = entity.SpeculationPathSet{Head: head, Paths: entries} + } + return snap +} + +func entryFor(path entity.SpeculationPath, status entity.SpeculationPathStatus) entity.SpeculationPathEntry { + return entity.SpeculationPathEntry{ID: path.ID(), Path: path, Status: status, Attempt: 1} +} + +// A well-formed build proposal on a speculating head survives. +func TestFilterProposals_KeepsValidProposals(t *testing.T) { + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + snap := checkSnapshot(entity.BatchStateSpeculating) + + kept, rejected := filterProposals([]entity.Speculation{ + {Path: path, Action: entity.PathActionBuild}, + }, snap) + + assert.Empty(t, rejected) + require.Len(t, kept, 1) + assert.Equal(t, path.ID(), kept[0].Path.ID()) +} + +func TestFilterProposals_Rejects(t *testing.T) { + valid := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + + tests := []struct { + name string + proposal entity.Speculation + snap snapshot + want rejection + }{ + { + name: "zero-value action", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionUnknown}, + snap: checkSnapshot(entity.BatchStateSpeculating), + want: rejectUnknownAction, + }, + { + name: "head this run never read", + proposal: entity.Speculation{ + Path: entity.SpeculationPath{Head: "q/batch/ghost"}, + Action: entity.PathActionBuild, + }, + snap: checkSnapshot(entity.BatchStateSpeculating), + want: rejectUnknownHead, + }, + { + name: "head already merging", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionBuild}, + snap: checkSnapshot(entity.BatchStateMerging), + want: rejectHeadNotSpeculating, + }, + { + name: "head already cancelling", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionBuild}, + snap: checkSnapshot(entity.BatchStateCancelling), + want: rejectHeadNotSpeculating, + }, + { + name: "path missing one of the head's dependencies", + proposal: entity.Speculation{ + Path: pathOver(entity.DependencyAssumptionSucceeds), + Action: entity.PathActionBuild, + }, + snap: checkSnapshot(entity.BatchStateSpeculating), + want: rejectMalformedPath, + }, + { + name: "cancel on a path that is not stored", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionCancel}, + snap: checkSnapshot(entity.BatchStateSpeculating), + want: rejectCancelNotInFlight, + }, + { + name: "cancel on a path whose build already finished", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionCancel}, + snap: checkSnapshot(entity.BatchStateSpeculating, + entryFor(valid, entity.SpeculationPathStatusFailed)), + want: rejectCancelNotInFlight, + }, + { + name: "cancel would discard a passed build", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionCancel}, + snap: checkSnapshot(entity.BatchStateSpeculating, + entryFor(valid, entity.SpeculationPathStatusPassed)), + want: rejectCancelPassed, + }, + { + name: "build would resurrect a finished path", + proposal: entity.Speculation{Path: valid, Action: entity.PathActionBuild}, + snap: checkSnapshot(entity.BatchStateSpeculating, + entryFor(valid, entity.SpeculationPathStatusPassed)), + want: rejectPathTerminal, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + kept, rejected := filterProposals([]entity.Speculation{tt.proposal}, tt.snap) + assert.Empty(t, kept) + require.Len(t, rejected, 1) + assert.Equal(t, tt.want, rejected[0]) + }) + } +} + +// A path a resolved dependency has already ruled out must not be funded, even +// if the Speculator proposes it. +func TestFilterProposals_RejectsBrokenPath(t *testing.T) { + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + + snap := checkSnapshot(entity.BatchStateSpeculating) + snap.batches[dep1] = entity.Batch{ID: dep1, State: entity.BatchStateFailed} + + kept, rejected := filterProposals([]entity.Speculation{ + {Path: path, Action: entity.PathActionBuild}, + }, snap) + + assert.Empty(t, kept) + require.Len(t, rejected, 1) + assert.Equal(t, rejectBrokenAssumption, rejected[0]) +} + +// Cancelling a path that is still running is the Speculator's one legitimate +// cancel, so it must survive. +func TestFilterProposals_KeepsCancelOfRunningPath(t *testing.T) { + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + + for _, status := range []entity.SpeculationPathStatus{ + entity.SpeculationPathStatusPending, + entity.SpeculationPathStatusBuilding, + } { + t.Run(string(status), func(t *testing.T) { + snap := checkSnapshot(entity.BatchStateSpeculating, entryFor(path, status)) + + kept, rejected := filterProposals([]entity.Speculation{ + {Path: path, Action: entity.PathActionCancel}, + }, snap) + + assert.Empty(t, rejected) + assert.Len(t, kept, 1) + }) + } +} + +// One bad proposal must not discard the good ones alongside it. +func TestFilterProposals_FiltersIndependently(t *testing.T) { + good := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + bad := pathOver(entity.DependencyAssumptionSucceeds) + + kept, rejected := filterProposals([]entity.Speculation{ + {Path: bad, Action: entity.PathActionBuild}, + {Path: good, Action: entity.PathActionBuild}, + }, checkSnapshot(entity.BatchStateSpeculating)) + + require.Len(t, kept, 1) + assert.Equal(t, good.ID(), kept[0].Path.ID()) + assert.Equal(t, []rejection{rejectMalformedPath}, rejected) +} +func TestIsWellFormed(t *testing.T) { + headBatch := entity.Batch{ID: head, Dependencies: []string{dep1, dep2}} + + tests := []struct { + name string + path entity.SpeculationPath + want bool + }{ + { + name: "one assumption per dependency", + path: pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails), + want: true, + }, + { + name: "order does not matter", + path: entity.SpeculationPath{Head: head, Dependencies: []entity.PathDependency{ + {Batch: dep2, Assumption: entity.DependencyAssumptionSucceeds}, + {Batch: dep1, Assumption: entity.DependencyAssumptionIgnored}, + }}, + want: true, + }, + { + name: "missing a dependency", + path: pathOver(entity.DependencyAssumptionSucceeds), + want: false, + }, + { + name: "duplicate dependency in place of another", + path: entity.SpeculationPath{Head: head, Dependencies: []entity.PathDependency{ + {Batch: dep1, Assumption: entity.DependencyAssumptionSucceeds}, + {Batch: dep1, Assumption: entity.DependencyAssumptionFails}, + }}, + want: false, + }, + { + name: "dependency the head does not have", + path: entity.SpeculationPath{Head: head, Dependencies: []entity.PathDependency{ + {Batch: dep1, Assumption: entity.DependencyAssumptionSucceeds}, + {Batch: "q/batch/stranger", Assumption: entity.DependencyAssumptionFails}, + }}, + want: false, + }, + { + name: "unknown assumption value", + path: pathOver(entity.DependencyAssumptionUnknown, entity.DependencyAssumptionFails), + want: false, + }, + { + name: "wrong head", + path: entity.SpeculationPath{Head: "q/batch/other", Dependencies: []entity.PathDependency{ + {Batch: dep1, Assumption: entity.DependencyAssumptionSucceeds}, + {Batch: dep2, Assumption: entity.DependencyAssumptionFails}, + }}, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, isWellFormed(tt.path, headBatch)) + }) + } +} diff --git a/submitqueue/orchestrator/controller/speculate/dispatch.go b/submitqueue/orchestrator/controller/speculate/dispatch.go new file mode 100644 index 00000000..f1854690 --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/dispatch.go @@ -0,0 +1,213 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import ( + "context" + "errors" + "fmt" + "time" + + "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/storage" +) + +// dispatch saves each head's decisions and hands the build stage its work. +// Everything decided this run — build results, broken-path cancellations and +// accepted proposals — is folded together per head, so a head costs one +// compare-and-swap however many of its paths changed. +func (c *Controller) dispatch(ctx context.Context, queue string, snap snapshot, kept []entity.Speculation) error { + nowMs := time.Now().UnixMilli() + + // Group the accepted proposals by head so each head is written once. + byHead := make(map[string][]entity.Speculation, len(kept)) + for _, proposal := range kept { + byHead[proposal.Path.Head] = append(byHead[proposal.Path.Head], proposal) + } + + for _, batch := range snap.speculating { + // A head with no stored set is one nothing has been funded for yet. It + // gets an empty set to fold this run's proposals into, which persist + // then creates; a head that ends the run with no paths writes nothing + // at all, because nothing marked it changed. + set, exists := snap.pathSets[batch.ID] + if !exists { + set = entity.SpeculationPathSet{Head: batch.ID} + } + + changed := snap.isDirty(batch.ID) + for _, proposal := range byHead[batch.ID] { + if applyProposal(&set, proposal, nowMs) { + changed = true + } + } + + if changed { + if err := c.persist(ctx, set, exists); err != nil { + if errors.Is(err, storage.ErrVersionMismatch) { + // Skipped rather than failed, and nothing is lost by that. + // + // The only other writer of a path set is another speculate + // run on this same queue — overlap happens when a message + // is redelivered or a partition moves between consumers. + // Whichever run wins has read the same queue and reached + // the same kind of conclusion, and its own dispatch step + // dispatches and publishes for this head, so the work + // happens either way. Failing the message instead would + // replay every other head's decisions to retry this one + // against a snapshot that is now stale anyway; any later + // signal on the queue re-plans this head from stored + // state. + metrics.NamedCounter(c.metricsScope, opName, "path_set_cas_lost", 1) + c.logger.Infow("lost a path set write; the next run re-plans this head", + "batch_id", batch.ID, + "queue", queue, + ) + continue + } + return err + } + } + + // Dispatch whenever the head has paths waiting to start, whether or not + // this run changed them: a pending path whose earlier dispatch was lost + // is re-sent until the build stage records it building or terminal. + // Cancelling paths need no dispatch at all — the poll loop reads the + // stop off the set and enacts it. Partitioned by batch, so heads + // dispatch in parallel while one head's dispatches stay ordered. + if hasActionablePaths(set) { + if err := c.publishBatchID(ctx, topickey.TopicKeyBuild, batch.ID, batch.ID); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) + return fmt.Errorf("failed to publish batch %s to build: %w", batch.ID, err) + } + } + } + + return nil +} + +// persist writes a head's path set, creating it if this run is the first to +// fund the head. +func (c *Controller) persist(ctx context.Context, set entity.SpeculationPathSet, exists bool) error { + store := c.store.GetSpeculationPathSetStore() + + if !exists { + set.Version = 1 + if err := store.Create(ctx, set); err != nil { + if errors.Is(err, storage.ErrAlreadyExists) { + // Another writer created it between this run's read and now. + // Treat it as a lost race: the next run reads the winner. + return storage.ErrVersionMismatch + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to create path set for batch %s: %w", set.Head, err) + } + return nil + } + + newVersion := set.Version + 1 + if err := store.Update(ctx, set, set.Version, newVersion); err != nil { + if errors.Is(err, storage.ErrVersionMismatch) { + return err + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to update path set for batch %s: %w", set.Head, err) + } + return nil +} + +// applyProposal folds one accepted proposal into the set and reports whether +// it changed anything. The return value says only whether the stored set +// needs rewriting, not whether to dispatch: the build stage is driven by the +// statuses in the set, so a path this returns false for is still dispatched +// if its status calls for it. +// +// The status transitions here are the ones drawn in the package doc's +// lifecycle diagram; the two rules worth calling out: +// +// - A build proposal for a path already in flight is a no-op, not a new +// attempt. The Allocator matches candidates to in-flight paths by ID +// precisely so a path keeps the slot it holds; re-funding it would start +// a second build for work already running. +// - A build proposal for a *terminal* path resurrects it: same path (its +// identity is its assumptions), new attempt. The attempt counter moves +// and the status returns to pending — the one backwards step in the +// lifecycle. The build link is keyed by (path, attempt), so the previous +// attempt's build stays addressable rather than being overwritten. +func applyProposal(set *entity.SpeculationPathSet, proposal entity.Speculation, nowMs int64) bool { + pathID := proposal.Path.ID() + + for i := range set.Paths { + entry := &set.Paths[i] + if entry.ID != pathID { + continue + } + + switch proposal.Action { + case entity.PathActionCancel: + if entry.Status == entity.SpeculationPathStatusCancelling { + return false + } + // Recording the intent is the whole enactment from this side: the + // poll loop reads the status off the set on every poll and asks + // the runner to stop the build. What ends the cancellation is CI + // actually stopping and a later run recording that on the path. + entry.Status = entity.SpeculationPathStatusCancelling + entry.UpdatedAtMs = nowMs + return true + + case entity.PathActionBuild: + if !entry.Status.IsTerminal() { + return false + } + entry.Status = entity.SpeculationPathStatusPending + entry.Attempt++ + entry.UpdatedAtMs = nowMs + return true + } + return false + } + + // A cancel can only apply to a path already in the set; filterProposals + // has already rejected the ones that are not. + if proposal.Action != entity.PathActionBuild { + return false + } + + set.Paths = append(set.Paths, entity.SpeculationPathEntry{ + ID: pathID, + Path: proposal.Path, + Status: entity.SpeculationPathStatusPending, + Attempt: 1, + Version: 1, + CreatedAtMs: nowMs, + UpdatedAtMs: nowMs, + }) + return true +} + +// hasActionablePaths reports whether a head has work for the build stage: a +// path waiting to start. Cancelling paths are not the build stage's work — +// the poll loop stops their builds — so they trigger no dispatch. +func hasActionablePaths(set entity.SpeculationPathSet) bool { + for _, entry := range set.Paths { + if entry.Status == entity.SpeculationPathStatusPending { + return true + } + } + return false +} diff --git a/submitqueue/orchestrator/controller/speculate/doc.go b/submitqueue/orchestrator/controller/speculate/doc.go new file mode 100644 index 00000000..10127d7a --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/doc.go @@ -0,0 +1,98 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package speculate plans a queue's speculative builds: which guesses about +// the queue's future are worth building, within the queue's cap on concurrent +// builds (the build budget). +// +// Batch outcomes — merge or fail — are still decided by the legacy per-batch +// finalizer in speculate.go, which waits on every dependency. Deriving them +// from the paths planned here replaces it in the next change; this package +// doc grows with it. +// +// # Why speculation +// +// Batches in a queue depend on the batches ahead of them, so without +// speculation everything is serial: C waits for B, B waits for A. Speculation +// builds a batch against a guess about how its dependencies turn out. When +// the guess holds, the head's build has already run by the time its +// dependencies resolve — it never waits for a build of its own to start +// afterwards. +// +// # Paths +// +// The batch being speculated on is the head. One complete guess about it is a +// path: one assumption per dependency, each "succeeds", "fails", or "ignored" +// (no claim either way). A path's ID hashes the head and its assumptions, so +// a path *is* its guess; building the same guess again is a new attempt of +// the same path, and (path ID, attempt) names the resulting build. +// +// # The life of a path +// +// A path's status tracks its current attempt: +// +// funded observed running build finished +// (no entry) ─────────► pending ─────────► building ──────┬──► passed +// │ │ └──► failed +// "stop this": │ │ +// broken or ▼ ▼ +// preempted ─────► cancelling ◄────┘ +// │ +// │ build observed stopped +// ▼ +// cancelled +// +// A path is funded — given one of the slots under the queue's cap on +// concurrent builds, the build budget — when its guess is judged worth +// building, and every pending, building, and cancelling path holds its slot +// until its build stops. A path is broken once a dependency's actual result +// proves one of its assumptions wrong: its guess can no longer come true, so +// its build is cancelled to free the slot. +// +// Cancelling is intent, not fact: the build keeps its slot until CI actually +// stops it, and only an observation of that stop moves the path to cancelled. +// The intent needs no dispatch of its own — the poll loop reads it off the set +// and asks the runner to stop the build. A terminal path can be resurrected by +// a new build proposal — status returns to pending and Attempt increments, the +// one backwards step in the diagram. +// +// # How a run works +// +// Every message is only a dirty signal — "this queue changed" — naming the +// batch that changed. The run then re-plans the whole queue from a single +// read; nothing carries over from earlier runs, so duplicated, delayed, or +// reordered signals are harmless, and a later run repairs whatever an +// earlier one left half-done. +// +// signal ──► read ──► cancel ──► ask ──► check ──► dispatch +// one broken the filter save changes, +// read of paths Specu- its hand builds to +// queue + lator proposals the build stage +// paths +// +// The Speculator is the extension that proposes which paths to fund or +// preempt. It only ever proposes: check.go filters its answer, and broken +// paths are cancelled before it is asked, so it reasons over facts as they +// now stand rather than over a picture the run is about to invalidate. +// +// # Ownership +// +// The path set has exactly one writer: this controller's run. The build stage +// starts builds and the buildsignal stage watches and stops them; both write +// only per-build records of their own (the build link and the build's status), +// which read folds into the snapshot, and buildsignal reads the set only as +// the kill list for builds nothing wants any more. One writer is what lets a +// run hold one version of a head's paths across its whole decision without a +// poll invalidating it mid-thought. +package speculate diff --git a/submitqueue/orchestrator/controller/speculate/run.go b/submitqueue/orchestrator/controller/speculate/run.go new file mode 100644 index 00000000..bfcb6db7 --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/run.go @@ -0,0 +1,294 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +// run re-plans a whole queue from a single read of its state, in the five +// steps the package doc lays out: read, cancel broken paths, ask, check, +// dispatch. +// +// The batch on the triggering message only says which queue woke up; nothing +// about the run depends on which batch it was, or on any earlier run. That is +// what makes duplicated, delayed, and reordered signals harmless, and what +// lets a later run repair anything an earlier one left half-done. +// +// Cancelling broken paths before asking is what keeps the Speculator's work +// from being wasted: it reasons over the queue as the facts have already left +// it, rather than over a picture this run is about to invalidate. +func (c *Controller) run(ctx context.Context, queue string) error { + snap, err := c.read(ctx, queue) + if err != nil { + return err + } + if len(snap.speculating) == 0 { + // No head is open to new work, so there is nothing to speculate about. + return nil + } + + c.cancelBrokenPaths(&snap) + + proposals, err := c.ask(ctx, queue, snap) + if err != nil { + return err + } + + kept, rejected := filterProposals(proposals, snap) + for _, reason := range rejected { + metrics.NamedCounter(c.metricsScope, opName, "speculation_rejected", 1, + metrics.NewTag("reason", string(reason))) + c.logger.Warnw("dropped a speculator proposal", + "queue", queue, + "reason", string(reason), + ) + } + + return c.dispatch(ctx, queue, snap, kept) +} + +// read builds the run's snapshot. Batches come first because their dependency +// lists say which finalized batches still have to be resolved, and their IDs +// say which path sets to load. +func (c *Controller) read(ctx context.Context, queue string) (snapshot, error) { + // TODO: this is the third caller of GetByQueueAndStates, which is backed by + // the only secondary index in the schema. It is being replaced by a + // per-queue active-batch aggregate; this call moves with the others. + inFlight, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, queue, entity.ActiveBatchStates()) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return snapshot{}, fmt.Errorf("failed to list in-flight batches of queue %s: %w", queue, err) + } + + snap := snapshot{ + batches: make(map[string]entity.Batch, len(inFlight)), + pathSets: make(map[string]entity.SpeculationPathSet, len(inFlight)), + dirty: make(map[string]bool, len(inFlight)), + } + + for _, batch := range inFlight { + snap.batches[batch.ID] = batch + if batch.State == entity.BatchStateSpeculating { + snap.speculating = append(snap.speculating, batch) + } + } + + // Also load dependencies that already finished: they are no longer in the + // in-flight list, but their final states are what decide whether a path's + // assumptions still hold. + for _, batch := range inFlight { + for _, depID := range batch.Dependencies { + if _, ok := snap.batches[depID]; ok { + continue + } + dep, err := c.store.GetBatchStore().Get(ctx, depID) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return snapshot{}, fmt.Errorf("failed to get dependency batch %s of %s: %w", depID, batch.ID, err) + } + snap.batches[depID] = dep + } + } + + for _, batch := range inFlight { + set, err := c.store.GetSpeculationPathSetStore().Get(ctx, batch.ID) + if err != nil { + if errors.Is(err, storage.ErrNotFound) { + // Nothing has been funded for this head yet. + continue + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return snapshot{}, fmt.Errorf("failed to get path set for batch %s: %w", batch.ID, err) + } + changed, err := c.updatePathsFromBuilds(ctx, &set) + if err != nil { + return snapshot{}, err + } + snap.pathSets[batch.ID] = set + if changed { + snap.markDirty(batch.ID) + } + } + + return snap, nil +} + +// updatePathsFromBuilds updates each live path's status in the in-memory set +// from what its build actually did: a build found running turns a pending +// path into building, and a finished build writes its result (passed, failed, +// cancelled) onto the path. +// +// Paths already in a final status are skipped — their result was written to +// the set by an earlier run — and the live ones are capped by the build +// budget, so this costs a bounded handful of point reads per run, not one per +// path. +// +// The update is applied in memory only. It is persisted by whichever write +// the run makes later, which keeps this controller the path set's single +// writer: the build stages record what CI did on per-build records, and the +// run alone decides when that becomes the path's status. +func (c *Controller) updatePathsFromBuilds(ctx context.Context, set *entity.SpeculationPathSet) (bool, error) { + changed := false + for i := range set.Paths { + entry := &set.Paths[i] + if entry.Status.IsTerminal() { + continue + } + + link, err := c.store.GetPathBuildStore().Get(ctx, entry.ID, entry.Attempt) + if errors.Is(err, storage.ErrNotFound) { //nolint:gocritic // reads better than a switch here + // No build is recorded for this attempt. Usually nothing was ever + // dispatched; at worst a dispatch is mid-flight, and a build that + // materializes after this read still gets its signal, so the poll + // loop finds it and stops it if its path no longer wants it. + // + // A pending path keeps the run's own intent: its dispatch is + // re-sent rather than abandoned. + continue + } + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return false, fmt.Errorf("failed to look up build for path %s attempt %d: %w", entry.ID, entry.Attempt, err) + } + + build, err := c.store.GetBuildStore().Get(ctx, link.BuildID) + if err != nil { + if errors.Is(err, storage.ErrNotFound) { + // Unreachable in a consistent store: the dispatch writes the + // build before the link, so a link always has one. Treated as + // not started rather than fatal, because the alternative is + // wedging a whole queue's run on one corrupt row. + continue + } + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return false, fmt.Errorf("failed to get build %s for path %s: %w", link.BuildID, entry.ID, err) + } + + if !build.Status.IsTerminal() { + // The build is with the runner, so a pending path is now building. + // Recording that is what stops the dispatch: hasActionablePaths + // counts pending as work to send, so a path left pending would be + // re-dispatched on every run for a build that is already running. + // + // A cancelling path stays cancelling. The stop is recorded in the + // set and the poll loop enacts it against the runner; the path keeps + // that status until a later run sees its build actually stopped, so + // a build still running must not overwrite the intent. + if entry.Status == entity.SpeculationPathStatusPending { + entry.Status = entity.SpeculationPathStatusBuilding + changed = true + } + continue + } + + result := terminalPathStatus(build.Status) + entry.Status = result + changed = true + metrics.NamedCounter(c.metricsScope, opName, "path_outcome_recorded", 1, + metrics.NewTag("status", string(result))) + } + return changed, nil +} + +// terminalPathStatus maps a finished build onto the path status that records +// its outcome. Callers check the build is terminal first, so the default is +// unreachable; it returns the zero value there because that value means "no +// status", not "still running". +func terminalPathStatus(status entity.BuildStatus) entity.SpeculationPathStatus { + switch status { + case entity.BuildStatusSucceeded: + return entity.SpeculationPathStatusPassed + case entity.BuildStatusFailed: + return entity.SpeculationPathStatusFailed + case entity.BuildStatusCancelled: + return entity.SpeculationPathStatusCancelled + default: + return entity.SpeculationPathStatusUnknown + } +} + +// cancelBrokenPaths marks cancelling, across every head in the snapshot, each +// path with a broken assumption. Such a path can never merge its head, so this +// is a fact, not a choice — and folding it in before the Speculator is asked +// keeps it from proposing work on top of paths that are already dead, which +// check would only throw away. +func (c *Controller) cancelBrokenPaths(snap *snapshot) { + nowMs := time.Now().UnixMilli() + + for _, batch := range snap.speculating { + set, exists := snap.pathSets[batch.ID] + if !exists { + continue + } + if cancelBrokenPathsInSet(&set, *snap, nowMs) { + snap.pathSets[batch.ID] = set + snap.markDirty(batch.ID) + } + } +} + +// cancelBrokenPathsInSet marks cancelling every live path in one set with a +// broken assumption, and reports whether anything changed. Cancelling rather +// than cancelled, because the path's build may still be occupying CI — only +// the signal that sees it stop can call it terminal. +func cancelBrokenPathsInSet(set *entity.SpeculationPathSet, snap snapshot, nowMs int64) bool { + changed := false + for i := range set.Paths { + entry := &set.Paths[i] + if entry.Status.IsTerminal() || entry.Status == entity.SpeculationPathStatusCancelling { + continue + } + if !assumptionBroken(entry.Path, snap) { + continue + } + entry.Status = entity.SpeculationPathStatusCancelling + entry.UpdatedAtMs = nowMs + changed = true + } + return changed +} + +// ask hands the snapshot to the queue's Speculator. Its answer is a proposal, +// not an instruction: check decides what is actually enacted. +func (c *Controller) ask(ctx context.Context, queue string, snap snapshot) ([]entity.Speculation, error) { + spec, err := c.speculators.For(speculator.Config{QueueName: queue}) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "speculator_errors", 1) + return nil, fmt.Errorf("failed to build speculator for queue %s: %w", queue, err) + } + + sets := make([]entity.SpeculationPathSet, 0, len(snap.pathSets)) + for _, batch := range snap.speculating { + if set, ok := snap.pathSets[batch.ID]; ok { + sets = append(sets, set) + } + } + + proposals, err := spec.Speculate(ctx, snap.speculating, sets) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "speculator_errors", 1) + return nil, fmt.Errorf("speculator failed for queue %s: %w", queue, err) + } + return proposals, nil +} diff --git a/submitqueue/orchestrator/controller/speculate/run_test.go b/submitqueue/orchestrator/controller/speculate/run_test.go new file mode 100644 index 00000000..ab42d303 --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/run_test.go @@ -0,0 +1,489 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber-go/tally" + entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" + "github.com/uber/submitqueue/platform/consumer" + queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" + "github.com/uber/submitqueue/submitqueue/core/topickey" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" + "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" +) + +// scriptedSpeculator returns a fixed answer and records what it was asked. +type scriptedSpeculator struct { + proposals []entity.Speculation + err error + + gotBatches []entity.Batch + gotSets []entity.SpeculationPathSet + calls int +} + +func (s *scriptedSpeculator) Speculate(_ context.Context, batches []entity.Batch, sets []entity.SpeculationPathSet) ([]entity.Speculation, error) { + s.calls++ + s.gotBatches = batches + s.gotSets = sets + return s.proposals, s.err +} + +type runHarness struct { + controller *Controller + batches *storagemock.MockBatchStore + pathSets *storagemock.MockSpeculationPathSetStore + pathBuilds *storagemock.MockPathBuildStore + builds *storagemock.MockBuildStore + spec *scriptedSpeculator + published []string +} + +// newRunHarness wires a controller whose queue read returns inFlight. +func newRunHarness(t *testing.T, ctrl *gomock.Controller, spec *scriptedSpeculator, inFlight []entity.Batch) *runHarness { + t.Helper() + + h := &runHarness{spec: spec} + + h.batches = storagemock.NewMockBatchStore(ctrl) + h.batches.EXPECT(). + GetByQueueAndStates(gomock.Any(), "q", entity.ActiveBatchStates()). + Return(inFlight, nil).AnyTimes() + + h.pathSets = storagemock.NewMockSpeculationPathSetStore(ctrl) + h.pathBuilds = storagemock.NewMockPathBuildStore(ctrl) + h.builds = storagemock.NewMockBuildStore(ctrl) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(h.batches).AnyTimes() + store.EXPECT().GetSpeculationPathSetStore().Return(h.pathSets).AnyTimes() + store.EXPECT().GetPathBuildStore().Return(h.pathBuilds).AnyTimes() + store.EXPECT().GetBuildStore().Return(h.builds).AnyTimes() + + pub := queuemock.NewMockPublisher(ctrl) + pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, topic string, msg entityqueue.Message) error { + h.published = append(h.published, topic) + return nil + }, + ).AnyTimes() + q := queuemock.NewMockQueue(ctrl) + q.EXPECT().Publisher().Return(pub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyBuild, Name: "build", Queue: q}, + {Key: topickey.TopicKeyMerge, Name: "submitqueue-merge", Queue: q}, + {Key: topickey.TopicKeyConclude, Name: "conclude", Queue: q}, + }) + require.NoError(t, err) + + h.controller = NewController( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, store, + staticSpeculatorFactory{s: spec}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate", + ) + return h +} + +// noBuildsDispatched makes the build lookup find nothing, i.e. no path has a build yet. +// It is opt-in rather than a harness default because a catch-all registered up +// front would shadow the specific expectations build-lookup tests set. +func (h *runHarness) noBuildsDispatched() { + h.pathBuilds.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()). + Return(entity.PathBuild{}, storage.ErrNotFound).AnyTimes() +} + +func speculatingHead() entity.Batch { + return entity.Batch{ + ID: head, Queue: "q", State: entity.BatchStateSpeculating, + Dependencies: []string{dep1, dep2}, Version: 1, + } +} + +// A funded path is persisted as pending and the head is dispatched to build. +func TestRun_FundsProposedPath(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{proposals: []entity.Speculation{{Path: path, Action: entity.PathActionBuild}}} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{}, storage.ErrNotFound) + + // First funding of a head creates its set at version 1. + h.pathSets.EXPECT().Create(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, set entity.SpeculationPathSet) error { + require.Len(t, set.Paths, 1) + assert.Equal(t, path.ID(), set.Paths[0].ID) + assert.Equal(t, entity.SpeculationPathStatusPending, set.Paths[0].Status) + assert.Equal(t, 1, set.Paths[0].Attempt) + assert.Equal(t, int32(1), set.Version) + return nil + }) + + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Equal(t, []string{"build"}, h.published) +} + +// The Speculator sees the speculating batches and their path sets. +func TestRun_PassesSnapshotToSpeculator(t *testing.T) { + ctrl := gomock.NewController(t) + spec := &scriptedSpeculator{} + + merging := entity.Batch{ID: "q/batch/merging", Queue: "q", State: entity.BatchStateMerging, Version: 1} + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead(), merging}) + 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.BatchStateSpeculating}, nil) + + existing := entity.SpeculationPathSet{Head: head, Version: 3} + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(existing, nil) + h.pathSets.EXPECT().Get(gomock.Any(), merging.ID).Return(entity.SpeculationPathSet{}, storage.ErrNotFound) + + require.NoError(t, h.controller.run(context.Background(), "q")) + + require.Equal(t, 1, spec.calls) + require.Len(t, spec.gotBatches, 1) + assert.Equal(t, head, spec.gotBatches[0].ID, "only speculating heads are action targets") + require.Len(t, spec.gotSets, 1) + assert.Equal(t, int32(3), spec.gotSets[0].Version) +} + +// A resolved dependency that breaks a running path's assumption cancels it, +// without the Speculator being consulted. +func TestRun_CancelsBrokenPath(t *testing.T) { + ctrl := gomock.NewController(t) + broken := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + // dep1 failed, so a path betting on its success can never pass. + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateFailed}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(broken, entity.SpeculationPathStatusBuilding)}, + Version: 2, + }, nil) + + h.pathSets.EXPECT().Update(gomock.Any(), gomock.Any(), int32(2), int32(3)). + DoAndReturn(func(_ context.Context, set entity.SpeculationPathSet, _, _ int32) error { + assert.Equal(t, entity.SpeculationPathStatusCancelling, set.Paths[0].Status) + return nil + }) + + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Empty(t, h.published, + "a cancelling path needs no dispatch; the poll loop reads the stop off the set") +} + +// A proposal the check rejects is dropped without touching storage, and the run +// still succeeds — a misbehaving Speculator must not stall the queue. +func TestRun_DropsRejectedProposal(t *testing.T) { + ctrl := gomock.NewController(t) + malformed := pathOver(entity.DependencyAssumptionSucceeds) // missing dep2 + spec := &scriptedSpeculator{proposals: []entity.Speculation{ + {Path: malformed, Action: entity.PathActionBuild}, + }} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{}, storage.ErrNotFound) + + // No Create and no Update: nothing was accepted. + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Empty(t, h.published) +} + +// Re-proposing a path that is already funded must not start a second build. +func TestRun_AlreadyFundedPathIsNotRefunded(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{proposals: []entity.Speculation{{Path: path, Action: entity.PathActionBuild}}} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(path, entity.SpeculationPathStatusBuilding)}, + Version: 1, + }, nil) + + // No Update: the path keeps the slot and the attempt it already has. + require.NoError(t, h.controller.run(context.Background(), "q")) + + // The head still has no actionable path, so nothing is dispatched. + assert.Empty(t, h.published) +} + +// A pending path is re-dispatched every run until the build stage moves it on, +// so a dispatch lost in flight is eventually re-sent. +func TestRun_RedispatchesPendingPath(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(path, entity.SpeculationPathStatusPending)}, + Version: 1, + }, nil) + + // Nothing changed, so no write — but the dispatch goes out again. + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Equal(t, []string{"build"}, h.published) +} + +// Losing the path-set race is not an error: the head is simply re-planned next +// run, and the rest of the queue is unaffected. +func TestRun_LostCASIsNotAnError(t *testing.T) { + ctrl := gomock.NewController(t) + broken := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateFailed}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(broken, entity.SpeculationPathStatusBuilding)}, + Version: 2, + }, nil) + h.pathSets.EXPECT().Update(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(storage.ErrVersionMismatch) + + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Empty(t, h.published, "a head whose write was lost is not dispatched on stale state") +} + +// A Speculator failure abandons the run; the next dirty signal retries with +// fresh state. +func TestRun_SpeculatorFailure(t *testing.T) { + ctrl := gomock.NewController(t) + spec := &scriptedSpeculator{err: fmt.Errorf("scorer unavailable")} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{}, storage.ErrNotFound) + + require.Error(t, h.controller.run(context.Background(), "q")) +} + +// A queue with nothing speculating never reaches the Speculator. +func TestRun_NoSpeculatingBatches(t *testing.T) { + ctrl := gomock.NewController(t) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, nil) + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Zero(t, spec.calls) +} + +// Verify the harness's speculator satisfies the extension contract. +var _ speculator.Speculator = (*scriptedSpeculator)(nil) + +// Reading build records is how the run learns what CI did without the build stages writing +// the path set. A finished build becomes the path's status in memory, and the +// run's own write is what persists it. +func TestRun_RecordsFinishedBuildsOnPaths(t *testing.T) { + tests := []struct { + name string + build entity.BuildStatus + want entity.SpeculationPathStatus + }{ + {"succeeded becomes passed", entity.BuildStatusSucceeded, entity.SpeculationPathStatusPassed}, + {"failed becomes failed", entity.BuildStatusFailed, entity.SpeculationPathStatusFailed}, + {"cancelled becomes cancelled", entity.BuildStatusCancelled, entity.SpeculationPathStatusCancelled}, + {"running becomes building", entity.BuildStatusRunning, entity.SpeculationPathStatusBuilding}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(path, entity.SpeculationPathStatusPending)}, + Version: 1, + }, nil) + + h.pathBuilds.EXPECT().Get(gomock.Any(), path.ID(), 1). + Return(entity.PathBuild{PathID: path.ID(), Attempt: 1, BuildID: "build-1"}, nil) + h.builds.EXPECT().Get(gomock.Any(), "build-1"). + Return(entity.Build{ID: "build-1", BatchID: head, Status: tt.build}, nil) + + h.pathSets.EXPECT().Update(gomock.Any(), gomock.Any(), int32(1), int32(2)). + DoAndReturn(func(_ context.Context, s entity.SpeculationPathSet, _, _ int32) error { + assert.Equal(t, tt.want, s.Paths[0].Status) + return nil + }) + + require.NoError(t, h.controller.run(context.Background(), "q")) + }) + } +} + +// A path the run wants stopped stays cancelling while its build is still +// running: intent outlives a build still seen running, and only CI actually +// stopping ends it. +func TestRun_BuildUpdateDoesNotOverrideCancellingIntent(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateSpeculating}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(path, entity.SpeculationPathStatusCancelling)}, + Version: 1, + }, nil) + + h.pathBuilds.EXPECT().Get(gomock.Any(), path.ID(), 1). + Return(entity.PathBuild{PathID: path.ID(), Attempt: 1, BuildID: "build-1"}, nil) + h.builds.EXPECT().Get(gomock.Any(), "build-1"). + Return(entity.Build{ID: "build-1", BatchID: head, Status: entity.BuildStatusRunning}, nil) + + // Nothing changed, so no write — and no dispatch either: the poll loop is + // what keeps asking the runner to stop, not the build stage. + require.NoError(t, h.controller.run(context.Background(), "q")) + assert.Empty(t, h.published) +} + +// A finished path is never looked up again: its outcome is already in the set. +func TestRun_DoesNotReReadFinishedPaths(t *testing.T) { + ctrl := gomock.NewController(t) + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + 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.BatchStateFailed}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(path, entity.SpeculationPathStatusPassed)}, + Version: 1, + }, nil) + + // The default pathBuilds expectation is AnyTimes, so assert on the build + // store: a finished entry must not reach it. + require.NoError(t, h.controller.run(context.Background(), "q")) +} + +// Broken paths are cancelled before the Speculator is asked, so it reasons over the +// queue as the facts have already left it. Asking first would have it propose +// work on top of a path this run is about to rule out, which check would only +// throw away. +func TestRun_BrokenPathsAreVisibleToTheSpeculator(t *testing.T) { + ctrl := gomock.NewController(t) + broken := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + spec := &scriptedSpeculator{} + + h := newRunHarness(t, ctrl, spec, []entity.Batch{speculatingHead()}) + h.noBuildsDispatched() + // dep1 failed, so a path assuming it succeeds can never pass. + h.batches.EXPECT().Get(gomock.Any(), dep1).Return(entity.Batch{ID: dep1, State: entity.BatchStateFailed}, nil) + h.batches.EXPECT().Get(gomock.Any(), dep2).Return(entity.Batch{ID: dep2, State: entity.BatchStateSpeculating}, nil) + h.pathSets.EXPECT().Get(gomock.Any(), head).Return(entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{entryFor(broken, entity.SpeculationPathStatusBuilding)}, + Version: 2, + }, nil) + h.pathSets.EXPECT().Update(gomock.Any(), gomock.Any(), int32(2), int32(3)).Return(nil) + + require.NoError(t, h.controller.run(context.Background(), "q")) + + require.Equal(t, 1, spec.calls) + require.Len(t, spec.gotSets, 1) + require.Len(t, spec.gotSets[0].Paths, 1) + assert.Equal(t, entity.SpeculationPathStatusCancelling, spec.gotSets[0].Paths[0].Status, + "the Speculator must see the broken path as already cancelling") +} + +// cancelBrokenPathsInSet marks broken paths cancelling, not cancelled: their builds +// may still be occupying CI, and only the signal that sees them stop can call +// it done. +func TestCancelBrokenPathsInSet(t *testing.T) { + broken := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + intact := pathOver(entity.DependencyAssumptionFails, entity.DependencyAssumptionIgnored) + + set := entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{ + {ID: broken.ID(), Path: broken, Status: entity.SpeculationPathStatusBuilding}, + {ID: intact.ID(), Path: intact, Status: entity.SpeculationPathStatusPending}, + }, + } + + // dep1 failed, so the path assuming it succeeds is broken. + snap := snapWith(entity.BatchStateFailed, entity.BatchStateSpeculating) + + require.True(t, cancelBrokenPathsInSet(&set, snap, 42)) + assert.Equal(t, entity.SpeculationPathStatusCancelling, set.Paths[0].Status) + assert.Equal(t, int64(42), set.Paths[0].UpdatedAtMs) + assert.Equal(t, entity.SpeculationPathStatusPending, set.Paths[1].Status) +} + +// A path whose build already finished is left alone: a recorded outcome is not +// something a later run gets to revise. +func TestCancelBrokenPathsInSet_LeavesFinishedPaths(t *testing.T) { + broken := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionIgnored) + snap := snapWith(entity.BatchStateFailed, entity.BatchStateSpeculating) + + for _, status := range []entity.SpeculationPathStatus{ + entity.SpeculationPathStatusPassed, + entity.SpeculationPathStatusFailed, + entity.SpeculationPathStatusCancelled, + entity.SpeculationPathStatusCancelling, + } { + t.Run(string(status), func(t *testing.T) { + set := entity.SpeculationPathSet{Head: head, Paths: []entity.SpeculationPathEntry{ + {ID: broken.ID(), Path: broken, Status: status}, + }} + assert.False(t, cancelBrokenPathsInSet(&set, snap, 42)) + assert.Equal(t, status, set.Paths[0].Status) + }) + } +} diff --git a/submitqueue/orchestrator/controller/speculate/snapshot.go b/submitqueue/orchestrator/controller/speculate/snapshot.go new file mode 100644 index 00000000..cd439766 --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/snapshot.go @@ -0,0 +1,86 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import "github.com/uber/submitqueue/submitqueue/entity" + +// snapshot is one run's working state: the queue as it was read, plus the +// changes the run has decided on so far. The store is read once — two +// decisions in the same run can never disagree about the state of the world — +// and everything the run concludes is folded back in here before anything +// else is derived from it. +type snapshot struct { + // batches is every batch the run can reason about, by ID: the queue's + // in-flight batches plus any finalized batch still named as a dependency + // of one of them. + batches map[string]entity.Batch + // speculating is the queue's Speculating batches, in queue order. These are + // the heads open to new work: what the Speculator is handed, and what the + // dispatch step walks. + speculating []entity.Batch + // pathSets is each head's in-memory path set, by head batch ID. Statuses + // already reflect what each path's build actually did — see + // (*Controller).updatePathsFromBuilds. + pathSets map[string]entity.SpeculationPathSet + // dirty marks heads whose in-memory set differs from what is stored. + // Always touch it through markDirty / isDirty so every step of the + // handshake is greppable; see those methods for the contract. + dirty map[string]bool +} + +// markDirty records that the head's in-memory path set has diverged from the +// stored copy and must be persisted before the run ends. Call it after every +// mutation of snap.pathSets[id]. +func (s *snapshot) markDirty(id string) { + s.dirty[id] = true +} + +// isDirty reports whether the head's set still needs persisting. A head with +// no entry is not dirty, so the unguarded map read is deliberate. +func (s snapshot) isDirty(id string) bool { + return s.dirty[id] +} + +// batchState returns a batch's state, or BatchStateUnknown for a batch the +// run never read — which resolves no assumption either way. +func (s snapshot) batchState(id string) entity.BatchState { + return s.batches[id].State +} + +// assumptionBroken reports whether a finished dependency has already proven +// one of the path's assumptions wrong: a dependency the path assumed would +// succeed ended some other way, or one it assumed would fail succeeded. A +// dependency still in flight proves nothing either way, and an ignored one +// never does — the path made no claim about it. +func assumptionBroken(path entity.SpeculationPath, snap snapshot) bool { + for _, dep := range path.Dependencies { + state := snap.batchState(dep.Batch) + if !state.IsTerminal() { + continue + } + + switch dep.Assumption { + case entity.DependencyAssumptionSucceeds: + if state != entity.BatchStateSucceeded { + return true + } + case entity.DependencyAssumptionFails: + if state == entity.BatchStateSucceeded { + return true + } + } + } + return false +} diff --git a/submitqueue/orchestrator/controller/speculate/snapshot_test.go b/submitqueue/orchestrator/controller/speculate/snapshot_test.go new file mode 100644 index 00000000..a8c9d68f --- /dev/null +++ b/submitqueue/orchestrator/controller/speculate/snapshot_test.go @@ -0,0 +1,101 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package speculate + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/uber/submitqueue/submitqueue/entity" +) + +const ( + head = "q/batch/head" + dep1 = "q/batch/dep1" + dep2 = "q/batch/dep2" +) + +// pathOver builds a path over the given assumptions, in dep1, dep2 order. +func pathOver(assumptions ...entity.DependencyAssumption) entity.SpeculationPath { + deps := []string{dep1, dep2} + p := entity.SpeculationPath{Head: head} + for i, a := range assumptions { + p.Dependencies = append(p.Dependencies, entity.PathDependency{Batch: deps[i], Assumption: a}) + } + return p +} + +// snapWith builds a snapshot where dep1 and dep2 are in the given states. +func snapWith(dep1State, dep2State entity.BatchState) snapshot { + return snapshot{ + batches: map[string]entity.Batch{ + head: {ID: head, State: entity.BatchStateSpeculating, Dependencies: []string{dep1, dep2}}, + dep1: {ID: dep1, State: dep1State}, + dep2: {ID: dep2, State: dep2State}, + }, + pathSets: map[string]entity.SpeculationPathSet{}, + } +} + +func TestAssumptionBroken(t *testing.T) { + const ( + running = entity.BatchStateSpeculating + succeeded = entity.BatchStateSucceeded + failed = entity.BatchStateFailed + cancelled = entity.BatchStateCancelled + ) + const ( + succeeds = entity.DependencyAssumptionSucceeds + fails = entity.DependencyAssumptionFails + ignored = entity.DependencyAssumptionIgnored + ) + + tests := []struct { + name string + assumption entity.DependencyAssumption + depState entity.BatchState + want bool + }{ + {"succeeds holds while unresolved", succeeds, running, false}, + {"succeeds holds when it succeeds", succeeds, succeeded, false}, + {"succeeds broken when it fails", succeeds, failed, true}, + {"succeeds broken when it is cancelled", succeeds, cancelled, true}, + + {"fails holds while unresolved", fails, running, false}, + {"fails holds when it fails", fails, failed, false}, + {"fails holds when it is cancelled", fails, cancelled, false}, + {"fails broken when it succeeds", fails, succeeded, true}, + + {"ignored survives success", ignored, succeeded, false}, + {"ignored survives failure", ignored, failed, false}, + {"ignored survives cancellation", ignored, cancelled, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + path := pathOver(tt.assumption, entity.DependencyAssumptionIgnored) + assert.Equal(t, tt.want, assumptionBroken(path, snapWith(tt.depState, entity.BatchStateSpeculating))) + }) + } +} + +// A dependency the run never read resolves nothing: the path is still a live +// guess rather than a broken one. +func TestAssumptionBroken_UnknownDependency(t *testing.T) { + path := pathOver(entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails) + snap := snapshot{batches: map[string]entity.Batch{}, pathSets: map[string]entity.SpeculationPathSet{}} + + assert.False(t, assumptionBroken(path, snap)) +} diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index c9f85d41..3f2e8cfb 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -20,43 +20,47 @@ import ( "fmt" "github.com/uber-go/tally" - entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/core/publish" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" "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 message is a dirty signal: it names the batch that changed, but only so +// the queue wakes up. The controller then re-plans that whole queue from a +// single read — see run — asking the Speculator which paths are worth building +// within the budget and cancelling the ones a resolved dependency has ruled +// out. Nothing carries over between runs, so duplicated or reordered signals +// are harmless and a later run repairs whatever an earlier one left half-done. // -// - Created → publish to build, transition to Speculating. -// - Speculating → if all deps are Succeeded, publish to merge and +// Batch verdicts are still the naive per-batch state machine below, which +// advances the triggering batch one step: +// +// - Created → admit to Speculating so the Speculator can act on it. +// - 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). -// - 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. +// - Cancelling → cancel any in-flight Build entity, respeculate +// dependents, CAS to terminal Cancelled, publish to conclude. +// - Merging → no-op (owned by the merge controller). +// - Terminal → re-fan-out to conclude for self-healing in case a +// prior publish was lost. // -// The controller is re-triggered on every relevant downstream event -// (buildsignal, merge), so each call simply re-evaluates the current -// state and either advances or waits. +// Waiting on every dependency is strictly stricter than waiting on the ones a +// passed path assumed, so this is correct while the path-aware finalization +// that replaces it is written — it just does not yet collect the speedup the +// paths are earning. type Controller struct { logger *zap.SugaredLogger metricsScope tally.Scope store storage.Storage + speculators speculator.Factory registry consumer.TopicRegistry topicKey consumer.TopicKey consumerGroup string @@ -65,7 +69,7 @@ type Controller struct { // Verify Controller implements consumer.Controller interface at compile time. var _ consumer.Controller = (*Controller)(nil) -// opName is the metric operation name shared by every emit in this file. +// opName is the metric operation name shared by every emit in this package. const opName = "process" // NewController creates a new speculate controller for the orchestrator. @@ -73,6 +77,7 @@ func NewController( logger *zap.SugaredLogger, scope tally.Scope, store storage.Storage, + speculators speculator.Factory, registry consumer.TopicRegistry, topicKey consumer.TopicKey, consumerGroup string, @@ -81,13 +86,15 @@ func NewController( logger: logger.Named("speculate_controller"), metricsScope: scope.SubScope("speculate_controller"), store: store, + speculators: speculators, registry: registry, topicKey: topicKey, consumerGroup: consumerGroup, } } -// Process advances a batch one step along the naive happy-path. +// Process re-plans the triggering batch's queue, then advances that batch one +// step along the legacy per-batch state machine (see the package doc). // Returns nil to ack (success), or error to nack (retry). func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) error { msg := delivery.Message() @@ -126,47 +133,67 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er return c.fanout(ctx, batch.ID, batch.Queue) } + // A Created batch is admitted before the run so the Speculator can see it: + // proposals may only target Speculating heads. + wasCreated := batch.State == entity.BatchStateCreated + if wasCreated { + admitted, err := c.admit(ctx, batch) + if err != nil { + return err + } + batch = admitted + } + + if err := c.run(ctx, batch.Queue); err != nil { + return err + } + + // A freshly admitted head stops here rather than falling through to the + // finalizer. The run above funded its first paths and nothing has been + // built yet — finalizing on the same message would let a batch with no + // dependencies merge before any build had run. + if wasCreated { + return nil + } + // Merging is owned by the merge controller, which has its own self-heal. if batch.State == entity.BatchStateMerging { metrics.NamedCounter(c.metricsScope, opName, "noop_merging", 1) return nil } - switch batch.State { - case entity.BatchStateCreated: - return c.startSpeculation(ctx, batch) - case entity.BatchStateSpeculating: + if batch.State == 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) } -} - -// startSpeculation kicks off CI for this batch on top of the speculative head -// (batch.Dependencies assumed to all pass), then transitions to Speculating. -func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) error { - c.logger.Infow("starting speculation", - "batch_id", batch.ID, - "speculation_chain", append(append([]string{}, batch.Dependencies...), batch.ID), - ) - if err := c.publish(ctx, topickey.TopicKeyBuild, batch.ID, batch.Queue); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) - return fmt.Errorf("failed to publish to build: %w", err) - } + metrics.NamedCounter(c.metricsScope, opName, "unexpected_state", 1) + return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) +} - // Optimistic CAS: if the version has already advanced (concurrent speculate), - // the next event will see the new state and behave correctly. +// admit moves a batch from Created to Speculating, which is what makes it +// visible to the Speculator as an action target. It returns the batch with the +// new state and version so the caller keeps writing against a current copy. +// +// It no longer dispatches anything itself: which paths to build for this head +// is the run's decision, taken over the whole queue rather than one batch at a +// time. +func (c *Controller) admit(ctx context.Context, batch entity.Batch) (entity.Batch, error) { newVersion := batch.Version + 1 batch.State = entity.BatchStateSpeculating if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update batch %s state to speculating: %w", batch.ID, err) + return batch, fmt.Errorf("failed to update batch %s state to speculating: %w", batch.ID, err) } + batch.Version = newVersion + batch.State = entity.BatchStateSpeculating - metrics.NamedCounter(c.metricsScope, opName, "started_speculation", 1) - return nil + metrics.NamedCounter(c.metricsScope, opName, "admitted", 1) + c.logger.Infow("admitted batch to speculation", + "batch_id", batch.ID, + "queue", batch.Queue, + "dependencies", batch.Dependencies, + ) + return batch, nil } // tryFinalize publishes to merge and transitions to Merging iff every @@ -215,7 +242,7 @@ func (c *Controller) tryFinalize(ctx context.Context, batch entity.Batch) error return nil } - if err := c.publish(ctx, topickey.TopicKeyMerge, batch.ID, batch.Queue); err != nil { + if err := c.publishBatchID(ctx, topickey.TopicKeyMerge, batch.ID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to merge: %w", err) } @@ -251,7 +278,7 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d } batch.Version = newVersion - if err := c.publish(ctx, topickey.TopicKeyConclude, batch.ID, batch.Queue); err != nil { + if err := c.publishBatchID(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) } @@ -319,7 +346,7 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error return err } - if err := c.publish(ctx, topickey.TopicKeyConclude, batch.ID, batch.Queue); err != nil { + if err := c.publishBatchID(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) } @@ -387,7 +414,7 @@ func (c *Controller) respeculateDependents(ctx context.Context, batch entity.Bat // reads, consumer-pool parallelism / backpressure, and the existing // state-machine dispatch in Process all argue for the publish. Revisit // if the extra message hop ever shows up as latency or cost. - if err := c.publish(ctx, topickey.TopicKeySpeculate, depID, batch.Queue); err != nil { + if err := c.publishBatchID(ctx, topickey.TopicKeySpeculate, depID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish dependent batch %s to speculate: %w", depID, err) } @@ -417,38 +444,22 @@ func (c *Controller) fetchDependencies(ctx context.Context, batch entity.Batch) // a terminal state. Used for self-healing when a previous publish was lost: // re-sending to conclude guarantees request-state reconciliation. func (c *Controller) fanout(ctx context.Context, batchID, partitionKey string) error { - if err := c.publish(ctx, topickey.TopicKeyConclude, batchID, partitionKey); err != nil { + if err := c.publishBatchID(ctx, topickey.TopicKeyConclude, batchID, partitionKey); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish to conclude: %w", err) } return nil } -// publish publishes a batch ID to the specified topic key. -func (c *Controller) publish(ctx context.Context, key consumer.TopicKey, batchID string, partitionKey string) error { - bid := entity.BatchID{ID: batchID} - payload, err := bid.ToBytes() +// publishBatchID publishes a batch ID to the topic behind key. The batch ID +// doubles as the message ID, so the queue deduplicates repeat publishes for +// the same batch against rows it has not collected yet. +func (c *Controller) publishBatchID(ctx context.Context, key consumer.TopicKey, batchID string, partitionKey string) error { + payload, err := entity.BatchID{ID: batchID}.ToBytes() if err != nil { return fmt.Errorf("failed to serialize batch ID: %w", err) } - - msg := entityqueue.NewMessage(batchID, payload, partitionKey, nil) - - q, ok := c.registry.Queue(key) - if !ok { - return fmt.Errorf("no queue registered for topic key %s", key) - } - - topicName, ok := c.registry.TopicName(key) - if !ok { - return fmt.Errorf("no topic name registered for topic key %s", key) - } - - if err := q.Publisher().Publish(ctx, topicName, msg); err != nil { - return fmt.Errorf("failed to publish message: %w", err) - } - - return nil + return publish.Message(ctx, c.registry, key, batchID, payload, partitionKey, 0) } // Name returns the controller name for logging and metrics. diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 4a0cdefd..15ce16ab 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate_test.go +++ b/submitqueue/orchestrator/controller/speculate/speculate_test.go @@ -28,6 +28,7 @@ import ( queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" "github.com/uber/submitqueue/submitqueue/extension/storage" storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" "go.uber.org/mock/gomock" @@ -39,6 +40,31 @@ func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch { return batch } +// quietSpeculator proposes nothing, which is what tests focused on the verdict +// state machine want: the run happens but changes no paths. +type quietSpeculator struct{} + +func (quietSpeculator) Speculate(context.Context, []entity.Batch, []entity.SpeculationPathSet) ([]entity.Speculation, error) { + return nil, nil +} + +// staticSpeculatorFactory returns a fixed Speculator for any queue. +type staticSpeculatorFactory struct{ s speculator.Speculator } + +func (f staticSpeculatorFactory) For(speculator.Config) (speculator.Speculator, error) { + return f.s, nil +} + +// stubQuietRun makes the speculation run a no-op: the queue lists no in-flight +// batches, so the run returns before reading any path set or asking the +// Speculator. Tests below exercise the verdict state machine, which the run is +// deliberately independent of; run_test.go covers the run itself. +func stubQuietRun(batchStore *storagemock.MockBatchStore) { + batchStore.EXPECT(). + GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, nil).AnyTimes() +} + // 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() @@ -84,7 +110,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock ) require.NoError(t, err) - return NewController(logger, scope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + return NewController(logger, scope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") } // runProcess builds a delivery for batchID and invokes Process once. @@ -128,6 +154,7 @@ func TestController_Process_StartSpeculation(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -146,6 +173,7 @@ func TestController_Process_FinalizeNoDeps(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -166,6 +194,7 @@ func TestController_Process_FinalizeAllDepsSucceeded(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -184,6 +213,7 @@ func TestController_Process_WaitingOnDep(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -204,6 +234,7 @@ func TestController_Process_FailedDepFailsBatch(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -226,6 +257,7 @@ func TestController_Process_CancelledDepSkipped(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -242,6 +274,7 @@ func TestController_Process_MergingNoOp(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -266,6 +299,7 @@ func TestController_Process_TerminalSelfHeals(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) // Require exactly one publish to the conclude topic for self-healing. mockPub := queuemock.NewMockPublisher(ctrl) @@ -282,7 +316,7 @@ func TestController_Process_TerminalSelfHeals(t *testing.T) { require.NoError(t, err) logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + controller := NewController(logger, tally.NoopScope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) }) @@ -310,6 +344,7 @@ func TestController_Process_CancelledTerminalSelfHealsDependents(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() // BuildStore must NOT be touched on the terminal self-heal path. @@ -337,7 +372,7 @@ func TestController_Process_CancelledTerminalSelfHealsDependents(t *testing.T) { require.NoError(t, err) logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + controller := NewController(logger, tally.NoopScope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -381,6 +416,7 @@ func TestController_Process_CancellingTerminalFlow(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -408,7 +444,7 @@ func TestController_Process_CancellingTerminalFlow(t *testing.T) { require.NoError(t, err) logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + controller := NewController(logger, tally.NoopScope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) @@ -444,6 +480,7 @@ func TestController_Process_CancellingBuildAlreadyTerminal(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -473,6 +510,7 @@ func TestController_Process_CancellingNoBuildYet(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -500,6 +538,7 @@ func TestController_Process_CancellingNoDependents(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -517,7 +556,7 @@ func TestController_Process_CancellingNoDependents(t *testing.T) { require.NoError(t, err) logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + controller := NewController(logger, tally.NoopScope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") require.NoError(t, runProcess(t, ctrl, controller, batch.ID)) } @@ -541,6 +580,7 @@ func TestController_Process_CancellingTerminalCASVersionMismatch(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() // BatchDependentStore must NOT be touched — terminal CAS failed before fan-out. @@ -558,7 +598,7 @@ func TestController_Process_CancellingTerminalCASVersionMismatch(t *testing.T) { require.NoError(t, err) logger := zaptest.NewLogger(t).Sugar() - controller := NewController(logger, tally.NoopScope, store, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") + controller := NewController(logger, tally.NoopScope, store, staticSpeculatorFactory{s: quietSpeculator{}}, registry, topickey.TopicKeySpeculate, "orchestrator-speculate") err = runProcess(t, ctrl, controller, batch.ID) require.Error(t, err) @@ -576,6 +616,7 @@ func TestController_Process_UnrecognizedState(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) require.Error(t, runProcess(t, ctrl, controller, batch.ID)) @@ -591,6 +632,7 @@ func TestController_Process_StorageFailure(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, nil) err := runProcess(t, ctrl, controller, "test-queue/batch/1") @@ -599,9 +641,11 @@ func TestController_Process_StorageFailure(t *testing.T) { } // Publish failure must not advance the batch state. +// A failed merge publish must abort before the batch is moved to Merging: +// a batch recorded as merging that Runway was never told about would stall. func TestController_Process_PublishFailure(t *testing.T) { ctrl := gomock.NewController(t) - batch := testBatch(entity.BatchStateCreated) + batch := testBatch(entity.BatchStateSpeculating) batchStore := storagemock.NewMockBatchStore(ctrl) batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) @@ -609,6 +653,7 @@ func TestController_Process_PublishFailure(t *testing.T) { store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + stubQuietRun(batchStore) controller := newTestController(t, ctrl, store, fmt.Errorf("publish failed")) require.Error(t, runProcess(t, ctrl, controller, batch.ID)) diff --git a/submitqueue/orchestrator/pipeline.go b/submitqueue/orchestrator/pipeline.go index 38b39033..a1230f0f 100644 --- a/submitqueue/orchestrator/pipeline.go +++ b/submitqueue/orchestrator/pipeline.go @@ -27,6 +27,7 @@ import ( "github.com/uber/submitqueue/submitqueue/extension/buildrunner" "github.com/uber/submitqueue/submitqueue/extension/changeprovider" "github.com/uber/submitqueue/submitqueue/extension/conflict" + "github.com/uber/submitqueue/submitqueue/extension/speculation/speculator" "github.com/uber/submitqueue/submitqueue/extension/storage" "github.com/uber/submitqueue/submitqueue/extension/validator" "github.com/uber/submitqueue/submitqueue/orchestrator/controller" @@ -71,6 +72,9 @@ type Deps struct { // Analyzer resolves the conflict analyzer for each queue. Analyzer conflict.Factory + // Speculator resolves the speculator for each queue. + Speculator speculator.Factory + // Validator resolves the validator for each queue. Validator validator.Factory } @@ -146,7 +150,7 @@ var Stages = []pipeline.Stage[Deps]{ Name: "speculate", ConsumerGroup: "orchestrator", New: func(d Deps, sc pipeline.StageContext) (consumer.Controller, error) { - return speculate.NewController(d.Logger, d.Scope, d.Storage, sc.Registry, sc.TopicKey, sc.ConsumerGroup), nil + return speculate.NewController(d.Logger, d.Scope, d.Storage, d.Speculator, sc.Registry, sc.TopicKey, sc.ConsumerGroup), nil }, DLQ: func(d Deps, sc pipeline.StageContext) (consumer.Controller, error) { return dlq.NewDLQBatchController(d.Logger, d.Scope, d.Storage, sc.Registry, sc.TopicKey, sc.ConsumerGroup), nil