feat(prioritization): add prioritizer + prioritization-limit extensions

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/... ./platform/consumer/... ./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/... ./submitqueue/extension/prioritization/prioritizer/... ./submitqueue/extension/prioritization/limit/... ./platform/consumer/... ./submitqueue/core/changeset/... ./stovepipe/extension/storage/... ./stovepipe/extension/sourcecontrol/...
 	@echo "Mocks generated successfully!"
 
 proto: ## Generate protobuf files from .proto definitions

submitqueue/entity/build.go

@@ -53,12 +53,6 @@ func (s BuildStatus) IsTerminal() bool {
 	return s == BuildStatusSucceeded || s == BuildStatusFailed || s == BuildStatusCancelled
 }
 
-// SpeculationPathInfo represents the base and head commits of a speculation path used in a build.
-type SpeculationPathInfo struct {
-	// Base is a list of batchIDs(in order) that form the base of this speculation path.
-	Base []string
-}
-
 // Build represents a build scheduled for a batch along a specific speculation path.
 // All fields except the Status are immutable after creation.
 type Build struct {
@@ -69,10 +63,9 @@ type Build struct {
 	BatchID string
 	// SpeculationPath is the speculation path that represents this build. For
 	// a given batch this path is crafted from the graph that is generated from the
-	// dependencies of this batch.
-	SpeculationPath SpeculationPathInfo
-	// Score represents the build prediction score for this speculation path.
-	Score float32
+	// dependencies of this batch. Its Head is the batch being verified (equal to
+	// BatchID) and its Base is the assumed-good prefix of predecessor batches.
+	SpeculationPath SpeculationPath
 	// Status represents the state of the build lifecycle this build is in.
 	Status BuildStatus
 }

submitqueue/entity/build_test.go

@@ -70,10 +70,10 @@ func TestBuild_ToBytes(t *testing.T) {
 	build := Build{
 		ID:      "build-1",
 		BatchID: "batch-1",
-		SpeculationPath: SpeculationPathInfo{
+		SpeculationPath: SpeculationPath{
 			Base: []string{"batch-0", "batch-prev"},
+			Head: "batch-1",
 		},
-		Score:  0.85,
 		Status: BuildStatusAccepted,
 	}
 
@@ -92,10 +92,10 @@ func TestBuildFromBytes(t *testing.T) {
 	original := Build{
 		ID:      "build-42",
 		BatchID: "batch-7",
-		SpeculationPath: SpeculationPathInfo{
+		SpeculationPath: SpeculationPath{
 			Base: []string{"batch-5", "batch-6"},
+			Head: "batch-7",
 		},
-		Score:  0.92,
 		Status: BuildStatusAccepted,
 	}
 
@@ -111,7 +111,6 @@ func TestBuildFromBytes(t *testing.T) {
 	assert.Equal(t, original.ID, deserialized.ID)
 	assert.Equal(t, original.BatchID, deserialized.BatchID)
 	assert.Equal(t, original.SpeculationPath.Base, deserialized.SpeculationPath.Base)
-	assert.Equal(t, original.Score, deserialized.Score)
 	assert.Equal(t, original.Status, deserialized.Status)
 }
 
@@ -132,7 +131,6 @@ func TestBuildFromBytes_EmptyData(t *testing.T) {
 	assert.Empty(t, build.ID)
 	assert.Empty(t, build.BatchID)
 	assert.Equal(t, BuildStatusUnknown, build.Status)
-	assert.Equal(t, float32(0), build.Score)
 }
 
 func TestBuild_SerializationRoundTrip(t *testing.T) {
@@ -145,10 +143,10 @@ func TestBuild_SerializationRoundTrip(t *testing.T) {
 			build: Build{
 				ID:      "build-100",
 				BatchID: "batch-50",
-				SpeculationPath: SpeculationPathInfo{
+				SpeculationPath: SpeculationPath{
 					Base: []string{"batch-48", "batch-49"},
+					Head: "batch-50",
 				},
-				Score:  0.75,
 				Status: BuildStatusAccepted,
 			},
 		},
@@ -157,19 +155,18 @@ func TestBuild_SerializationRoundTrip(t *testing.T) {
 			build: Build{
 				ID:      "build-200",
 				BatchID: "batch-60",
-				Score:   1.0,
 				Status:  BuildStatusSucceeded,
 			},
 		},
 		{
-			name: "failed build with zero score",
+			name: "failed build",
 			build: Build{
 				ID:      "build-300",
 				BatchID: "batch-70",
-				SpeculationPath: SpeculationPathInfo{
+				SpeculationPath: SpeculationPath{
 					Base: []string{"batch-65"},
+					Head: "batch-70",
 				},
-				Score:  0,
 				Status: BuildStatusFailed,
 			},
 		},

submitqueue/entity/speculation_tree.go

@@ -14,38 +14,120 @@
 
 package entity
 
-// SpeculationPathAction defines the possible actions for a speculation path.
+// SpeculationPath is a single speculation path: an assumed-good prefix of
+// predecessor batches (Base) on top of which the batch under verification
+// (Head) is built and validated.
+//
+// This is the unit the build stage consumes: Base maps to the build runner's
+// base changes (an assumed-good prefix to apply) and Head maps to the changes
+// being validated.
+type SpeculationPath struct {
+	// Base is the ordered list of predecessor batch IDs assumed to have passed.
+	// Empty means the path builds the head batch directly on the target branch.
+	Base []string
+	// Head is the batch ID being verified by this path.
+	Head string
+}
+
+// SpeculationPathStatus is the observed lifecycle state of a speculation path.
+// It is written only by the orchestrator's speculate controller (into the
+// speculation tree store) and read by the path selector as input; enumerators
+// and selectors never write it.
+type SpeculationPathStatus string
+
+const (
+	// SpeculationPathStatusUnknown is the unreachable zero value, set by default
+	// on init. A persisted path always carries a real status (candidate onward),
+	// so this should never be seen in the store.
+	SpeculationPathStatusUnknown SpeculationPathStatus = ""
+	// SpeculationPathStatusCandidate is a freshly enumerated path the controller
+	// has persisted but not yet sent to build.
+	SpeculationPathStatusCandidate SpeculationPathStatus = "candidate"
+	// SpeculationPathStatusSelected is a path the controller has sent to the build
+	// controller (in response to a selector Build action) but for which no build
+	// signal has arrived yet — the build system may not have started it
+	// (resource-gated), so whether it is actually building is not yet known.
+	SpeculationPathStatusSelected SpeculationPathStatus = "selected"
+	// SpeculationPathStatusBuilding is a path a build signal has confirmed is in
+	// flight; its BuildID is known.
+	SpeculationPathStatusBuilding SpeculationPathStatus = "building"
+	// SpeculationPathStatusPassed is a path whose build succeeded.
+	SpeculationPathStatusPassed SpeculationPathStatus = "passed"
+	// SpeculationPathStatusFailed is a path whose build failed.
+	SpeculationPathStatusFailed SpeculationPathStatus = "failed"
+	// SpeculationPathStatusCancelled is a path that is no longer pursued — its
+	// base was invalidated, its build was cancelled, or the selector dropped it.
+	SpeculationPathStatusCancelled SpeculationPathStatus = "cancelled"
+)
+
+// SpeculationPathAction is the action a path selector asks the controller to
+// take for a path. It is the selector's only output: ephemeral (recomputed
+// every time the selector runs) and never persisted. The controller enacts it
+// and records the resulting SpeculationPathStatus.
 type SpeculationPathAction string
 
 const (
-	// SpeculationPathActionUnknown is the default zero value for SpeculationPathAction.
+	// SpeculationPathActionUnknown is the unreachable zero value. A real decision
+	// always carries Build or Cancel; the selector expresses "leave this path
+	// as-is" by omitting it from its decisions, not by returning this.
 	SpeculationPathActionUnknown SpeculationPathAction = ""
-	// TODO: Add comprehensive list of actions
+	// SpeculationPathActionBuild asks the controller to send this path to the
+	// build controller (which triggers a build subject to resources). The path moves
+	// to Selected on send, then Building once a build signal confirms it.
+	SpeculationPathActionBuild SpeculationPathAction = "build"
+	// SpeculationPathActionCancel asks the controller to drop this path and
+	// cancel any build in flight for it.
+	SpeculationPathActionCancel SpeculationPathAction = "cancel"
 )
 
-// SpeculationInfo represents metadata about a single speculation path, including the path through the dependency graph, its current state, and the predicted build score.
-type SpeculationInfo struct {
-	// Path represents the speculation path; which is an ordered list of batches.
-	Path []string
-	// Action is a state that this path is in.
-	Action SpeculationPathAction
-	// Score is score for this speculation path.
+// SpeculationPathInfo is the per-path entry in a speculation tree: a path, its
+// latest predicted-success score, its controller-owned status, and a link to
+// the build dispatched for it (if any).
+type SpeculationPathInfo struct {
+	// Path is the Base/Head split this entry covers.
+	Path SpeculationPath
+	// Score is the path's predicted-success score. It is computed by the scorer
+	// and persisted by the controller, not set at enumeration — the enumerator
+	// produces structure only. It is dynamic: the controller re-runs the scorer
+	// on every respeculate (as dependencies land, dependency builds pass, or
+	// sibling paths fail), so the value tracks the latest state rather than a
+	// figure frozen when the path was first enumerated (~0 until the first pass).
 	Score float32
+	// Status is the observed lifecycle state of the path. Written only by the
+	// controller; read by the selector.
+	Status SpeculationPathStatus
+	// BuildID links this path to its build. Empty until a build signal confirms
+	// the build and the controller records it (Selected -> Building); the
+	// controller never knows the ID at send time.
+	BuildID string
+}
+
+// SpeculationPathDecision is a path selector's decision for a single path: the
+// action the controller should take for it. It is the selector's output and is
+// not persisted.
+type SpeculationPathDecision struct {
+	// Path identifies the speculation path the action applies to.
+	Path SpeculationPath
+	// Action is what the controller should do for the path.
+	Action SpeculationPathAction
 }
 
-// SpeculationTree represents the set of speculation paths constructed for a batch based on its dependency graph.
+// SpeculationTree is the set of candidate speculation paths for a batch, built
+// from its dependency graph.
 type SpeculationTree struct {
 	// BatchID is the batch for which this speculation tree is constructed.
 	BatchID string
-	// Speculations is a list of speculation paths for this batch based on a graph of its
-	// dependents.
+	// Paths is the candidate speculation paths for this batch, derived from a
+	// graph of its dependencies. Each entry's per-path dynamic state (Score,
+	// Status, BuildID) is documented on SpeculationPathInfo.
 	//
 	// For e.g - Consider batches - queueA/batch/1, queueA/batch/2, queueA/batch/3
-	// such that - queueA/batch/2 and queueA/batch/3 depend on queueA/batch/1
+	// such that - queueA/batch/2 and queueA/batch/3 depend on queueA/batch/1.
+	// Each dependent batch gets two paths: build alone, or build on the
+	// assumed-good predecessor. Just after enumeration every path is a candidate:
 	//
-	// Speculations for queueA/batch/1 - [{Path: []string{"queueA/batch/1"}, State: "scheduled", Score: 0.1}]
-	// Speculations for queueA/batch/2 - [{Path: []string{"queueA/batch/2"}, State: "scheduled", Score: 0.9}, {Path: []string{"queueA/batch/1", "queueA/batch/2"}, State: "scheduled", Score: 0.3}]
-	// Speculations for queueA/batch/3 - [{Path: []string{"queueA/batch/3"}, State: "scheduled", Score: 0.9}, {Path: []string{"queueA/batch/1", "queueA/batch/3"}, State: "scheduled", Score: 0.3}]
+	// Paths for queueA/batch/2 - [{Path: {Base: [], Head: "queueA/batch/2"}, Status: "candidate"}, {Path: {Base: ["queueA/batch/1"], Head: "queueA/batch/2"}, Status: "candidate"}]
+	// Paths for queueA/batch/3 - [{Path: {Base: [], Head: "queueA/batch/3"}, Status: "candidate"}, {Path: {Base: ["queueA/batch/1"], Head: "queueA/batch/3"}, Status: "candidate"}]
 	//
-	Speculations []SpeculationInfo
+	Paths []SpeculationPathInfo
 }

submitqueue/extension/prioritization/limit/BUILD.bazel

@@ -0,0 +1,8 @@
+load("@rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["limit.go"],
+    importpath = "github.com/uber/submitqueue/submitqueue/extension/prioritization/limit",
+    visibility = ["//visibility:public"],
+)

submitqueue/extension/prioritization/limit/README.md

@@ -0,0 +1,17 @@
+# Prioritization Limit
+
+Vendor-agnostic "how much" policy that bounds how many builds a queue may run at once — the queue's concurrent-build budget.
+
+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.
+
+## Prioritization Limit
+
+The prioritization limit is the [prioritizer](../prioritizer)'s companion. The prioritizer decides **which** of the queue's pending builds run — its ranking across all in-flight batches; the prioritization limit decides **how many** fit at once. It is the queue-wide resource knob, the ultimate cap on speculation's demand on CI.
+
+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 cost budgets, time of day, or an experiment toggle.
+
+It is **injected into the prioritizer** at construction and called by it, never passed as a method parameter — following the repo's extension-contract pattern, keeping the prioritizer interface limit-free and stable, and letting the limit be swapped independently of prioritizer 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, cost budgets, config — are injected at construction by the integrator in the wiring layer, which is also where the limit is handed to the prioritizer. Computing the limit itself takes no further inputs.

submitqueue/extension/prioritization/limit/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/prioritization/limit/fake",
+    visibility = ["//visibility:public"],
+    deps = ["//submitqueue/extension/prioritization/limit: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",
+    ],
+)

submitqueue/extension/prioritization/limit/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 limit.PrioritizationLimit 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/prioritization/limit"
+)
+
+// PrioritizationLimit is a programmable limit.PrioritizationLimit.
+type PrioritizationLimit struct {
+	limit int
+	err   error
+}
+
+// New returns a fake PrioritizationLimit whose Limit returns the given value.
+func New(value int) *PrioritizationLimit {
+	return &PrioritizationLimit{limit: value}
+}
+
+// FailWith makes every Limit call return err.
+func (l *PrioritizationLimit) FailWith(err error) *PrioritizationLimit {
+	l.err = err
+	return l
+}
+
+// Limit returns the configured value, or the injected error if FailWith was set.
+func (l *PrioritizationLimit) Limit(_ context.Context) (int, error) {
+	if l.err != nil {
+		return 0, l.err
+	}
+	return l.limit, nil
+}
+
+// ensure the fake satisfies the interface.
+var _ limit.PrioritizationLimit = (*PrioritizationLimit)(nil)

submitqueue/extension/prioritization/limit/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(8).Limit(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, 8, got)
+}
+
+func TestLimit_FailWith(t *testing.T) {
+	sentinel := errors.New("boom")
+	_, err := New(8).FailWith(sentinel).Limit(context.Background())
+	require.ErrorIs(t, err, sentinel)
+}