From 50877962f6d4f30f9ae34f5256423df428893fdc Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Tue, 9 Jun 2026 12:34:46 -0700 Subject: [PATCH] feat(speculation): add enumerator and selector extensions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the two pluggable seams from the speculation RFC, as vendor-agnostic extension interfaces under submitqueue/extension/speculation/. enumerator: given a batch and its dependency batches (carrying per-batch Score), mechanically lists the candidate Base/Head paths and scores each — pure, deterministic, status-free. selector: given a speculation tree with controller-stamped status, returns a per-path action (Build/Cancel) — the policy seam; reads status, emits actions, never writes status. Each follows the repo extension contract (conflict.Analyzer reference shape): Factory.For(Config) (T, error) with Config carrying only QueueName; behavioral knobs are integrator-injected at construction. Includes READMEs, gomock packages, and Makefile mock-gen wiring. Interfaces only; concrete impls and controller wiring are deferred. --- Makefile | 2 +- .../speculation/enumerator/BUILD.bazel | 9 ++ .../speculation/enumerator/README.md | 19 ++++ .../speculation/enumerator/enumerator.go | 65 +++++++++++++ .../speculation/enumerator/mock/BUILD.bazel | 13 +++ .../enumerator/mock/enumerator_mock.go | 97 +++++++++++++++++++ .../speculation/selector/BUILD.bazel | 9 ++ .../extension/speculation/selector/README.md | 19 ++++ .../speculation/selector/mock/BUILD.bazel | 13 +++ .../selector/mock/selector_mock.go | 97 +++++++++++++++++++ .../speculation/selector/selector.go | 60 ++++++++++++ 11 files changed, 402 insertions(+), 1 deletion(-) create mode 100644 submitqueue/extension/speculation/enumerator/BUILD.bazel create mode 100644 submitqueue/extension/speculation/enumerator/README.md create mode 100644 submitqueue/extension/speculation/enumerator/enumerator.go create mode 100644 submitqueue/extension/speculation/enumerator/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/enumerator/mock/enumerator_mock.go create mode 100644 submitqueue/extension/speculation/selector/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selector/README.md create mode 100644 submitqueue/extension/speculation/selector/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selector/mock/selector_mock.go create mode 100644 submitqueue/extension/speculation/selector/selector.go diff --git a/Makefile b/Makefile index af3b8aa3..d96e158d 100644 --- a/Makefile +++ b/Makefile @@ -336,7 +336,7 @@ local-stovepipe-gateway-start: build-stovepipe-gateway-linux ## Start Stovepipe mocks: ## Generate mock files using mockgen @echo "Generating mocks..." - @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./extension/counter/... ./extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/core/consumer/... ./submitqueue/core/changeset/... + @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./extension/counter/... ./extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/extension/speculation/enumerator/... ./submitqueue/extension/speculation/selector/... ./submitqueue/core/consumer/... ./submitqueue/core/changeset/... @echo "Mocks generated successfully!" proto: ## Generate protobuf files from .proto definitions diff --git a/submitqueue/extension/speculation/enumerator/BUILD.bazel b/submitqueue/extension/speculation/enumerator/BUILD.bazel new file mode 100644 index 00000000..00f293f2 --- /dev/null +++ b/submitqueue/extension/speculation/enumerator/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "enumerator", + srcs = ["enumerator.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity"], +) diff --git a/submitqueue/extension/speculation/enumerator/README.md b/submitqueue/extension/speculation/enumerator/README.md new file mode 100644 index 00000000..a435cfad --- /dev/null +++ b/submitqueue/extension/speculation/enumerator/README.md @@ -0,0 +1,19 @@ +# Speculation Tree Enumerator + +Vendor-agnostic interface for enumerating the **speculation tree** of a batch — the set of candidate speculation paths the orchestrator may build, each scored with its predicted probability of success. + +See the [Speculation RFC](../../../../doc/rfc/submitqueue/speculation.md) for the end-to-end design and how enumeration fits into the orchestrator pipeline. + +## Enumerator + +An enumerator is deliberately **dumb**: *given a batch and its dependency batches, it mechanically lists the candidate paths and scores them.* It does **not** decide which paths to build — that is the [selector](../selector)'s job — it does **not** set path status, and it does **not** decide how far back to speculate. Speculation depth is the controller's responsibility: the controller trims the dependency list before calling the enumerator, which then enumerates over exactly the list it is handed. + +Each candidate is a path: an assumed-good prefix of predecessor batches (the base) on top of which the batch under verification (the head) is built. The base maps directly onto the build stage's base changes and the head onto the changes being validated. + +Enumeration is **pure and deterministic**: the same batch and dependency list always produce the same tree. This lets the controller regenerate a tree whenever the dependency graph changes without tracking incremental state in the enumerator. Keeping enumeration tractable for a very wide dependency list is the enumerator's only real concern. + +Scores ride in on the inputs. Each dependency is passed as a full `entity.Batch`, which already carries its per-batch success probability (`Batch.Score`) from the score stage; the enumerator combines the scores of a path's base batches into the path's score. No separate scoring backend or injected probability source is needed, and tests just set `.Score` on literal batches. The head is passed as an ID — its score is constant across all of its own paths. + +## Factory + +`Factory.For(Config) (Enumerator, error)` returns the enumerator for a queue, following the repo's extension contract (`conflict.Analyzer` is the reference shape). `Config` carries only the queue identity (`QueueName`); the system hands the factory nothing else. Everything an implementation needs — including behavioral knobs like speculation depth — is injected at construction by the integrator in the wiring layer, which resolves per-queue settings through `queueconfig`. `Enumerate` itself stays config-free. diff --git a/submitqueue/extension/speculation/enumerator/enumerator.go b/submitqueue/extension/speculation/enumerator/enumerator.go new file mode 100644 index 00000000..beb35587 --- /dev/null +++ b/submitqueue/extension/speculation/enumerator/enumerator.go @@ -0,0 +1,65 @@ +// 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 enumerator + +//go:generate mockgen -source=enumerator.go -destination=mock/enumerator_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// Enumerator builds the speculation tree for a batch: the set of candidate +// speculation paths to consider, each scored with its predicted success +// probability. +// +// Enumeration answers "what futures are possible" for a batch. It is +// deliberately dumb: it mechanically lists candidate paths from the dependency +// batches it is handed and attaches a Score to each. It does not decide which +// paths to build — that is the selector's job (see +// extension/speculation/selector) — and it does not decide how far back to +// speculate: the controller trims the dependency list by speculation depth +// before calling Enumerate. +type Enumerator interface { + // Enumerate returns the speculation tree for the batch identified by batchID, + // given its dependency batches in arrival order. Each returned path carries a + // Base/Head split and a predicted success Score; the returned paths leave + // Status unset (the controller stamps it on persist). + // + // Path scores are derived from the dependency batches' Score field (the + // per-batch success probability set by the score stage), so no separate + // scoring backend is needed. The combination formula is the implementation's + // concern. + // + // Enumeration is pure and deterministic: the same (batchID, deps) always + // yields the same tree, so callers may regenerate safely. + Enumerate(ctx context.Context, batchID string, deps []entity.Batch) (entity.SpeculationTree, error) +} + +// Config carries the per-queue identity handed to a Factory. The system knows +// only the queue name; everything an implementation needs (including behavioral +// knobs such as speculation depth) is injected at construction by the integrator. +type Config struct { + // QueueName identifies the queue this Enumerator serves. + QueueName string +} + +// Factory builds the Enumerator for a queue. Implementations are provided by +// integrators (and tests) and inject whatever they need at construction. +type Factory interface { + // For returns the Enumerator for the given queue. + For(cfg Config) (Enumerator, error) +} diff --git a/submitqueue/extension/speculation/enumerator/mock/BUILD.bazel b/submitqueue/extension/speculation/enumerator/mock/BUILD.bazel new file mode 100644 index 00000000..a44e299a --- /dev/null +++ b/submitqueue/extension/speculation/enumerator/mock/BUILD.bazel @@ -0,0 +1,13 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "mock", + srcs = ["enumerator_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity", + "//submitqueue/extension/speculation/enumerator", + "@org_uber_go_mock//gomock", + ], +) diff --git a/submitqueue/extension/speculation/enumerator/mock/enumerator_mock.go b/submitqueue/extension/speculation/enumerator/mock/enumerator_mock.go new file mode 100644 index 00000000..6ad1cc75 --- /dev/null +++ b/submitqueue/extension/speculation/enumerator/mock/enumerator_mock.go @@ -0,0 +1,97 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: enumerator.go +// +// Generated by this command: +// +// mockgen -source=enumerator.go -destination=mock/enumerator_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + entity "github.com/uber/submitqueue/submitqueue/entity" + enumerator "github.com/uber/submitqueue/submitqueue/extension/speculation/enumerator" + gomock "go.uber.org/mock/gomock" +) + +// MockEnumerator is a mock of Enumerator interface. +type MockEnumerator struct { + ctrl *gomock.Controller + recorder *MockEnumeratorMockRecorder + isgomock struct{} +} + +// MockEnumeratorMockRecorder is the mock recorder for MockEnumerator. +type MockEnumeratorMockRecorder struct { + mock *MockEnumerator +} + +// NewMockEnumerator creates a new mock instance. +func NewMockEnumerator(ctrl *gomock.Controller) *MockEnumerator { + mock := &MockEnumerator{ctrl: ctrl} + mock.recorder = &MockEnumeratorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockEnumerator) EXPECT() *MockEnumeratorMockRecorder { + return m.recorder +} + +// Enumerate mocks base method. +func (m *MockEnumerator) Enumerate(ctx context.Context, batchID string, deps []entity.Batch) (entity.SpeculationTree, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Enumerate", ctx, batchID, deps) + ret0, _ := ret[0].(entity.SpeculationTree) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Enumerate indicates an expected call of Enumerate. +func (mr *MockEnumeratorMockRecorder) Enumerate(ctx, batchID, deps any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Enumerate", reflect.TypeOf((*MockEnumerator)(nil).Enumerate), ctx, batchID, deps) +} + +// MockFactory is a mock of Factory interface. +type MockFactory struct { + ctrl *gomock.Controller + recorder *MockFactoryMockRecorder + isgomock struct{} +} + +// MockFactoryMockRecorder is the mock recorder for MockFactory. +type MockFactoryMockRecorder struct { + mock *MockFactory +} + +// NewMockFactory creates a new mock instance. +func NewMockFactory(ctrl *gomock.Controller) *MockFactory { + mock := &MockFactory{ctrl: ctrl} + mock.recorder = &MockFactoryMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockFactory) EXPECT() *MockFactoryMockRecorder { + return m.recorder +} + +// For mocks base method. +func (m *MockFactory) For(cfg enumerator.Config) (enumerator.Enumerator, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "For", cfg) + ret0, _ := ret[0].(enumerator.Enumerator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// For indicates an expected call of For. +func (mr *MockFactoryMockRecorder) For(cfg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "For", reflect.TypeOf((*MockFactory)(nil).For), cfg) +} diff --git a/submitqueue/extension/speculation/selector/BUILD.bazel b/submitqueue/extension/speculation/selector/BUILD.bazel new file mode 100644 index 00000000..90b25fec --- /dev/null +++ b/submitqueue/extension/speculation/selector/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "selector", + srcs = ["selector.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selector", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity"], +) diff --git a/submitqueue/extension/speculation/selector/README.md b/submitqueue/extension/speculation/selector/README.md new file mode 100644 index 00000000..1b304326 --- /dev/null +++ b/submitqueue/extension/speculation/selector/README.md @@ -0,0 +1,19 @@ +# Speculation Path Selector + +Vendor-agnostic interface for deciding what the orchestrator should do with each path in a batch's enumerated speculation tree. + +See the [Speculation RFC](../../../../doc/rfc/submitqueue/speculation.md) for the end-to-end design and how selection fits into the orchestrator pipeline. + +## Selector + +A selector is the **policy** — the part that decides how aggressively to spend build resources. *Given the candidate paths in a tree and their current status, what should we do with each, right now?* It consumes the tree produced by an [enumerator](../enumerator) and returns an **action** per path — `Build` or `Cancel`. Strategies span a spectrum: build only the single optimistic path (cheapest — bet on the happy case), build every candidate (maximum parallelism, maximum build cost), or a top-K / budget-bounded subset in between. + +The selector decides only where to spend build resources. It does **not** decide merging: a path becomes mergeable when its build passed and its base matches what actually landed, which is deterministic, not a policy choice — so the controller finalizes it on its own. + +The selector's only output is actions; it **never** writes status. The controller owns every status write into the store — it reconciles each path's status (candidate, building, passed, failed, cancelled) from the latest builds and dependency states, then feeds the up-to-date tree back in. So the tree is the selector's **complete input**: it never reads storage, builds, or scores directly. This keeps it a pure, deterministic policy that is trivial to test against a literal tree. + +Because it is re-run on every build signal, a selector can start narrow — build the optimistic path first — and widen later, committing more paths only once earlier bets resolve. Returning no action for a path leaves it as-is. Policy parameters — a top-K cap, a build budget, an experiment toggle — are configured when the selector is constructed rather than passed through this contract. + +## Factory + +`Factory.For(Config) (Selector, error)` returns the selector for a queue, following the repo's extension contract (`conflict.Analyzer` is the reference shape). `Config` carries only the queue identity (`QueueName`); the system hands the factory nothing else. Policy knobs — a top-K cap, a build budget, an experiment toggle — are injected at construction by the integrator in the wiring layer, which resolves per-queue settings through `queueconfig`. `Select` itself stays config-free. diff --git a/submitqueue/extension/speculation/selector/mock/BUILD.bazel b/submitqueue/extension/speculation/selector/mock/BUILD.bazel new file mode 100644 index 00000000..4b2a5546 --- /dev/null +++ b/submitqueue/extension/speculation/selector/mock/BUILD.bazel @@ -0,0 +1,13 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "mock", + srcs = ["selector_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity", + "//submitqueue/extension/speculation/selector", + "@org_uber_go_mock//gomock", + ], +) diff --git a/submitqueue/extension/speculation/selector/mock/selector_mock.go b/submitqueue/extension/speculation/selector/mock/selector_mock.go new file mode 100644 index 00000000..755c1dbc --- /dev/null +++ b/submitqueue/extension/speculation/selector/mock/selector_mock.go @@ -0,0 +1,97 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: selector.go +// +// Generated by this command: +// +// mockgen -source=selector.go -destination=mock/selector_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + entity "github.com/uber/submitqueue/submitqueue/entity" + selector "github.com/uber/submitqueue/submitqueue/extension/speculation/selector" + gomock "go.uber.org/mock/gomock" +) + +// MockSelector is a mock of Selector interface. +type MockSelector struct { + ctrl *gomock.Controller + recorder *MockSelectorMockRecorder + isgomock struct{} +} + +// MockSelectorMockRecorder is the mock recorder for MockSelector. +type MockSelectorMockRecorder struct { + mock *MockSelector +} + +// NewMockSelector creates a new mock instance. +func NewMockSelector(ctrl *gomock.Controller) *MockSelector { + mock := &MockSelector{ctrl: ctrl} + mock.recorder = &MockSelectorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSelector) EXPECT() *MockSelectorMockRecorder { + return m.recorder +} + +// Select mocks base method. +func (m *MockSelector) Select(ctx context.Context, tree entity.SpeculationTree) ([]entity.SpeculationPathDecision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Select", ctx, tree) + ret0, _ := ret[0].([]entity.SpeculationPathDecision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Select indicates an expected call of Select. +func (mr *MockSelectorMockRecorder) Select(ctx, tree any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Select", reflect.TypeOf((*MockSelector)(nil).Select), ctx, tree) +} + +// MockFactory is a mock of Factory interface. +type MockFactory struct { + ctrl *gomock.Controller + recorder *MockFactoryMockRecorder + isgomock struct{} +} + +// MockFactoryMockRecorder is the mock recorder for MockFactory. +type MockFactoryMockRecorder struct { + mock *MockFactory +} + +// NewMockFactory creates a new mock instance. +func NewMockFactory(ctrl *gomock.Controller) *MockFactory { + mock := &MockFactory{ctrl: ctrl} + mock.recorder = &MockFactoryMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockFactory) EXPECT() *MockFactoryMockRecorder { + return m.recorder +} + +// For mocks base method. +func (m *MockFactory) For(cfg selector.Config) (selector.Selector, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "For", cfg) + ret0, _ := ret[0].(selector.Selector) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// For indicates an expected call of For. +func (mr *MockFactoryMockRecorder) For(cfg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "For", reflect.TypeOf((*MockFactory)(nil).For), cfg) +} diff --git a/submitqueue/extension/speculation/selector/selector.go b/submitqueue/extension/speculation/selector/selector.go new file mode 100644 index 00000000..6190268b --- /dev/null +++ b/submitqueue/extension/speculation/selector/selector.go @@ -0,0 +1,60 @@ +// 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 selector + +//go:generate mockgen -source=selector.go -destination=mock/selector_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// Selector decides what the controller should do with each path in a batch's +// speculation tree. +// +// Selection is the policy: it answers "which futures do we spend build resources +// on, and how many, right now". It reads the tree — including each path's controller-stamped +// Status (Candidate / Building / Passed / Failed / Cancelled) and Score — and +// returns an action per path it wants to act on. +// +// The selector's only output is actions; it never writes Status. The controller +// owns every Status write (into the store) and feeds the up-to-date tree back in +// on the next call, so the tree is the selector's complete input. This keeps the +// selector a pure, deterministic policy. Policy knobs such as a top-K limit or +// budget belong to the implementation's construction, not this method. +type Selector interface { + // Select returns the actions to take for the given tree. Returning multiple + // Build decisions dispatches several speculative builds in parallel; an empty + // result means nothing should be done right now. Paths the selector has no + // opinion on are simply omitted (leave-as-is). + Select(ctx context.Context, tree entity.SpeculationTree) ([]entity.SpeculationPathDecision, error) +} + +// Config carries the per-queue identity handed to a Factory. The system knows +// only the queue name; everything an implementation needs (including policy +// knobs such as a top-K cap or build budget) is injected at construction by the +// integrator. +type Config struct { + // QueueName identifies the queue this Selector serves. + QueueName string +} + +// Factory builds the Selector for a queue. Implementations are provided by +// integrators (and tests) and inject whatever they need at construction. +type Factory interface { + // For returns the Selector for the given queue. + For(cfg Config) (Selector, error) +}