From 6bc842643fac7c1604498d4b2d9fd5a66aee2d8d Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Tue, 9 Jun 2026 12:34:18 -0700 Subject: [PATCH] feat(entity,storage): rework speculation tree model and store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reshape the speculation tree data model around the Base/Head path that the build stage actually consumes, and align its store. entity: SpeculationPath{Base, Head} becomes the unit; SpeculationPathInfo carries the path plus its enumerator Score, controller-owned Status (candidate/selected/building/passed/failed/cancelled), and BuildID. Adds SpeculationPathAction and SpeculationPathDecision for the selector seam. SpeculationTree.Speculations becomes Paths. The Build entity drops its own local SpeculationPathInfo and uses the shared SpeculationPath (Head = the batch under verification), with the build controller updated to match. storage: SpeculationTreeStore.UpdateSpeculations(batchID, []SpeculationInfo) becomes Update(ctx, SpeculationTree) — symmetric with Create, keyed by tree.BatchID. MySQL impl and mock updated. --- submitqueue/entity/build.go | 11 +- submitqueue/entity/build_test.go | 12 +- submitqueue/entity/speculation_tree.go | 108 +++++++++++++++--- .../mock/speculation_tree_store_mock.go | 12 +- .../storage/mysql/speculation_tree_store.go | 23 ++-- .../storage/speculation_tree_store.go | 6 +- .../orchestrator/controller/build/build.go | 2 +- 7 files changed, 124 insertions(+), 50 deletions(-) diff --git a/submitqueue/entity/build.go b/submitqueue/entity/build.go index 86039f43..87400774 100644 --- a/submitqueue/entity/build.go +++ b/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,8 +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 + // 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 // Score represents the build prediction score for this speculation path. Score float32 // Status represents the state of the build lifecycle this build is in. diff --git a/submitqueue/entity/build_test.go b/submitqueue/entity/build_test.go index 97de0193..ce33095e 100644 --- a/submitqueue/entity/build_test.go +++ b/submitqueue/entity/build_test.go @@ -70,8 +70,9 @@ 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,8 +93,9 @@ 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, @@ -145,8 +147,9 @@ 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, @@ -166,8 +169,9 @@ func TestBuild_SerializationRoundTrip(t *testing.T) { build: Build{ ID: "build-300", BatchID: "batch-70", - SpeculationPath: SpeculationPathInfo{ + SpeculationPath: SpeculationPath{ Base: []string{"batch-65"}, + Head: "batch-70", }, Score: 0, Status: BuildStatusFailed, diff --git a/submitqueue/entity/speculation_tree.go b/submitqueue/entity/speculation_tree.go index d6bb1750..74baf61e 100644 --- a/submitqueue/entity/speculation_tree.go +++ b/submitqueue/entity/speculation_tree.go @@ -14,38 +14,112 @@ 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, the +// enumerator's predicted score for it, 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 enumerator's predicted success score for this path. 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. // // 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 // - // 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"}, Score: 0.9, Status: "candidate"}, {Path: {Base: ["queueA/batch/1"], Head: "queueA/batch/2"}, Score: 0.3, Status: "candidate"}] + // Paths for queueA/batch/3 - [{Path: {Base: [], Head: "queueA/batch/3"}, Score: 0.9, Status: "candidate"}, {Path: {Base: ["queueA/batch/1"], Head: "queueA/batch/3"}, Score: 0.3, Status: "candidate"}] // - Speculations []SpeculationInfo + Paths []SpeculationPathInfo } diff --git a/submitqueue/extension/storage/mock/speculation_tree_store_mock.go b/submitqueue/extension/storage/mock/speculation_tree_store_mock.go index e7fba22b..c5055fd6 100644 --- a/submitqueue/extension/storage/mock/speculation_tree_store_mock.go +++ b/submitqueue/extension/storage/mock/speculation_tree_store_mock.go @@ -70,16 +70,16 @@ func (mr *MockSpeculationTreeStoreMockRecorder) Get(ctx, batchID any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockSpeculationTreeStore)(nil).Get), ctx, batchID) } -// UpdateSpeculations mocks base method. -func (m *MockSpeculationTreeStore) UpdateSpeculations(ctx context.Context, batchID string, speculations []entity.SpeculationInfo) error { +// Update mocks base method. +func (m *MockSpeculationTreeStore) Update(ctx context.Context, speculationTree entity.SpeculationTree) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateSpeculations", ctx, batchID, speculations) + ret := m.ctrl.Call(m, "Update", ctx, speculationTree) ret0, _ := ret[0].(error) return ret0 } -// UpdateSpeculations indicates an expected call of UpdateSpeculations. -func (mr *MockSpeculationTreeStoreMockRecorder) UpdateSpeculations(ctx, batchID, speculations any) *gomock.Call { +// Update indicates an expected call of Update. +func (mr *MockSpeculationTreeStoreMockRecorder) Update(ctx, speculationTree any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateSpeculations", reflect.TypeOf((*MockSpeculationTreeStore)(nil).UpdateSpeculations), ctx, batchID, speculations) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockSpeculationTreeStore)(nil).Update), ctx, speculationTree) } diff --git a/submitqueue/extension/storage/mysql/speculation_tree_store.go b/submitqueue/extension/storage/mysql/speculation_tree_store.go index 2366dadc..96922fe2 100644 --- a/submitqueue/extension/storage/mysql/speculation_tree_store.go +++ b/submitqueue/extension/storage/mysql/speculation_tree_store.go @@ -59,7 +59,7 @@ func (s *speculationTreeStore) Get(ctx context.Context, batchID string) (ret ent return entity.SpeculationTree{}, fmt.Errorf("failed to get speculation tree entity batchID=%s from the database: %w", batchID, err) } - if err := json.Unmarshal(speculationsJSON, &st.Speculations); err != nil { + if err := json.Unmarshal(speculationsJSON, &st.Paths); err != nil { return entity.SpeculationTree{}, fmt.Errorf("failed to unmarshal speculations for speculation tree entity batchID=%s from the database: %w", batchID, err) } @@ -71,7 +71,7 @@ func (s *speculationTreeStore) Create(ctx context.Context, speculationTree entit op := metrics.Begin(s.scope, "create") defer func() { op.Complete(retErr) }() - speculationsJSON, err := json.Marshal(speculationTree.Speculations) + speculationsJSON, err := json.Marshal(speculationTree.Paths) if err != nil { return fmt.Errorf("failed to marshal speculations batchID=%s for Create speculation tree entity: %w", speculationTree.BatchID, err) } @@ -91,31 +91,32 @@ func (s *speculationTreeStore) Create(ctx context.Context, speculationTree entit return nil } -// UpdateSpeculations updates the speculations of a speculation tree. Returns ErrNotFound if the speculation tree is not found. -func (s *speculationTreeStore) UpdateSpeculations(ctx context.Context, batchID string, speculations []entity.SpeculationInfo) (retErr error) { - op := metrics.Begin(s.scope, "update_speculations") +// Update overwrites the paths of an existing speculation tree, identified by +// speculationTree.BatchID. Returns ErrNotFound if the speculation tree is not found. +func (s *speculationTreeStore) Update(ctx context.Context, speculationTree entity.SpeculationTree) (retErr error) { + op := metrics.Begin(s.scope, "update") defer func() { op.Complete(retErr) }() - speculationsJSON, err := json.Marshal(speculations) + speculationsJSON, err := json.Marshal(speculationTree.Paths) if err != nil { - return fmt.Errorf("failed to marshal speculations batchID=%s for UpdateSpeculations: %w", batchID, err) + return fmt.Errorf("failed to marshal paths batchID=%s for Update: %w", speculationTree.BatchID, err) } result, err := s.db.ExecContext(ctx, "UPDATE speculation_tree SET speculations = ? WHERE batch_id = ?", - speculationsJSON, batchID, + speculationsJSON, speculationTree.BatchID, ) if err != nil { - return fmt.Errorf("failed to update speculations for batchID=%q: %w", batchID, err) + return fmt.Errorf("failed to update speculation tree for batchID=%q: %w", speculationTree.BatchID, err) } rowsAffected, err := result.RowsAffected() if err != nil { - return fmt.Errorf("failed to get rows affected from update for batchID=%q: %w", batchID, err) + return fmt.Errorf("failed to get rows affected from update for batchID=%q: %w", speculationTree.BatchID, err) } if rowsAffected != 1 { - return storage.WrapNotFound(fmt.Errorf("speculation tree entity batchID=%s", batchID)) + return storage.WrapNotFound(fmt.Errorf("speculation tree entity batchID=%s", speculationTree.BatchID)) } return nil diff --git a/submitqueue/extension/storage/speculation_tree_store.go b/submitqueue/extension/storage/speculation_tree_store.go index 0b021bba..b282d60b 100644 --- a/submitqueue/extension/storage/speculation_tree_store.go +++ b/submitqueue/extension/storage/speculation_tree_store.go @@ -32,7 +32,7 @@ type SpeculationTreeStore interface { // Returns ErrAlreadyExists if the entry already exists. Create(ctx context.Context, speculationTree entity.SpeculationTree) error - // UpdateSpeculations updates the speculations of a speculation tree. - // Returns ErrNotFound if the speculation tree is not found. - UpdateSpeculations(ctx context.Context, batchID string, speculations []entity.SpeculationInfo) error + // Update overwrites the paths of an existing speculation tree, identified by + // speculationTree.BatchID. Returns ErrNotFound if the speculation tree is not found. + Update(ctx context.Context, speculationTree entity.SpeculationTree) error } diff --git a/submitqueue/orchestrator/controller/build/build.go b/submitqueue/orchestrator/controller/build/build.go index 4606d923..5491bfee 100644 --- a/submitqueue/orchestrator/controller/build/build.go +++ b/submitqueue/orchestrator/controller/build/build.go @@ -143,7 +143,7 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r build := entity.Build{ ID: buildID.ID, BatchID: batch.ID, - SpeculationPath: entity.SpeculationPathInfo{Base: append([]string{}, batch.Dependencies...)}, + SpeculationPath: entity.SpeculationPath{Base: append([]string{}, batch.Dependencies...), Head: batch.ID}, Status: entity.BuildStatusAccepted, }