From 0ae2849bd97f1fc5ba64d8697ada5960ef0052f6 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Wed, 8 Jul 2026 10:15:21 -0700 Subject: [PATCH] feat(speculation): add selector + selection-limit extensions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the selector seam and its limit counterpart from the speculation RFC, as vendor-agnostic extension interfaces under submitqueue/extension/speculation/. selector: the controller hands it the batch's speculation tree; it returns a per-path decision (Promote/Cancel) for the paths it chooses to act on. It is the policy seam — it reads the tree it is given and emits decisions only, never writing status, and does not decide merging. The controller maps each decision to a guarded status transition (Promote → Selected, Cancel → Cancelling/Cancelled) and persists it, staying the single writer of tree state. selectionlimit: the "how much" policy bounding how many paths a batch may build in parallel. It is the selector's companion — the selector decides which paths, the limit decides how many. Injected into the selector at construction and called by it, keeping the selector interface limit-free; the value is dynamic, not a fixed constant. Each follows the repo extension contract: Factory.For(Config) (T, error) with Config carrying only QueueName. Includes READMEs, gomock packages, and programmable fakes. Interfaces only; concrete impls and controller wiring are deferred. --- Makefile | 2 +- .../speculation/selectionlimit/BUILD.bazel | 8 ++ .../speculation/selectionlimit/README.md | 17 ++++ .../selectionlimit/fake/BUILD.bazel | 19 ++++ .../speculation/selectionlimit/fake/fake.go | 52 ++++++++++ .../selectionlimit/fake/fake_test.go | 36 +++++++ .../selectionlimit/mock/BUILD.bazel | 12 +++ .../mock/selectionlimit_mock.go | 96 ++++++++++++++++++ .../selectionlimit/selectionlimit.go | 57 +++++++++++ .../speculation/selector/BUILD.bazel | 9 ++ .../extension/speculation/selector/README.md | 19 ++++ .../speculation/selector/fake/BUILD.bazel | 23 +++++ .../speculation/selector/fake/fake.go | 61 ++++++++++++ .../speculation/selector/fake/fake_test.go | 46 +++++++++ .../speculation/selector/mock/BUILD.bazel | 13 +++ .../selector/mock/selector_mock.go | 97 +++++++++++++++++++ .../speculation/selector/selector.go | 64 ++++++++++++ 17 files changed, 630 insertions(+), 1 deletion(-) create mode 100644 submitqueue/extension/speculation/selectionlimit/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selectionlimit/README.md create mode 100644 submitqueue/extension/speculation/selectionlimit/fake/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selectionlimit/fake/fake.go create mode 100644 submitqueue/extension/speculation/selectionlimit/fake/fake_test.go create mode 100644 submitqueue/extension/speculation/selectionlimit/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selectionlimit/mock/selectionlimit_mock.go create mode 100644 submitqueue/extension/speculation/selectionlimit/selectionlimit.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/fake/BUILD.bazel create mode 100644 submitqueue/extension/speculation/selector/fake/fake.go create mode 100644 submitqueue/extension/speculation/selector/fake/fake_test.go 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 2d0e6e39..328d00ec 100644 --- a/Makefile +++ b/Makefile @@ -364,7 +364,7 @@ local-stovepipe-stop: ## Stop the Stovepipe service mocks: ## Generate mock files using mockgen @echo "Generating mocks..." - @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./platform/extension/counter/... ./platform/extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/extension/speculation/enumerator/... ./submitqueue/extension/speculation/dependencylimit/... ./submitqueue/extension/speculation/scorer/... ./platform/consumer/... ./submitqueue/core/changeset/... ./stovepipe/extension/storage/... ./stovepipe/extension/sourcecontrol/... + @$(BAZEL) run @rules_go//go -- generate ./submitqueue/extension/storage/... ./submitqueue/extension/buildrunner/... ./submitqueue/extension/changeprovider/... ./platform/extension/counter/... ./platform/extension/messagequeue/... ./submitqueue/extension/queueconfig/... ./submitqueue/extension/mergechecker/... ./submitqueue/extension/pusher/... ./submitqueue/extension/scorer/... ./submitqueue/extension/conflict/... ./submitqueue/extension/speculation/enumerator/... ./submitqueue/extension/speculation/dependencylimit/... ./submitqueue/extension/speculation/scorer/... ./submitqueue/extension/speculation/selector/... ./submitqueue/extension/speculation/selectionlimit/... ./platform/consumer/... ./submitqueue/core/changeset/... ./stovepipe/extension/storage/... ./stovepipe/extension/sourcecontrol/... @echo "Mocks generated successfully!" proto: ## Generate protobuf files from .proto definitions diff --git a/submitqueue/extension/speculation/selectionlimit/BUILD.bazel b/submitqueue/extension/speculation/selectionlimit/BUILD.bazel new file mode 100644 index 00000000..8dcc050e --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/BUILD.bazel @@ -0,0 +1,8 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["selectionlimit.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selectionlimit", + visibility = ["//visibility:public"], +) diff --git a/submitqueue/extension/speculation/selectionlimit/README.md b/submitqueue/extension/speculation/selectionlimit/README.md new file mode 100644 index 00000000..15e56f98 --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/README.md @@ -0,0 +1,17 @@ +# Speculation Selection Limit + +Vendor-agnostic "how much" policy that bounds how many paths a batch may build in parallel. + +See the [Speculation RFC](../../../../doc/rfc/submitqueue/speculation.md) for the end-to-end design and how limits fit into the two-layer speculation model. + +## Selection Limit + +The selection limit is the [selector](../selector)'s companion. The selector decides **which** of a batch's paths are worth building — its ranking over the tree; the selection limit decides **how many** of them may run at once. Keeping "which" and "how much" separate keeps selector logic free of resource accounting and lets the bound scale with build resources without touching that logic. + +The value is **signal-driven**, not a fixed constant. Its primary input is the build system's available capacity, but a policy may also weigh historical pass rates, cost budgets, time of day, or an experiment toggle. + +Unlike the dependency limit — which the controller holds and applies as an eligibility gate — the selection limit is **injected into the seam that uses it**: the selector is constructed with it and calls it itself, never receiving it as a method parameter. This follows the repo's extension-contract pattern (dependencies injected at the `Factory`), keeps the selector interface limit-free and stable, and lets the limit be swapped independently of selector logic. + +## Factory + +A per-queue factory returns the limit policy for a queue, following the repo's extension contract. It is handed only the queue identity; the signals a policy weighs — a capacity feed, historical metrics, config — are injected at construction by the integrator in the wiring layer, which is also where the limit is handed to the selector. Computing the limit itself takes no further inputs. diff --git a/submitqueue/extension/speculation/selectionlimit/fake/BUILD.bazel b/submitqueue/extension/speculation/selectionlimit/fake/BUILD.bazel new file mode 100644 index 00000000..cf5ff1ec --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/fake/BUILD.bazel @@ -0,0 +1,19 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["fake.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selectionlimit/fake", + visibility = ["//visibility:public"], + deps = ["//submitqueue/extension/speculation/selectionlimit:go_default_library"], +) + +go_test( + name = "go_default_test", + srcs = ["fake_test.go"], + embed = [":go_default_library"], + deps = [ + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/selectionlimit/fake/fake.go b/submitqueue/extension/speculation/selectionlimit/fake/fake.go new file mode 100644 index 00000000..4311f831 --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/fake/fake.go @@ -0,0 +1,52 @@ +// 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 fake provides a programmable selectionlimit.SelectionLimit for tests +// and examples. New sets the value returned by Limit; FailWith injects an error +// on every call. It is intended for examples and tests only, never production. +package fake + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/extension/speculation/selectionlimit" +) + +// SelectionLimit is a programmable selectionlimit.SelectionLimit. +type SelectionLimit struct { + limit int + err error +} + +// New returns a fake SelectionLimit whose Limit returns the given value. +func New(limit int) *SelectionLimit { + return &SelectionLimit{limit: limit} +} + +// FailWith makes every Limit call return err. +func (l *SelectionLimit) FailWith(err error) *SelectionLimit { + l.err = err + return l +} + +// Limit returns the configured value, or the injected error if FailWith was set. +func (l *SelectionLimit) Limit(_ context.Context) (int, error) { + if l.err != nil { + return 0, l.err + } + return l.limit, nil +} + +// ensure the fake satisfies the interface. +var _ selectionlimit.SelectionLimit = (*SelectionLimit)(nil) diff --git a/submitqueue/extension/speculation/selectionlimit/fake/fake_test.go b/submitqueue/extension/speculation/selectionlimit/fake/fake_test.go new file mode 100644 index 00000000..a69497b1 --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/fake/fake_test.go @@ -0,0 +1,36 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package fake + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLimit_ReturnsConfiguredValue(t *testing.T) { + got, err := New(2).Limit(context.Background()) + require.NoError(t, err) + assert.Equal(t, 2, got) +} + +func TestLimit_FailWith(t *testing.T) { + sentinel := errors.New("boom") + _, err := New(2).FailWith(sentinel).Limit(context.Background()) + require.ErrorIs(t, err, sentinel) +} diff --git a/submitqueue/extension/speculation/selectionlimit/mock/BUILD.bazel b/submitqueue/extension/speculation/selectionlimit/mock/BUILD.bazel new file mode 100644 index 00000000..7d9fceec --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/mock/BUILD.bazel @@ -0,0 +1,12 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["selectionlimit_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selectionlimit/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/extension/speculation/selectionlimit:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/selectionlimit/mock/selectionlimit_mock.go b/submitqueue/extension/speculation/selectionlimit/mock/selectionlimit_mock.go new file mode 100644 index 00000000..f4cbbf14 --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/mock/selectionlimit_mock.go @@ -0,0 +1,96 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: selectionlimit.go +// +// Generated by this command: +// +// mockgen -source=selectionlimit.go -destination=mock/selectionlimit_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + selectionlimit "github.com/uber/submitqueue/submitqueue/extension/speculation/selectionlimit" + gomock "go.uber.org/mock/gomock" +) + +// MockSelectionLimit is a mock of SelectionLimit interface. +type MockSelectionLimit struct { + ctrl *gomock.Controller + recorder *MockSelectionLimitMockRecorder + isgomock struct{} +} + +// MockSelectionLimitMockRecorder is the mock recorder for MockSelectionLimit. +type MockSelectionLimitMockRecorder struct { + mock *MockSelectionLimit +} + +// NewMockSelectionLimit creates a new mock instance. +func NewMockSelectionLimit(ctrl *gomock.Controller) *MockSelectionLimit { + mock := &MockSelectionLimit{ctrl: ctrl} + mock.recorder = &MockSelectionLimitMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSelectionLimit) EXPECT() *MockSelectionLimitMockRecorder { + return m.recorder +} + +// Limit mocks base method. +func (m *MockSelectionLimit) Limit(ctx context.Context) (int, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Limit", ctx) + ret0, _ := ret[0].(int) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Limit indicates an expected call of Limit. +func (mr *MockSelectionLimitMockRecorder) Limit(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Limit", reflect.TypeOf((*MockSelectionLimit)(nil).Limit), ctx) +} + +// 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 selectionlimit.Config) (selectionlimit.SelectionLimit, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "For", cfg) + ret0, _ := ret[0].(selectionlimit.SelectionLimit) + 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/selectionlimit/selectionlimit.go b/submitqueue/extension/speculation/selectionlimit/selectionlimit.go new file mode 100644 index 00000000..68cfdcd0 --- /dev/null +++ b/submitqueue/extension/speculation/selectionlimit/selectionlimit.go @@ -0,0 +1,57 @@ +// 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 selectionlimit + +//go:generate mockgen -source=selectionlimit.go -destination=mock/selectionlimit_mock.go -package=mock + +import "context" + +// SelectionLimit is the "how much" policy that bounds how many paths a batch may +// build in parallel. +// +// It is the selector's companion: the selector decides *which* of a batch's +// paths are worth building (its ranking); the selection limit decides *how many* +// of them may run at once. Separating the two keeps selector logic free of +// resource accounting and lets the bound scale with build resources without +// touching that logic. +// +// The value is dynamic: it may change between calls, so the selector reads it +// each pass rather than caching it. +// +// Unlike the dependency limit, this limit is injected into the seam that uses it +// — the selector is constructed with it and calls it itself — never passed as a +// method parameter, keeping the selector interface limit-free and stable. +type SelectionLimit interface { + // Limit returns the current maximum number of paths a batch may build in + // parallel. The selector caps its Promote decisions at this. It takes no + // parameters; anything an implementation needs is injected at construction. + Limit(ctx context.Context) (int, error) +} + +// Config carries the per-queue identity handed to a Factory. The system knows +// only the queue name; everything a policy needs to compute the limit (a +// capacity feed, historical metrics, config) is injected at construction by the +// integrator. +type Config struct { + // QueueName identifies the queue this SelectionLimit serves. + QueueName string +} + +// Factory builds the SelectionLimit for a queue. Implementations are provided by +// integrators (and tests) and inject whatever signals they need at construction. +type Factory interface { + // For returns the SelectionLimit for the given queue. + For(cfg Config) (SelectionLimit, error) +} diff --git a/submitqueue/extension/speculation/selector/BUILD.bazel b/submitqueue/extension/speculation/selector/BUILD.bazel new file mode 100644 index 00000000..aca92b3d --- /dev/null +++ b/submitqueue/extension/speculation/selector/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["selector.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selector", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity:go_default_library"], +) diff --git a/submitqueue/extension/speculation/selector/README.md b/submitqueue/extension/speculation/selector/README.md new file mode 100644 index 00000000..cd13c426 --- /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 the batch's tree and their current status, what should we do with each, right now?* It returns an **action** per path — `Promote` (advance it one stage toward running) 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 controller hands the selector the batch's **speculation tree** directly — the subject it decides over. The controller is the single writer: it reconciles each path's status (candidate, selected, prioritized, building, passed, failed, cancelling, cancelled) from the latest builds and dependency states, and it maps each of the selector's decisions to a status transition — `Promote` → `Selected`, `Cancel` → `Cancelling` (or `Cancelled` when nothing is building) — applied under the tree's optimistic lock and persisted. The selector's only output is decisions; it **never** writes status. This keeps it a deterministic policy over the tree it is given. + +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 + +A per-queue factory returns the selector for a queue, following the repo's extension contract. It is handed only the queue identity and 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`. Selection itself stays config-free. diff --git a/submitqueue/extension/speculation/selector/fake/BUILD.bazel b/submitqueue/extension/speculation/selector/fake/BUILD.bazel new file mode 100644 index 00000000..7013e840 --- /dev/null +++ b/submitqueue/extension/speculation/selector/fake/BUILD.bazel @@ -0,0 +1,23 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["fake.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/fake", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/selector:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["fake_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/selector/fake/fake.go b/submitqueue/extension/speculation/selector/fake/fake.go new file mode 100644 index 00000000..05de8c0c --- /dev/null +++ b/submitqueue/extension/speculation/selector/fake/fake.go @@ -0,0 +1,61 @@ +// 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 fake provides a programmable selector.Selector for tests and examples. +// It returns the decisions seeded via SetDecisions (none by default, i.e. leave +// every path as-is). FailWith injects an error on every call. It is intended for +// examples and tests only, never production. +package fake + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/selector" +) + +// Selector is a programmable selector.Selector. +type Selector struct { + decisions []entity.SpeculationPathDecision + err error +} + +// New returns a fake Selector that decides nothing (leaves every path as-is). +// Seed decisions with SetDecisions. +func New() *Selector { + return &Selector{} +} + +// SetDecisions seeds the decisions returned by Select. +func (s *Selector) SetDecisions(decisions ...entity.SpeculationPathDecision) *Selector { + s.decisions = decisions + return s +} + +// FailWith makes every Select call return err. +func (s *Selector) FailWith(err error) *Selector { + s.err = err + return s +} + +// Select returns the seeded decisions. The tree argument is ignored. +func (s *Selector) Select(_ context.Context, _ entity.SpeculationTree) ([]entity.SpeculationPathDecision, error) { + if s.err != nil { + return nil, s.err + } + return s.decisions, nil +} + +// ensure the fake satisfies the interface. +var _ selector.Selector = (*Selector)(nil) diff --git a/submitqueue/extension/speculation/selector/fake/fake_test.go b/submitqueue/extension/speculation/selector/fake/fake_test.go new file mode 100644 index 00000000..f116ada1 --- /dev/null +++ b/submitqueue/extension/speculation/selector/fake/fake_test.go @@ -0,0 +1,46 @@ +// 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 fake + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/submitqueue/entity" +) + +func TestSelect_DefaultDecidesNothing(t *testing.T) { + got, err := New().Select(context.Background(), entity.SpeculationTree{BatchID: "q/batch/2"}) + require.NoError(t, err) + assert.Empty(t, got) +} + +func TestSelect_ReturnsSeededDecisions(t *testing.T) { + want := []entity.SpeculationPathDecision{ + {Path: entity.SpeculationPath{Head: "q/batch/2"}, Action: entity.SpeculationPathActionPromote}, + } + got, err := New().SetDecisions(want...).Select(context.Background(), entity.SpeculationTree{BatchID: "q/batch/2"}) + require.NoError(t, err) + assert.Equal(t, want, got) +} + +func TestSelect_FailWith(t *testing.T) { + sentinel := errors.New("boom") + _, err := New().FailWith(sentinel).Select(context.Background(), entity.SpeculationTree{BatchID: "q/batch/2"}) + require.ErrorIs(t, err, sentinel) +} diff --git a/submitqueue/extension/speculation/selector/mock/BUILD.bazel b/submitqueue/extension/speculation/selector/mock/BUILD.bazel new file mode 100644 index 00000000..e31a192b --- /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 = "go_default_library", + srcs = ["selector_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/selector/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/selector:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) 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..38ee5b7a --- /dev/null +++ b/submitqueue/extension/speculation/selector/selector.go @@ -0,0 +1,64 @@ +// 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 batch's tree — each path's +// controller-stamped Status (Candidate / Selected / Prioritized / Building / +// Passed / Failed / Cancelling / Cancelled) and Score — and returns a decision +// per path it wants to act on. +// +// The controller hands the selector the batch's speculation tree directly — the +// subject it decides over. Its only output is decisions (Promote / Cancel); it +// never writes Status. The controller maps each decision to a status +// transition — Promote → Selected, Cancel → Cancelling (or Cancelled) — applied +// under the tree's optimistic lock, and persists it, staying the single writer +// of tree state. This keeps the selector a deterministic policy over the tree it +// is given. 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 decisions to take for the given tree. Returning multiple + // Promote decisions selects several paths for building in parallel — each + // waits for the prioritizer to clear it under the queue's build budget before + // a build is dispatched. 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; policy knobs such as a top-K cap or build budget are +// 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) +}