From 293ad914345aa9521e84289f85ead0d07c2c2b5a Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 9 Jul 2026 12:25:50 -0700 Subject: [PATCH] =?UTF-8?q?feat(storage,entity):=20speculation=20path?= =?UTF-8?q?=E2=86=92build=20mapping=20store?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? Integrating speculation needs the controllers to answer "which build belongs to this path" in both directions, without ever looking a row up by non-key attribute — the storage contract stays get/put-by-key so any KV backend can satisfy it. The build system mints its own build identifiers, and those stay the build store's primary key: the runner's ID is the natural name for the build row. What's missing is a durable link between a tree path and its build that exists independent of any in-flight message, hung on the path identity (SpeculationPathInfo.ID) introduced at the bottom of this stack. ### What? `entity.SpeculationPathBuild` is the new mapping entity ({PathID, BuildID, BatchID, Version, CreatedAt}, named after its `speculation_path_build` table) with a `SpeculationPathBuildStore` (Create/Get keyed by PathID): the forward path→build lookup, written only by the build controller, at most one build per path with ErrAlreadyExists making the existing row the truth on races. BatchID makes the row self-describing without parsing PathID's format; Version follows the repo's optimistic-locking convention (write-once today, reserved for future conditional re-pointing flows). `entity.Build` keeps the runner-minted `ID` as its primary key and gains `SpeculationPathID` as a plain column — the reverse build→path lookup; the previously embedded `SpeculationPath` value is dropped as a redundant denormalized copy of what the tree already stores. `SpeculationPath.Equal` (order-sensitive Base + Head) provides structural path identity for the few controller spots where only structure can identify a path (deduplicating enumerator output, carrying entries over across re-enumeration). `entity.QueueID` with ToBytes/QueueIDFromBytes mirrors the BatchID payload pattern for queue-scoped stages. MySQL gains the `speculation_path_build` table and the build table swaps `speculation_path`/`runner_id` for `speculation_path_id`; schema files are picked up automatically (the schema dir is globbed by both Bazel and the testutil ApplySchema helper). ## Test Plan ✅ `bazel test //submitqueue/entity/... //submitqueue/extension/storage/... //submitqueue/orchestrator/...` and the storage integration suite exercises Create/Get round-trips and the duplicate-create race. --- submitqueue/entity/BUILD.bazel | 4 + submitqueue/entity/build.go | 16 ++-- submitqueue/entity/build_test.go | 51 +++++------ submitqueue/entity/queue.go | 37 ++++++++ submitqueue/entity/queue_test.go | 72 ++++++++++++++++ submitqueue/entity/speculation_path_build.go | 43 ++++++++++ submitqueue/entity/speculation_tree.go | 28 ++++++- submitqueue/entity/speculation_tree_test.go | 75 +++++++++++++++++ submitqueue/extension/storage/BUILD.bazel | 1 + .../extension/storage/mock/BUILD.bazel | 1 + .../mock/speculation_path_build_store_mock.go | 71 ++++++++++++++++ .../extension/storage/mock/storage_mock.go | 14 ++++ .../extension/storage/mysql/BUILD.bazel | 1 + .../extension/storage/mysql/build_store.go | 21 ++--- .../extension/storage/mysql/schema/build.sql | 2 +- .../mysql/schema/speculation_path_build.sql | 8 ++ .../mysql/speculation_path_build_store.go | 84 +++++++++++++++++++ .../extension/storage/mysql/storage.go | 39 +++++---- .../storage/speculation_path_build_store.go | 41 +++++++++ submitqueue/extension/storage/storage.go | 3 + .../orchestrator/controller/build/build.go | 7 +- .../controller/build/build_test.go | 1 - .../submitqueue/extension/storage/suite.go | 84 +++++++++++++++++++ 23 files changed, 625 insertions(+), 79 deletions(-) create mode 100644 submitqueue/entity/queue.go create mode 100644 submitqueue/entity/queue_test.go create mode 100644 submitqueue/entity/speculation_path_build.go create mode 100644 submitqueue/entity/speculation_tree_test.go create mode 100644 submitqueue/extension/storage/mock/speculation_path_build_store_mock.go create mode 100644 submitqueue/extension/storage/mysql/schema/speculation_path_build.sql create mode 100644 submitqueue/extension/storage/mysql/speculation_path_build_store.go create mode 100644 submitqueue/extension/storage/speculation_path_build_store.go diff --git a/submitqueue/entity/BUILD.bazel b/submitqueue/entity/BUILD.bazel index 47f52b4d..4b291ada 100644 --- a/submitqueue/entity/BUILD.bazel +++ b/submitqueue/entity/BUILD.bazel @@ -14,9 +14,11 @@ go_library( "land_request.go", "merge_result.go", "push_result.go", + "queue.go", "queue_config.go", "request.go", "request_log.go", + "speculation_path_build.go", "speculation_tree.go", ], importpath = "github.com/uber/submitqueue/submitqueue/entity", @@ -34,8 +36,10 @@ go_test( "build_test.go", "cancel_request_test.go", "land_request_test.go", + "queue_test.go", "request_log_test.go", "request_test.go", + "speculation_tree_test.go", ], embed = [":go_default_library"], deps = [ diff --git a/submitqueue/entity/build.go b/submitqueue/entity/build.go index 2d8db24e..3cf5bf39 100644 --- a/submitqueue/entity/build.go +++ b/submitqueue/entity/build.go @@ -56,16 +56,18 @@ func (s BuildStatus) IsTerminal() bool { // 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 { - // ID represents the build ID. It is the responsibility of a build management system to ensure - // that this is unique. + // ID is the identifier minted by the queue's build runner when the build + // is triggered; this is the primary storage key. ID string // BatchID is the batch for which this build is scheduled. 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. Its Head is the batch being verified (equal to - // BatchID) and its Base is the assumed-good prefix of predecessor batches. - SpeculationPath SpeculationPath + // SpeculationPathID is the ID of the speculation-tree path this build + // verifies (SpeculationPathInfo.ID). The path's structure (Base/Head) is + // not embedded here — it lives on the tree entry and is looked up via the + // tree (SpeculationPathInfo.Path). This field enables the reverse lookup + // from a build row to its path; the forward lookup (path->build) lives in + // the separate SpeculationPathBuild mapping (see speculation_path_build.go). + SpeculationPathID string // Status represents the state of the build lifecycle this build is in. Status BuildStatus } diff --git a/submitqueue/entity/build_test.go b/submitqueue/entity/build_test.go index 755443d4..48cc6764 100644 --- a/submitqueue/entity/build_test.go +++ b/submitqueue/entity/build_test.go @@ -68,13 +68,10 @@ func TestBuildStatus_IsTerminal(t *testing.T) { func TestBuild_ToBytes(t *testing.T) { build := Build{ - ID: "build-1", - BatchID: "batch-1", - SpeculationPath: SpeculationPath{ - Base: []string{"batch-0", "batch-prev"}, - Head: "batch-1", - }, - Status: BuildStatusAccepted, + ID: "build-1", + BatchID: "batch-1", + Status: BuildStatusAccepted, + SpeculationPathID: "path-1", } data, err := build.ToBytes() @@ -86,17 +83,15 @@ func TestBuild_ToBytes(t *testing.T) { assert.Contains(t, jsonStr, "build-1") assert.Contains(t, jsonStr, "batch-1") assert.Contains(t, jsonStr, "accepted") + assert.Contains(t, jsonStr, "path-1") } func TestBuildFromBytes(t *testing.T) { original := Build{ - ID: "build-42", - BatchID: "batch-7", - SpeculationPath: SpeculationPath{ - Base: []string{"batch-5", "batch-6"}, - Head: "batch-7", - }, - Status: BuildStatusAccepted, + ID: "build-42", + BatchID: "batch-7", + Status: BuildStatusAccepted, + SpeculationPathID: "path-42", } // Serialize @@ -110,8 +105,8 @@ func TestBuildFromBytes(t *testing.T) { // Verify all fields match 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.Status, deserialized.Status) + assert.Equal(t, original.SpeculationPathID, deserialized.SpeculationPathID) } func TestBuildFromBytes_InvalidJSON(t *testing.T) { @@ -139,19 +134,16 @@ func TestBuild_SerializationRoundTrip(t *testing.T) { build Build }{ { - name: "accepted build with speculation path", + name: "accepted build with speculation path id", build: Build{ - ID: "build-100", - BatchID: "batch-50", - SpeculationPath: SpeculationPath{ - Base: []string{"batch-48", "batch-49"}, - Head: "batch-50", - }, - Status: BuildStatusAccepted, + ID: "build-100", + BatchID: "batch-50", + Status: BuildStatusAccepted, + SpeculationPathID: "path-100", }, }, { - name: "succeeded build with no speculation base", + name: "succeeded build with no speculation path id", build: Build{ ID: "build-200", BatchID: "batch-60", @@ -161,13 +153,10 @@ func TestBuild_SerializationRoundTrip(t *testing.T) { { name: "failed build", build: Build{ - ID: "build-300", - BatchID: "batch-70", - SpeculationPath: SpeculationPath{ - Base: []string{"batch-65"}, - Head: "batch-70", - }, - Status: BuildStatusFailed, + ID: "build-300", + BatchID: "batch-70", + Status: BuildStatusFailed, + SpeculationPathID: "path-300", }, }, } diff --git a/submitqueue/entity/queue.go b/submitqueue/entity/queue.go new file mode 100644 index 00000000..b7a680dd --- /dev/null +++ b/submitqueue/entity/queue.go @@ -0,0 +1,37 @@ +// 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 entity + +import "encoding/json" + +// QueueID is the queue-message payload for queue-scoped pipeline stages. It +// carries only the queue name; consumers resolve the state they need from +// storage. +type QueueID struct { + // Name is the merge-queue name the message targets. + Name string `json:"name"` +} + +// ToBytes serializes the QueueID to JSON bytes for queue message payload. +func (q QueueID) ToBytes() ([]byte, error) { + return json.Marshal(q) +} + +// QueueIDFromBytes deserializes a QueueID from JSON bytes. +func QueueIDFromBytes(data []byte) (QueueID, error) { + var qid QueueID + err := json.Unmarshal(data, &qid) + return qid, err +} diff --git a/submitqueue/entity/queue_test.go b/submitqueue/entity/queue_test.go new file mode 100644 index 00000000..543ea07b --- /dev/null +++ b/submitqueue/entity/queue_test.go @@ -0,0 +1,72 @@ +// 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 entity + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestQueueID_SerializationRoundTrip(t *testing.T) { + tests := []struct { + name string + queueID QueueID + }{ + { + name: "simple queue name", + queueID: QueueID{Name: "queueA"}, + }, + { + name: "another queue name", + queueID: QueueID{Name: "queueB"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := tt.queueID.ToBytes() + require.NoError(t, err) + + deserialized, err := QueueIDFromBytes(data) + require.NoError(t, err) + + assert.Equal(t, tt.queueID, deserialized) + }) + } +} + +func TestQueueIDFromBytes_InvalidJSON(t *testing.T) { + _, err := QueueIDFromBytes([]byte(`{"invalid": json"}`)) + assert.Error(t, err) +} + +func TestQueueIDFromBytes_EmptyJSON(t *testing.T) { + queueID, err := QueueIDFromBytes([]byte(`{}`)) + require.NoError(t, err) + + assert.Empty(t, queueID.Name) +} + +func TestQueueIDFromBytes_EmptyBytes(t *testing.T) { + _, err := QueueIDFromBytes([]byte{}) + assert.Error(t, err) +} + +func TestQueueIDFromBytes_NilBytes(t *testing.T) { + _, err := QueueIDFromBytes(nil) + assert.Error(t, err) +} diff --git a/submitqueue/entity/speculation_path_build.go b/submitqueue/entity/speculation_path_build.go new file mode 100644 index 00000000..2755f2f7 --- /dev/null +++ b/submitqueue/entity/speculation_path_build.go @@ -0,0 +1,43 @@ +// 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 entity + +// SpeculationPathBuild is the path->build mapping: given a speculation path's +// ID, it records the build that path resolved to. It is written by the build +// controller when it triggers a build for a speculation path; this is the +// forward lookup (path->build), while the reverse lookup (build->path) is +// Build.SpeculationPathID. Today a row is write-once — only Create and Get +// exist, no Update — so Version is always 1 on every row; the field is +// reserved for a future repair/re-trigger flow that conditionally re-points a +// path to a newer build without a schema migration (not used yet). +type SpeculationPathBuild struct { + // PathID is the speculation path's ID (SpeculationPathInfo.ID). It is the + // primary key of this mapping. + PathID string + // BuildID is the runner-minted build ID (Build.ID) this path resolved to. + BuildID string + // BatchID is the batch whose speculation tree contains this path. It + // makes the row self-describing without parsing PathID's internal format. + BatchID string + // Version is the version of the object, used for optimistic locking: + // updates are conditional on the persisted version matching the caller's + // expected version. Versioning starts at 1; version arithmetic is owned + // by the controller, the store performs a pure conditional write. Not + // used yet — reserved for a future conditional-update flow. + Version int32 + // CreatedAt is the creation time of this mapping, in milliseconds since + // epoch. + CreatedAt int64 +} diff --git a/submitqueue/entity/speculation_tree.go b/submitqueue/entity/speculation_tree.go index 40d01ad0..39e7c46d 100644 --- a/submitqueue/entity/speculation_tree.go +++ b/submitqueue/entity/speculation_tree.go @@ -29,6 +29,27 @@ type SpeculationPath struct { Head string } +// Equal reports whether p and other are structurally the same speculation +// path. It is true iff Head matches and Base has the same elements in the same +// order — Base order is the build order and is significant. The controller uses +// it where only structure can identify a path (deduplicating enumerator output, +// carrying entries over across re-enumeration); everything else references a +// persisted path by its assigned ID (SpeculationPathInfo.ID). +func (p SpeculationPath) Equal(other SpeculationPath) bool { + if p.Head != other.Head { + return false + } + if len(p.Base) != len(other.Base) { + return false + } + for i := range p.Base { + if p.Base[i] != other.Base[i] { + return false + } + } + return true +} + // 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 decision seams (selector, prioritizer) @@ -108,6 +129,7 @@ type SpeculationPathInfo struct { // meaning — never parse it. Everything outside the tree names a path by this // ID: seam outputs (path scores, path decisions) and durable links from // other entities all refer to it rather than restating the Base/Head split. + // The path→build mapping row is keyed by it (SpeculationPathBuild.PathID). ID string // Path is the Base/Head split this entry covers. Immutable: it identifies // the entry and never changes after the path is first persisted. @@ -124,9 +146,9 @@ type SpeculationPathInfo struct { // only by the controller; read by the decision seams (scorer, selector, // prioritizer). Status SpeculationPathStatus - // BuildID links this path to its build. Updateable: empty until a build - // signal confirms the build and the controller records it (Prioritized -> - // Building); the controller never knows the ID at send time. + // BuildID holds the runner-minted build identifier (also the build store's + // primary key) for this path. Updateable: it is empty until the speculate + // controller's reconcile stamps it once a build exists for this path. BuildID string } diff --git a/submitqueue/entity/speculation_tree_test.go b/submitqueue/entity/speculation_tree_test.go new file mode 100644 index 00000000..8b6b0428 --- /dev/null +++ b/submitqueue/entity/speculation_tree_test.go @@ -0,0 +1,75 @@ +// 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 entity + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSpeculationPath_Equal(t *testing.T) { + tests := []struct { + name string + path SpeculationPath + other SpeculationPath + equal bool + }{ + { + name: "equal paths", + path: SpeculationPath{Base: []string{"q/batch/1", "q/batch/2"}, Head: "q/batch/3"}, + other: SpeculationPath{Base: []string{"q/batch/1", "q/batch/2"}, Head: "q/batch/3"}, + equal: true, + }, + { + name: "different head", + path: SpeculationPath{Base: []string{"q/batch/1"}, Head: "q/batch/2"}, + other: SpeculationPath{Base: []string{"q/batch/1"}, Head: "q/batch/3"}, + equal: false, + }, + { + name: "different base order", + path: SpeculationPath{Base: []string{"q/batch/1", "q/batch/2"}, Head: "q/batch/3"}, + other: SpeculationPath{Base: []string{"q/batch/2", "q/batch/1"}, Head: "q/batch/3"}, + equal: false, + }, + { + name: "different base length", + path: SpeculationPath{Base: []string{"q/batch/1"}, Head: "q/batch/3"}, + other: SpeculationPath{Base: []string{"q/batch/1", "q/batch/2"}, Head: "q/batch/3"}, + equal: false, + }, + { + name: "both empty", + path: SpeculationPath{}, + other: SpeculationPath{}, + equal: true, + }, + { + name: "nil base equals empty base", + path: SpeculationPath{Base: nil, Head: "q/batch/1"}, + other: SpeculationPath{Base: []string{}, Head: "q/batch/1"}, + equal: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.equal, tt.path.Equal(tt.other)) + // Equal must be symmetric. + assert.Equal(t, tt.equal, tt.other.Equal(tt.path)) + }) + } +} diff --git a/submitqueue/extension/storage/BUILD.bazel b/submitqueue/extension/storage/BUILD.bazel index aa02fefb..ccc2bcef 100644 --- a/submitqueue/extension/storage/BUILD.bazel +++ b/submitqueue/extension/storage/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "change_store.go", "request_log_store.go", "request_store.go", + "speculation_path_build_store.go", "speculation_tree_store.go", "storage.go", ], diff --git a/submitqueue/extension/storage/mock/BUILD.bazel b/submitqueue/extension/storage/mock/BUILD.bazel index 9224f547..97ceeec5 100644 --- a/submitqueue/extension/storage/mock/BUILD.bazel +++ b/submitqueue/extension/storage/mock/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "change_store_mock.go", "request_log_store_mock.go", "request_store_mock.go", + "speculation_path_build_store_mock.go", "speculation_tree_store_mock.go", "storage_mock.go", ], diff --git a/submitqueue/extension/storage/mock/speculation_path_build_store_mock.go b/submitqueue/extension/storage/mock/speculation_path_build_store_mock.go new file mode 100644 index 00000000..91e200db --- /dev/null +++ b/submitqueue/extension/storage/mock/speculation_path_build_store_mock.go @@ -0,0 +1,71 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: speculation_path_build_store.go +// +// Generated by this command: +// +// mockgen -source=speculation_path_build_store.go -destination=mock/speculation_path_build_store_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" + gomock "go.uber.org/mock/gomock" +) + +// MockSpeculationPathBuildStore is a mock of SpeculationPathBuildStore interface. +type MockSpeculationPathBuildStore struct { + ctrl *gomock.Controller + recorder *MockSpeculationPathBuildStoreMockRecorder + isgomock struct{} +} + +// MockSpeculationPathBuildStoreMockRecorder is the mock recorder for MockSpeculationPathBuildStore. +type MockSpeculationPathBuildStoreMockRecorder struct { + mock *MockSpeculationPathBuildStore +} + +// NewMockSpeculationPathBuildStore creates a new mock instance. +func NewMockSpeculationPathBuildStore(ctrl *gomock.Controller) *MockSpeculationPathBuildStore { + mock := &MockSpeculationPathBuildStore{ctrl: ctrl} + mock.recorder = &MockSpeculationPathBuildStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSpeculationPathBuildStore) EXPECT() *MockSpeculationPathBuildStoreMockRecorder { + return m.recorder +} + +// Create mocks base method. +func (m *MockSpeculationPathBuildStore) Create(ctx context.Context, pathBuild entity.SpeculationPathBuild) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Create", ctx, pathBuild) + ret0, _ := ret[0].(error) + return ret0 +} + +// Create indicates an expected call of Create. +func (mr *MockSpeculationPathBuildStoreMockRecorder) Create(ctx, pathBuild any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockSpeculationPathBuildStore)(nil).Create), ctx, pathBuild) +} + +// Get mocks base method. +func (m *MockSpeculationPathBuildStore) Get(ctx context.Context, pathID string) (entity.SpeculationPathBuild, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Get", ctx, pathID) + ret0, _ := ret[0].(entity.SpeculationPathBuild) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Get indicates an expected call of Get. +func (mr *MockSpeculationPathBuildStoreMockRecorder) Get(ctx, pathID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockSpeculationPathBuildStore)(nil).Get), ctx, pathID) +} diff --git a/submitqueue/extension/storage/mock/storage_mock.go b/submitqueue/extension/storage/mock/storage_mock.go index 4133bc2a..c144bc50 100644 --- a/submitqueue/extension/storage/mock/storage_mock.go +++ b/submitqueue/extension/storage/mock/storage_mock.go @@ -138,6 +138,20 @@ func (mr *MockStorageMockRecorder) GetRequestStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestStore", reflect.TypeOf((*MockStorage)(nil).GetRequestStore)) } +// GetSpeculationPathBuildStore mocks base method. +func (m *MockStorage) GetSpeculationPathBuildStore() storage.SpeculationPathBuildStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetSpeculationPathBuildStore") + ret0, _ := ret[0].(storage.SpeculationPathBuildStore) + return ret0 +} + +// GetSpeculationPathBuildStore indicates an expected call of GetSpeculationPathBuildStore. +func (mr *MockStorageMockRecorder) GetSpeculationPathBuildStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSpeculationPathBuildStore", reflect.TypeOf((*MockStorage)(nil).GetSpeculationPathBuildStore)) +} + // GetSpeculationTreeStore mocks base method. func (m *MockStorage) GetSpeculationTreeStore() storage.SpeculationTreeStore { m.ctrl.T.Helper() diff --git a/submitqueue/extension/storage/mysql/BUILD.bazel b/submitqueue/extension/storage/mysql/BUILD.bazel index 5cb431e4..cc04b8c6 100644 --- a/submitqueue/extension/storage/mysql/BUILD.bazel +++ b/submitqueue/extension/storage/mysql/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "change_store.go", "request_log_store.go", "request_store.go", + "speculation_path_build_store.go", "speculation_tree_store.go", "storage.go", ], diff --git a/submitqueue/extension/storage/mysql/build_store.go b/submitqueue/extension/storage/mysql/build_store.go index c5ea9fc3..b50b8168 100644 --- a/submitqueue/extension/storage/mysql/build_store.go +++ b/submitqueue/extension/storage/mysql/build_store.go @@ -17,7 +17,6 @@ package mysql import ( "context" "database/sql" - "encoding/json" "errors" "fmt" @@ -45,12 +44,11 @@ func (s *buildStore) Get(ctx context.Context, id string) (ret entity.Build, retE defer func() { op.Complete(retErr) }() var build entity.Build - var speculationPathJSON []byte err := s.db.QueryRowContext(ctx, - "SELECT id, batch_id, speculation_path, status FROM build WHERE id = ?", + "SELECT id, batch_id, status, speculation_path_id FROM build WHERE id = ?", id, - ).Scan(&build.ID, &build.BatchID, &speculationPathJSON, &build.Status) + ).Scan(&build.ID, &build.BatchID, &build.Status, &build.SpeculationPathID) if errors.Is(err, sql.ErrNoRows) { return entity.Build{}, storage.WrapNotFound(err) @@ -59,10 +57,6 @@ func (s *buildStore) Get(ctx context.Context, id string) (ret entity.Build, retE return entity.Build{}, fmt.Errorf("failed to get build entity id=%s from the database: %w", id, err) } - if err := json.Unmarshal(speculationPathJSON, &build.SpeculationPath); err != nil { - return entity.Build{}, fmt.Errorf("failed to unmarshal speculation_path for build entity id=%s from the database: %w", id, err) - } - return build, nil } @@ -71,14 +65,9 @@ func (s *buildStore) Create(ctx context.Context, build entity.Build) (retErr err op := metrics.Begin(s.scope, "create") defer func() { op.Complete(retErr) }() - speculationPathJSON, err := json.Marshal(build.SpeculationPath) - if err != nil { - return fmt.Errorf("failed to marshal speculation_path id=%s for Create build entity: %w", build.ID, err) - } - - _, err = s.db.ExecContext(ctx, - "INSERT INTO build (id, batch_id, speculation_path, status) VALUES (?, ?, ?, ?)", - build.ID, build.BatchID, speculationPathJSON, build.Status, + _, err := s.db.ExecContext(ctx, + "INSERT INTO build (id, batch_id, status, speculation_path_id) VALUES (?, ?, ?, ?)", + build.ID, build.BatchID, build.Status, build.SpeculationPathID, ) if err != nil { var mysqlErr *mysql.MySQLError diff --git a/submitqueue/extension/storage/mysql/schema/build.sql b/submitqueue/extension/storage/mysql/schema/build.sql index 8aea8230..946e45e4 100644 --- a/submitqueue/extension/storage/mysql/schema/build.sql +++ b/submitqueue/extension/storage/mysql/schema/build.sql @@ -1,7 +1,7 @@ CREATE TABLE IF NOT EXISTS build ( id VARCHAR(255) NOT NULL, batch_id VARCHAR(255) NOT NULL, - speculation_path JSON NOT NULL, status VARCHAR(64) NOT NULL, + speculation_path_id VARCHAR(255) NOT NULL, PRIMARY KEY (id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/submitqueue/extension/storage/mysql/schema/speculation_path_build.sql b/submitqueue/extension/storage/mysql/schema/speculation_path_build.sql new file mode 100644 index 00000000..b4a6e817 --- /dev/null +++ b/submitqueue/extension/storage/mysql/schema/speculation_path_build.sql @@ -0,0 +1,8 @@ +CREATE TABLE IF NOT EXISTS speculation_path_build ( + path_id VARCHAR(255) NOT NULL, + build_id VARCHAR(255) NOT NULL, + batch_id VARCHAR(255) NOT NULL, + version INT NOT NULL, + created_at BIGINT NOT NULL, + PRIMARY KEY (path_id) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/submitqueue/extension/storage/mysql/speculation_path_build_store.go b/submitqueue/extension/storage/mysql/speculation_path_build_store.go new file mode 100644 index 00000000..edb9d7e8 --- /dev/null +++ b/submitqueue/extension/storage/mysql/speculation_path_build_store.go @@ -0,0 +1,84 @@ +// 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 mysql + +import ( + "context" + "database/sql" + "errors" + "fmt" + + "github.com/go-sql-driver/mysql" + "github.com/uber-go/tally" + + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +type speculationPathBuildStore struct { + db *sql.DB + scope tally.Scope +} + +// NewSpeculationPathBuildStore creates a new MySQL-backed SpeculationPathBuildStore. +func NewSpeculationPathBuildStore(db *sql.DB, scope tally.Scope) storage.SpeculationPathBuildStore { + return &speculationPathBuildStore{db: db, scope: scope} +} + +// Create creates a new path->build mapping. Returns ErrAlreadyExists if a +// mapping for the given PathID already exists. The caller-supplied Version is +// persisted as-is; version arithmetic is owned by the controller, not the store. +func (s *speculationPathBuildStore) Create(ctx context.Context, pathBuild entity.SpeculationPathBuild) (retErr error) { + op := metrics.Begin(s.scope, "create") + defer func() { op.Complete(retErr) }() + + _, err := s.db.ExecContext(ctx, + "INSERT INTO speculation_path_build (path_id, build_id, batch_id, version, created_at) VALUES (?, ?, ?, ?, ?)", + pathBuild.PathID, pathBuild.BuildID, pathBuild.BatchID, pathBuild.Version, pathBuild.CreatedAt, + ) + if err != nil { + var mysqlErr *mysql.MySQLError + if errors.As(err, &mysqlErr) && mysqlErr.Number == 1062 { + return fmt.Errorf("path build mapping pathID=%s: %w", pathBuild.PathID, storage.ErrAlreadyExists) + } + return fmt.Errorf("failed to insert path build mapping pathID=%s: %w", pathBuild.PathID, err) + } + + return nil +} + +// Get retrieves the path->build mapping for the given path ID. Returns +// ErrNotFound if no mapping exists for pathID. +func (s *speculationPathBuildStore) Get(ctx context.Context, pathID string) (ret entity.SpeculationPathBuild, retErr error) { + op := metrics.Begin(s.scope, "get") + defer func() { op.Complete(retErr) }() + + var pb entity.SpeculationPathBuild + + err := s.db.QueryRowContext(ctx, + "SELECT path_id, build_id, batch_id, version, created_at FROM speculation_path_build WHERE path_id = ?", + pathID, + ).Scan(&pb.PathID, &pb.BuildID, &pb.BatchID, &pb.Version, &pb.CreatedAt) + + if errors.Is(err, sql.ErrNoRows) { + return entity.SpeculationPathBuild{}, storage.WrapNotFound(err) + } + if err != nil { + return entity.SpeculationPathBuild{}, fmt.Errorf("failed to get path build mapping pathID=%s from the database: %w", pathID, err) + } + + return pb, nil +} diff --git a/submitqueue/extension/storage/mysql/storage.go b/submitqueue/extension/storage/mysql/storage.go index 4ba0cf41..40e128f2 100644 --- a/submitqueue/extension/storage/mysql/storage.go +++ b/submitqueue/extension/storage/mysql/storage.go @@ -24,27 +24,29 @@ import ( ) type mysqlStorage struct { - db *sql.DB - requestStore storage.RequestStore - changeStore storage.ChangeStore - batchStore storage.BatchStore - batchDependentStore storage.BatchDependentStore - buildStore storage.BuildStore - speculationTreeStore storage.SpeculationTreeStore - requestLogStore storage.RequestLogStore + db *sql.DB + requestStore storage.RequestStore + changeStore storage.ChangeStore + batchStore storage.BatchStore + batchDependentStore storage.BatchDependentStore + buildStore storage.BuildStore + speculationPathBuildStore storage.SpeculationPathBuildStore + speculationTreeStore storage.SpeculationTreeStore + requestLogStore storage.RequestLogStore } // NewStorage creates a new MySQL storage. func NewStorage(db *sql.DB, scope tally.Scope) (storage.Storage, error) { return &mysqlStorage{ - db: db, - requestStore: NewRequestStore(db, scope.SubScope("request_store")), - changeStore: NewChangeStore(db, scope.SubScope("change_store")), - batchStore: NewBatchStore(db, scope.SubScope("batch_store")), - batchDependentStore: NewBatchDependentStore(db, scope.SubScope("batch_dependent_store")), - buildStore: NewBuildStore(db, scope.SubScope("build_store")), - speculationTreeStore: NewSpeculationTreeStore(db, scope.SubScope("speculation_tree_store")), - requestLogStore: NewRequestLogStore(db, scope.SubScope("request_log_store")), + db: db, + requestStore: NewRequestStore(db, scope.SubScope("request_store")), + changeStore: NewChangeStore(db, scope.SubScope("change_store")), + batchStore: NewBatchStore(db, scope.SubScope("batch_store")), + batchDependentStore: NewBatchDependentStore(db, scope.SubScope("batch_dependent_store")), + buildStore: NewBuildStore(db, scope.SubScope("build_store")), + speculationPathBuildStore: NewSpeculationPathBuildStore(db, scope.SubScope("speculation_path_build_store")), + speculationTreeStore: NewSpeculationTreeStore(db, scope.SubScope("speculation_tree_store")), + requestLogStore: NewRequestLogStore(db, scope.SubScope("request_log_store")), }, nil } @@ -73,6 +75,11 @@ func (f *mysqlStorage) GetBuildStore() storage.BuildStore { return f.buildStore } +// GetSpeculationPathBuildStore returns the MySQL-backed SpeculationPathBuildStore. +func (f *mysqlStorage) GetSpeculationPathBuildStore() storage.SpeculationPathBuildStore { + return f.speculationPathBuildStore +} + // GetSpeculationTreeStore returns the MySQL-backed SpeculationTreeStore. func (f *mysqlStorage) GetSpeculationTreeStore() storage.SpeculationTreeStore { return f.speculationTreeStore diff --git a/submitqueue/extension/storage/speculation_path_build_store.go b/submitqueue/extension/storage/speculation_path_build_store.go new file mode 100644 index 00000000..64e82ab1 --- /dev/null +++ b/submitqueue/extension/storage/speculation_path_build_store.go @@ -0,0 +1,41 @@ +// 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 storage + +//go:generate mockgen -source=speculation_path_build_store.go -destination=mock/speculation_path_build_store_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// SpeculationPathBuildStore is an interface that defines methods for managing +// the path->build mapping in the database. A row is written only by the build +// controller, when it triggers a build for a speculation path, and is +// immutable once written: a path maps to at most one build. If the mapping +// for a path is lost (e.g. redelivery races), the build controller +// re-triggers a build and the new mapping wins — so Create on an existing +// PathID returns ErrAlreadyExists, and the caller treats the already-persisted +// row as truth rather than overwriting it. +type SpeculationPathBuildStore interface { + // Create creates a new path->build mapping. Returns ErrAlreadyExists if a + // mapping for the given PathID already exists. + Create(ctx context.Context, pathBuild entity.SpeculationPathBuild) error + + // Get retrieves the path->build mapping for the given path ID. Returns + // ErrNotFound if no mapping exists for pathID. + Get(ctx context.Context, pathID string) (entity.SpeculationPathBuild, error) +} diff --git a/submitqueue/extension/storage/storage.go b/submitqueue/extension/storage/storage.go index a02bef73..a1ec91ef 100644 --- a/submitqueue/extension/storage/storage.go +++ b/submitqueue/extension/storage/storage.go @@ -59,6 +59,9 @@ type Storage interface { // GetBuildStore returns the BuildStore instance. GetBuildStore() BuildStore + // GetSpeculationPathBuildStore returns the SpeculationPathBuildStore instance. + GetSpeculationPathBuildStore() SpeculationPathBuildStore + // GetSpeculationTreeStore returns the SpeculationTreeStore instance. GetSpeculationTreeStore() SpeculationTreeStore diff --git a/submitqueue/orchestrator/controller/build/build.go b/submitqueue/orchestrator/controller/build/build.go index 85f1f68e..8745f772 100644 --- a/submitqueue/orchestrator/controller/build/build.go +++ b/submitqueue/orchestrator/controller/build/build.go @@ -138,10 +138,9 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r } build := entity.Build{ - ID: buildID.ID, - BatchID: batch.ID, - SpeculationPath: entity.SpeculationPath{Base: append([]string{}, batch.Dependencies...), Head: batch.ID}, - Status: entity.BuildStatusAccepted, + ID: buildID.ID, + BatchID: batch.ID, + Status: entity.BuildStatusAccepted, } // Persist the initial Build snapshot so the buildsignal poll loop has a diff --git a/submitqueue/orchestrator/controller/build/build_test.go b/submitqueue/orchestrator/controller/build/build_test.go index 4b04c89b..7347ae29 100644 --- a/submitqueue/orchestrator/controller/build/build_test.go +++ b/submitqueue/orchestrator/controller/build/build_test.go @@ -216,7 +216,6 @@ func TestController_Process_TriggersWithBaseAndHead(t *testing.T) { assert.Equal(t, "build-xyz", created.ID) assert.Equal(t, headBatch.ID, created.BatchID) assert.Equal(t, entity.BuildStatusAccepted, created.Status) - assert.Equal(t, []string{depBatch.ID}, created.SpeculationPath.Base) assert.Equal(t, published.ID, created.ID) } diff --git a/test/integration/submitqueue/extension/storage/suite.go b/test/integration/submitqueue/extension/storage/suite.go index 933dd96e..94a46996 100644 --- a/test/integration/submitqueue/extension/storage/suite.go +++ b/test/integration/submitqueue/extension/storage/suite.go @@ -504,3 +504,87 @@ func (s *StorageContractSuite) TestStorage_SpeculationUpdateNotFound() { err := s.storage.GetSpeculationTreeStore().Update(ctx, tree.BatchID, tree.Version, tree.Version+1, tree.Paths) assert.ErrorIs(t, err, storage.ErrVersionMismatch, "Update for unknown batch should return ErrVersionMismatch") } + +// TestStorage_BuildCreateAndGet verifies a build round-trips through the +// store: ID is the runner-minted build identifier (assigned by the build +// runner, not derived from the path), SpeculationPathID links the build back +// to the speculation-tree path it verifies, and Get by ID finds the build +// with both round-tripped intact. +func (s *StorageContractSuite) TestStorage_BuildCreateAndGet() { + t := s.T() + ctx := s.ctx + const batchID = "q/batch/create-and-get" + const buildID = "runner/create-and-get/42" + const pathID = "spec/create-and-get/path-1" + + build := entity.Build{ + ID: buildID, + BatchID: batchID, + SpeculationPathID: pathID, + Status: entity.BuildStatusRunning, + } + + require.NoError(t, s.storage.GetBuildStore().Create(ctx, build)) + + got, err := s.storage.GetBuildStore().Get(ctx, buildID) + require.NoError(t, err, "Get by ID should find the build") + assert.Equal(t, build.ID, got.ID) + assert.Equal(t, build.BatchID, got.BatchID) + assert.Equal(t, build.Status, got.Status) + assert.Equal(t, build.SpeculationPathID, got.SpeculationPathID, "SpeculationPathID should round-trip") + assert.NotEqual(t, build.ID, build.SpeculationPathID, "build ID and speculation path ID are distinct identifiers") +} + +// TestStorage_SpeculationPathBuildCreateAndGet verifies a path->build mapping +// round-trips through the store unchanged. +func (s *StorageContractSuite) TestStorage_SpeculationPathBuildCreateAndGet() { + t := s.T() + ctx := s.ctx + + pb := entity.SpeculationPathBuild{ + PathID: "spec/path-build/create-get", + BuildID: "runner/path-build/create-get/1", + BatchID: "q/batch/path-build-create-get", + Version: 1, + CreatedAt: 1700000000000, + } + + require.NoError(t, s.storage.GetSpeculationPathBuildStore().Create(ctx, pb)) + + got, err := s.storage.GetSpeculationPathBuildStore().Get(ctx, pb.PathID) + require.NoError(t, err) + assert.Equal(t, pb, got, "path build mapping should round-trip unchanged") +} + +// TestStorage_SpeculationPathBuildCreateDuplicate verifies a repeated Create +// for the same PathID returns ErrAlreadyExists — a path maps to at most one +// build, and the first-written mapping wins. +func (s *StorageContractSuite) TestStorage_SpeculationPathBuildCreateDuplicate() { + t := s.T() + ctx := s.ctx + + pb := entity.SpeculationPathBuild{ + PathID: "spec/path-build/duplicate", + BuildID: "runner/path-build/duplicate/1", + BatchID: "q/batch/path-build-duplicate", + Version: 1, + CreatedAt: 1700000000000, + } + + require.NoError(t, s.storage.GetSpeculationPathBuildStore().Create(ctx, pb)) + + other := pb + other.BuildID = "runner/path-build/duplicate/2" + err := s.storage.GetSpeculationPathBuildStore().Create(ctx, other) + assert.ErrorIs(t, err, storage.ErrAlreadyExists, "duplicate create should return ErrAlreadyExists") +} + +// TestStorage_SpeculationPathBuildGetNotFound verifies Get for an unknown +// path ID returns ErrNotFound. +func (s *StorageContractSuite) TestStorage_SpeculationPathBuildGetNotFound() { + t := s.T() + ctx := s.ctx + + _, err := s.storage.GetSpeculationPathBuildStore().Get(ctx, "spec/path-build/nonexistent") + assert.ErrorIs(t, err, storage.ErrNotFound, "Get for unknown path ID should return ErrNotFound") +}