From 10e097e749777f345808ee6b236c06ad9c5d614c Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 24 Jul 2026 10:53:04 -0700 Subject: [PATCH] feat(speculation): allocator contract and sticky impl Add submitqueue/extension/speculation/allocator, the budget-spending composition point (Allocator) the standard Speculator hands its candidate stream to, plus the sticky implementation and mocks. sticky fills only free budget slots and never preempts a running build, trading responsiveness for never discarding work already started; its policy counterpart is a preempting allocator. --- .../speculation/allocator/BUILD.bazel | 12 ++ .../extension/speculation/allocator/README.md | 15 ++ .../speculation/allocator/allocator.go | 39 ++++ .../speculation/allocator/mock/BUILD.bazel | 13 ++ .../allocator/mock/allocator_mock.go | 58 +++++ .../speculation/allocator/sticky/BUILD.bazel | 29 +++ .../speculation/allocator/sticky/sticky.go | 108 ++++++++++ .../allocator/sticky/sticky_test.go | 202 ++++++++++++++++++ 8 files changed, 476 insertions(+) create mode 100644 submitqueue/extension/speculation/allocator/BUILD.bazel create mode 100644 submitqueue/extension/speculation/allocator/README.md create mode 100644 submitqueue/extension/speculation/allocator/allocator.go create mode 100644 submitqueue/extension/speculation/allocator/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/allocator/mock/allocator_mock.go create mode 100644 submitqueue/extension/speculation/allocator/sticky/BUILD.bazel create mode 100644 submitqueue/extension/speculation/allocator/sticky/sticky.go create mode 100644 submitqueue/extension/speculation/allocator/sticky/sticky_test.go diff --git a/submitqueue/extension/speculation/allocator/BUILD.bazel b/submitqueue/extension/speculation/allocator/BUILD.bazel new file mode 100644 index 00000000..f1003432 --- /dev/null +++ b/submitqueue/extension/speculation/allocator/BUILD.bazel @@ -0,0 +1,12 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["allocator.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/allocator", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/allocator/README.md b/submitqueue/extension/speculation/allocator/README.md new file mode 100644 index 00000000..3d5586ac --- /dev/null +++ b/submitqueue/extension/speculation/allocator/README.md @@ -0,0 +1,15 @@ +# allocator + +The `allocator` package defines a composition point used by the `standard` `Speculator`: an `Allocator` spends the queue's build budget over a candidate iterator. Like `generator`, it is **not** a controller-facing extension — an alternate `Speculator` need not use or expose this split — so there is no `Config` or `Factory` here; an `Allocator` is chosen when the `standard` `Speculator` is constructed. + +`Allocate` pulls candidates from the iterator in best-first order and matches them against the queue's current path sets by path ID, so a pending or building path stays funded rather than starting a new attempt. Pending, building, and cancelling paths all charge the budget (a cancelling build holds its slot until it reaches a terminal state); terminal paths charge nothing. It returns the build and cancel actions that spend the budget. + +Because cancellation is best-effort, an `Allocator` should not spend capacity it merely expects a cancel to release — a build cancelled to make room keeps charging the budget until its cancel reaches a terminal state, so the queue converges over successive runs rather than oversubscribing the hard cap on concurrent builds in a single pass. + +## `sticky` + +The `sticky` `Allocator` fills only free budget slots and leaves every in-flight build running. Against the budget it counts the paths that still hold a build slot — `pending`, `building`, and `cancelling` (a cancel is a request, and the build keeps its slot until it actually stops). It then proposes a build for each new candidate, in order, until the budget fills; it never proposes a cancel. + +A stored path with no status at all is a data defect. It holds no slot, since a status that never turns terminal would pin one forever, and `sticky` logs a warning so it does not pass unnoticed. + +The trade-off: because `sticky` never preempts, a better candidate that arrives while the budget is full waits for a running build to finish rather than displacing it — it never discards work already started, but it also cannot react to a newly-attractive path mid-flight. Its policy counterpart is a preempting `Allocator` that cancels a low-value in-flight path to fund a better candidate, spending a cancel now to converge faster over the next run. diff --git a/submitqueue/extension/speculation/allocator/allocator.go b/submitqueue/extension/speculation/allocator/allocator.go new file mode 100644 index 00000000..45e8857d --- /dev/null +++ b/submitqueue/extension/speculation/allocator/allocator.go @@ -0,0 +1,39 @@ +// 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 allocator defines the budget-spending composition point used by the +// standard Speculator. An Allocator spends the queue's build budget over a +// candidate iterator. It is not a controller-facing extension — an alternate +// Speculator need not use or expose this split. +package allocator + +//go:generate mockgen -source=allocator.go -destination=mock/allocator_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// Allocator decides how to spend the queue's build budget over the generator's +// candidate stream, returning the build and cancel actions to take this run. How +// it rations the budget across new candidates and paths already in flight — and +// whether it preempts running builds — is the implementation's policy. +type Allocator interface { + // Allocate returns the build and cancel actions to take, given the candidate + // stream and the queue's current path sets (so it can reconcile candidates + // against paths already in flight). + Allocate(ctx context.Context, pathSets []entity.SpeculationPathSet, iter generator.PathIterator) ([]entity.Speculation, error) +} diff --git a/submitqueue/extension/speculation/allocator/mock/BUILD.bazel b/submitqueue/extension/speculation/allocator/mock/BUILD.bazel new file mode 100644 index 00000000..2e703953 --- /dev/null +++ b/submitqueue/extension/speculation/allocator/mock/BUILD.bazel @@ -0,0 +1,13 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["allocator_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/allocator/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/allocator/mock/allocator_mock.go b/submitqueue/extension/speculation/allocator/mock/allocator_mock.go new file mode 100644 index 00000000..fd5c8e1f --- /dev/null +++ b/submitqueue/extension/speculation/allocator/mock/allocator_mock.go @@ -0,0 +1,58 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: allocator.go +// +// Generated by this command: +// +// mockgen -source=allocator.go -destination=mock/allocator_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" + generator "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" + gomock "go.uber.org/mock/gomock" +) + +// MockAllocator is a mock of Allocator interface. +type MockAllocator struct { + ctrl *gomock.Controller + recorder *MockAllocatorMockRecorder + isgomock struct{} +} + +// MockAllocatorMockRecorder is the mock recorder for MockAllocator. +type MockAllocatorMockRecorder struct { + mock *MockAllocator +} + +// NewMockAllocator creates a new mock instance. +func NewMockAllocator(ctrl *gomock.Controller) *MockAllocator { + mock := &MockAllocator{ctrl: ctrl} + mock.recorder = &MockAllocatorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockAllocator) EXPECT() *MockAllocatorMockRecorder { + return m.recorder +} + +// Allocate mocks base method. +func (m *MockAllocator) Allocate(ctx context.Context, pathSets []entity.SpeculationPathSet, iter generator.PathIterator) ([]entity.Speculation, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Allocate", ctx, pathSets, iter) + ret0, _ := ret[0].([]entity.Speculation) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Allocate indicates an expected call of Allocate. +func (mr *MockAllocatorMockRecorder) Allocate(ctx, pathSets, iter any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Allocate", reflect.TypeOf((*MockAllocator)(nil).Allocate), ctx, pathSets, iter) +} diff --git a/submitqueue/extension/speculation/allocator/sticky/BUILD.bazel b/submitqueue/extension/speculation/allocator/sticky/BUILD.bazel new file mode 100644 index 00000000..54579b29 --- /dev/null +++ b/submitqueue/extension/speculation/allocator/sticky/BUILD.bazel @@ -0,0 +1,29 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["sticky.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/allocator/sticky", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/allocator:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@org_uber_go_zap//:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["sticky_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", + "@org_uber_go_zap//:go_default_library", + "@org_uber_go_zap//zapcore:go_default_library", + "@org_uber_go_zap//zaptest:go_default_library", + "@org_uber_go_zap//zaptest/observer:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/allocator/sticky/sticky.go b/submitqueue/extension/speculation/allocator/sticky/sticky.go new file mode 100644 index 00000000..022a6e79 --- /dev/null +++ b/submitqueue/extension/speculation/allocator/sticky/sticky.go @@ -0,0 +1,108 @@ +// 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 sticky provides an allocator.Allocator that fills only free budget +// slots and leaves every in-flight build running — it never preempts. Its policy +// counterpart is a preempting allocator, which would cancel a low-value in-flight +// path to fund a better candidate; sticky trades that responsiveness for never +// discarding a build it has already started. +package sticky + +import ( + "context" + + "go.uber.org/zap" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/allocator" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// alloc is a sticky allocator.Allocator. budget is the queue's cap on concurrent +// builds; pending, building, and cancelling paths all charge against it. +type alloc struct { + logger *zap.SugaredLogger + budget int +} + +// New returns a sticky allocator.Allocator with the given concurrent-build +// budget. +func New(logger *zap.SugaredLogger, budget int) allocator.Allocator { + return alloc{logger: logger.Named("sticky"), budget: budget} +} + +// Allocate keeps every in-flight path funded and pulls candidates in order into +// the remaining free budget, proposing a build for each new path until the +// budget fills. It never proposes a cancel. +func (a alloc) Allocate(ctx context.Context, pathSets []entity.SpeculationPathSet, iter generator.PathIterator) ([]entity.Speculation, error) { + // Paths already holding a build slot. This set does double duty: it is both + // the count of budget already spent and the guard against proposing a second + // build for something already running, so a status belongs here if either + // reason applies. + funded := make(map[string]bool) + for _, set := range pathSets { + for _, entry := range set.Paths { + if entry.Status == entity.SpeculationPathStatusUnknown { + // A stored path should always carry a real status, so this is a + // data defect. Skip it rather than counting it: a status that + // never turns terminal would pin a build slot forever. That leaves + // the path free to be proposed again, which is self-correcting, + // but the malformed entry is worth knowing about. + a.logger.Warnw("speculation path has no status", + "batch_id", set.BatchID, "path_id", entry.ID) + continue + } + // Anything else that has not finished is still holding a slot: + // pending and building obviously, and cancelling too, since the build + // keeps its slot until it actually stops. + if !entry.Status.IsTerminal() { + funded[entry.ID] = true + } + } + } + + free := a.budget - len(funded) + + var out []entity.Speculation + proposed := make(map[string]bool) + for free > 0 { + c, ok, err := iter.Next(ctx) + if err != nil { + return nil, err + } + if !ok { + // ok=false means the generator has nothing left to offer: every path + // the queue could speculate on has been seen. Stopping here with + // budget still free is normal, not a failure — there is simply + // nothing else worth building right now. + break + } + id := c.Path.ID() + if funded[id] || proposed[id] { + // funded: this path is already building. The generator only skips + // paths that have *finished*, so one still in flight comes back round + // as a candidate. It needs no new action and is already charging the + // budget, so skip it without spending anything. + // + // proposed: a generator is not supposed to repeat a candidate, so + // this should never fire. It is a cheap guard against one that does, + // since a duplicate would otherwise be paid for twice. + continue + } + out = append(out, entity.Speculation{Path: c.Path, Action: entity.PathActionBuild}) + proposed[id] = true + free-- + } + return out, nil +} diff --git a/submitqueue/extension/speculation/allocator/sticky/sticky_test.go b/submitqueue/extension/speculation/allocator/sticky/sticky_test.go new file mode 100644 index 00000000..6deee997 --- /dev/null +++ b/submitqueue/extension/speculation/allocator/sticky/sticky_test.go @@ -0,0 +1,202 @@ +// 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 sticky + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/zap" + "go.uber.org/zap/zapcore" + "go.uber.org/zap/zaptest" + "go.uber.org/zap/zaptest/observer" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// fakeIter yields a fixed slice of candidates, optionally failing on the first pull. +type fakeIter struct { + items []entity.CandidatePath + idx int + fail bool +} + +func (f *fakeIter) Next(context.Context) (entity.CandidatePath, bool, error) { + if f.fail { + return entity.CandidatePath{}, false, errors.New("iterator boom") + } + if f.idx >= len(f.items) { + return entity.CandidatePath{}, false, nil + } + c := f.items[f.idx] + f.idx++ + return c, true, nil +} + +// candidate builds a one-dependency candidate path for the given head and bet. +func candidate(head string, bet entity.DependencyBetType, dep string) entity.CandidatePath { + return entity.CandidatePath{ + Path: entity.SpeculationPath{ + Head: head, + Bets: []entity.DependencyBet{{Batch: dep, Bet: bet}}, + }, + } +} + +// fundedSet returns a path set holding one entry for head at the given status, +// whose ID matches candidate(head, bet, dep) so it lines up with an iterator item. +func fundedSet(head string, bet entity.DependencyBetType, dep string, status entity.SpeculationPathStatus) entity.SpeculationPathSet { + c := candidate(head, bet, dep) + return entity.SpeculationPathSet{ + BatchID: head, + Paths: []entity.SpeculationPathEntry{{ID: c.Path.ID(), Path: c.Path, Status: status}}, + } +} + +func TestSticky_Allocate(t *testing.T) { + tests := []struct { + name string + budget int + pathSets []entity.SpeculationPathSet + items []entity.CandidatePath + iterFail bool + wantErr bool + // wantIdx are indices into items expected to be proposed as builds, in order. + wantIdx []int + }{ + { + name: "fills free budget in order", + budget: 2, + items: []entity.CandidatePath{ + candidate("q/2", entity.BetIncluded, "q/1"), + candidate("q/2", entity.BetExcluded, "q/1"), + candidate("q/3", entity.BetIncluded, "q/1"), + }, + wantIdx: []int{0, 1}, + }, + { + name: "keeps in-flight path funded and fills the free slot", + budget: 2, + pathSets: []entity.SpeculationPathSet{fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusBuilding)}, + items: []entity.CandidatePath{ + candidate("q/2", entity.BetIncluded, "q/1"), // already funded -> skipped + candidate("q/2", entity.BetExcluded, "q/1"), // fills the one free slot + }, + wantIdx: []int{1}, + }, + { + name: "terminal entry does not charge budget", + budget: 1, + pathSets: []entity.SpeculationPathSet{fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusPassed)}, + items: []entity.CandidatePath{candidate("q/3", entity.BetExcluded, "q/1")}, + wantIdx: []int{0}, + }, + { + name: "pending entry charges budget", + budget: 1, + pathSets: []entity.SpeculationPathSet{fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusPending)}, + items: []entity.CandidatePath{candidate("q/3", entity.BetIncluded, "q/1")}, + wantIdx: nil, + }, + { + // Cancelling is non-terminal, so it still charges the budget: sticky + // must not spend a slot it only expects the cancel to free, which would + // oversubscribe the hard cap on concurrent builds. + name: "cancelling entry charges budget", + budget: 1, + pathSets: []entity.SpeculationPathSet{fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusCancelling)}, + items: []entity.CandidatePath{candidate("q/3", entity.BetIncluded, "q/1")}, + wantIdx: nil, + }, + { + // A zero-valued status is a data defect and holds no slot: one that + // never turns terminal would otherwise pin a slot forever. The budget + // stays free and the path is proposed again. + name: "unknown entry does not charge budget", + budget: 1, + pathSets: []entity.SpeculationPathSet{fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusUnknown)}, + items: []entity.CandidatePath{ + candidate("q/2", entity.BetIncluded, "q/1"), // the unknown entry itself + candidate("q/3", entity.BetIncluded, "q/1"), + }, + wantIdx: []int{0}, + }, + { + name: "deduplicates a repeated candidate", + budget: 5, + items: []entity.CandidatePath{ + candidate("q/2", entity.BetIncluded, "q/1"), + candidate("q/2", entity.BetIncluded, "q/1"), + }, + wantIdx: []int{0}, + }, + { + name: "propagates iterator error", + budget: 2, + iterFail: true, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + iter := &fakeIter{items: tt.items, fail: tt.iterFail} + got, err := New(zaptest.NewLogger(t).Sugar(), tt.budget). + Allocate(context.Background(), tt.pathSets, iter) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + + require.Len(t, got, len(tt.wantIdx)) + for i, idx := range tt.wantIdx { + assert.Equal(t, entity.PathActionBuild, got[i].Action) + assert.Equal(t, tt.items[idx].Path.ID(), got[i].Path.ID()) + } + }) + } +} + +func TestSticky_WarnsOnPathWithNoStatus(t *testing.T) { + // A zero-valued status is a data defect. It must not be swallowed silently, + // since nothing else in the pipeline will notice it. + core, logs := observer.New(zapcore.WarnLevel) + pathSets := []entity.SpeculationPathSet{ + fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusUnknown), + } + iter := &fakeIter{} + + _, err := New(zap.New(core).Sugar(), 1).Allocate(context.Background(), pathSets, iter) + require.NoError(t, err) + + assert.Equal(t, 1, logs.Len(), "a path with no status should warn exactly once") +} + +func TestSticky_QuietWhenEveryPathHasAStatus(t *testing.T) { + core, logs := observer.New(zapcore.WarnLevel) + pathSets := []entity.SpeculationPathSet{ + fundedSet("q/2", entity.BetIncluded, "q/1", entity.SpeculationPathStatusBuilding), + } + iter := &fakeIter{} + + _, err := New(zap.New(core).Sugar(), 1).Allocate(context.Background(), pathSets, iter) + require.NoError(t, err) + + assert.Zero(t, logs.Len()) +}