From c4f175ec46b336b1f577f51a805e99617d46dda1 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Thu, 30 Jul 2026 10:04:51 -0700 Subject: [PATCH] feat(storage): speculation path set store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? The speculate run reads every head's path set at the start of a run and persists what it funds at the end. `entity.SpeculationPathSet` exists, but nothing can store one. ### What? Adds `storage.SpeculationPathSetStore`, keyed on the head batch ID, with a MySQL implementation, schema, mocks, and contract tests. The mutation is a conditional whole-item put — `Update` replaces the set, guarded on version — because a conditional put on a key is the primitive every backend offers directly; a field-level update is the one shape non-SQL backends would have to emulate with a read-modify-write. Version arguments are explicit and the entity's own `Version` field is ignored, per the storage README. No secondary index: callers that want a queue's live sets enumerate the heads from the batch listing they already hold and read each set by key. ## Test Plan ✅ `bazel test //submitqueue/extension/storage/...` — sqlmock coverage of get/create/update and every error contract, plus a pin that `Update` ignores the entity's `Version` field. ✅ `bazel test //test/integration/submitqueue/extension/storage/mysql:go_default_test` — against real MySQL: round-trip (each stored ID still hashes its stored path), missing head, duplicate create, and a compare-and-swap race where the loser must not restore the path the winner dropped. ✅ `make fmt`, `make gazelle`, `make mocks` --- submitqueue/extension/storage/BUILD.bazel | 1 + .../extension/storage/mock/BUILD.bazel | 1 + .../mock/speculation_path_set_store_mock.go | 85 +++++ .../extension/storage/mock/storage_mock.go | 14 + .../extension/storage/mysql/BUILD.bazel | 2 + .../mysql/schema/speculation_path_set.sql | 6 + .../mysql/speculation_path_set_store.go | 134 ++++++++ .../mysql/speculation_path_set_store_test.go | 310 ++++++++++++++++++ .../extension/storage/mysql/storage.go | 51 +-- .../storage/speculation_path_set_store.go | 56 ++++ submitqueue/extension/storage/storage.go | 3 + .../submitqueue/extension/storage/suite.go | 106 ++++++ 12 files changed, 747 insertions(+), 22 deletions(-) create mode 100644 submitqueue/extension/storage/mock/speculation_path_set_store_mock.go create mode 100644 submitqueue/extension/storage/mysql/schema/speculation_path_set.sql create mode 100644 submitqueue/extension/storage/mysql/speculation_path_set_store.go create mode 100644 submitqueue/extension/storage/mysql/speculation_path_set_store_test.go create mode 100644 submitqueue/extension/storage/speculation_path_set_store.go diff --git a/submitqueue/extension/storage/BUILD.bazel b/submitqueue/extension/storage/BUILD.bazel index 2c1b3bb9..25b945be 100644 --- a/submitqueue/extension/storage/BUILD.bazel +++ b/submitqueue/extension/storage/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "request_store.go", "request_summary_store.go", "request_uri_store.go", + "speculation_path_set_store.go", "storage.go", ], importpath = "github.com/uber/submitqueue/submitqueue/extension/storage", diff --git a/submitqueue/extension/storage/mock/BUILD.bazel b/submitqueue/extension/storage/mock/BUILD.bazel index 8f7c1f8b..986a3d4d 100644 --- a/submitqueue/extension/storage/mock/BUILD.bazel +++ b/submitqueue/extension/storage/mock/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "request_store_mock.go", "request_summary_store_mock.go", "request_uri_store_mock.go", + "speculation_path_set_store_mock.go", "storage_mock.go", ], importpath = "github.com/uber/submitqueue/submitqueue/extension/storage/mock", diff --git a/submitqueue/extension/storage/mock/speculation_path_set_store_mock.go b/submitqueue/extension/storage/mock/speculation_path_set_store_mock.go new file mode 100644 index 00000000..f6d9fc59 --- /dev/null +++ b/submitqueue/extension/storage/mock/speculation_path_set_store_mock.go @@ -0,0 +1,85 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: speculation_path_set_store.go +// +// Generated by this command: +// +// mockgen -source=speculation_path_set_store.go -destination=mock/speculation_path_set_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" +) + +// MockSpeculationPathSetStore is a mock of SpeculationPathSetStore interface. +type MockSpeculationPathSetStore struct { + ctrl *gomock.Controller + recorder *MockSpeculationPathSetStoreMockRecorder + isgomock struct{} +} + +// MockSpeculationPathSetStoreMockRecorder is the mock recorder for MockSpeculationPathSetStore. +type MockSpeculationPathSetStoreMockRecorder struct { + mock *MockSpeculationPathSetStore +} + +// NewMockSpeculationPathSetStore creates a new mock instance. +func NewMockSpeculationPathSetStore(ctrl *gomock.Controller) *MockSpeculationPathSetStore { + mock := &MockSpeculationPathSetStore{ctrl: ctrl} + mock.recorder = &MockSpeculationPathSetStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSpeculationPathSetStore) EXPECT() *MockSpeculationPathSetStoreMockRecorder { + return m.recorder +} + +// Create mocks base method. +func (m *MockSpeculationPathSetStore) Create(ctx context.Context, set entity.SpeculationPathSet) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Create", ctx, set) + ret0, _ := ret[0].(error) + return ret0 +} + +// Create indicates an expected call of Create. +func (mr *MockSpeculationPathSetStoreMockRecorder) Create(ctx, set any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockSpeculationPathSetStore)(nil).Create), ctx, set) +} + +// Get mocks base method. +func (m *MockSpeculationPathSetStore) Get(ctx context.Context, head string) (entity.SpeculationPathSet, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Get", ctx, head) + ret0, _ := ret[0].(entity.SpeculationPathSet) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Get indicates an expected call of Get. +func (mr *MockSpeculationPathSetStoreMockRecorder) Get(ctx, head any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockSpeculationPathSetStore)(nil).Get), ctx, head) +} + +// Update mocks base method. +func (m *MockSpeculationPathSetStore) Update(ctx context.Context, set entity.SpeculationPathSet, oldVersion, newVersion int32) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Update", ctx, set, oldVersion, newVersion) + ret0, _ := ret[0].(error) + return ret0 +} + +// Update indicates an expected call of Update. +func (mr *MockSpeculationPathSetStoreMockRecorder) Update(ctx, set, oldVersion, newVersion any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockSpeculationPathSetStore)(nil).Update), ctx, set, oldVersion, newVersion) +} diff --git a/submitqueue/extension/storage/mock/storage_mock.go b/submitqueue/extension/storage/mock/storage_mock.go index d32eb70a..3f73c35e 100644 --- a/submitqueue/extension/storage/mock/storage_mock.go +++ b/submitqueue/extension/storage/mock/storage_mock.go @@ -193,3 +193,17 @@ func (mr *MockStorageMockRecorder) GetRequestURIStore() *gomock.Call { mr.mock.ctrl.T.Helper() return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestURIStore", reflect.TypeOf((*MockStorage)(nil).GetRequestURIStore)) } + +// GetSpeculationPathSetStore mocks base method. +func (m *MockStorage) GetSpeculationPathSetStore() storage.SpeculationPathSetStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetSpeculationPathSetStore") + ret0, _ := ret[0].(storage.SpeculationPathSetStore) + return ret0 +} + +// GetSpeculationPathSetStore indicates an expected call of GetSpeculationPathSetStore. +func (mr *MockStorageMockRecorder) GetSpeculationPathSetStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSpeculationPathSetStore", reflect.TypeOf((*MockStorage)(nil).GetSpeculationPathSetStore)) +} diff --git a/submitqueue/extension/storage/mysql/BUILD.bazel b/submitqueue/extension/storage/mysql/BUILD.bazel index af189d37..bf1c93c3 100644 --- a/submitqueue/extension/storage/mysql/BUILD.bazel +++ b/submitqueue/extension/storage/mysql/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "request_store.go", "request_summary_store.go", "request_uri_store.go", + "speculation_path_set_store.go", "storage.go", ], importpath = "github.com/uber/submitqueue/submitqueue/extension/storage/mysql", @@ -39,6 +40,7 @@ go_test( "request_store_test.go", "request_summary_store_test.go", "request_uri_store_test.go", + "speculation_path_set_store_test.go", "storage_test.go", ], embed = [":go_default_library"], diff --git a/submitqueue/extension/storage/mysql/schema/speculation_path_set.sql b/submitqueue/extension/storage/mysql/schema/speculation_path_set.sql new file mode 100644 index 00000000..32dd2813 --- /dev/null +++ b/submitqueue/extension/storage/mysql/schema/speculation_path_set.sql @@ -0,0 +1,6 @@ +CREATE TABLE IF NOT EXISTS speculation_path_set ( + head VARCHAR(255) NOT NULL, + paths JSON NOT NULL, + version INT NOT NULL, + PRIMARY KEY (head) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/submitqueue/extension/storage/mysql/speculation_path_set_store.go b/submitqueue/extension/storage/mysql/speculation_path_set_store.go new file mode 100644 index 00000000..2db49848 --- /dev/null +++ b/submitqueue/extension/storage/mysql/speculation_path_set_store.go @@ -0,0 +1,134 @@ +// 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" + "encoding/json" + "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 speculationPathSetStore struct { + db *sql.DB + scope tally.Scope +} + +// NewSpeculationPathSetStore creates a new MySQL-backed SpeculationPathSetStore. +func NewSpeculationPathSetStore(db *sql.DB, scope tally.Scope) storage.SpeculationPathSetStore { + return &speculationPathSetStore{db: db, scope: scope} +} + +// Get retrieves a head's path set, where head is the head batch's ID. +// Returns ErrNotFound if the head has no set. +func (s *speculationPathSetStore) Get(ctx context.Context, head string) (ret entity.SpeculationPathSet, retErr error) { + op := metrics.Begin(s.scope, "get", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + var set entity.SpeculationPathSet + var pathsJSON []byte + + err := s.db.QueryRowContext(ctx, + "SELECT head, paths, version FROM speculation_path_set WHERE head = ?", + head, + ).Scan(&set.Head, &pathsJSON, &set.Version) + + if errors.Is(err, sql.ErrNoRows) { + return entity.SpeculationPathSet{}, storage.WrapNotFound(err) + } + if err != nil { + return entity.SpeculationPathSet{}, fmt.Errorf("failed to get speculation path set entity head=%s from the database: %w", head, err) + } + + if err := json.Unmarshal(pathsJSON, &set.Paths); err != nil { + return entity.SpeculationPathSet{}, fmt.Errorf("failed to unmarshal paths for speculation path set entity head=%s from the database: %w", head, err) + } + + return set, nil +} + +// Create stores a head's first path set. Returns ErrAlreadyExists if the head already has one. +func (s *speculationPathSetStore) Create(ctx context.Context, set entity.SpeculationPathSet) (retErr error) { + op := metrics.Begin(s.scope, "create", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + pathsJSON, err := json.Marshal(set.Paths) + if err != nil { + return fmt.Errorf("failed to marshal paths head=%s for Create speculation path set entity: %w", set.Head, err) + } + + _, err = s.db.ExecContext(ctx, + "INSERT INTO speculation_path_set (head, paths, version) VALUES (?, ?, ?)", + set.Head, pathsJSON, set.Version, + ) + if err != nil { + var mysqlErr *mysql.MySQLError + if errors.As(err, &mysqlErr) && mysqlErr.Number == mysqlErrDuplicateEntry { + return fmt.Errorf("speculation path set entity head=%s: %w", set.Head, storage.ErrAlreadyExists) + } + return fmt.Errorf("failed to insert speculation path set entity head=%s: %w", set.Head, err) + } + + return nil +} + +// Update replaces the stored set and writes newVersion if the persisted version matches +// oldVersion. If versions do not match, returns ErrVersionMismatch. set.Version is ignored: +// version arithmetic is owned by the caller and this is a pure conditional write. +func (s *speculationPathSetStore) Update(ctx context.Context, set entity.SpeculationPathSet, oldVersion, newVersion int32) (retErr error) { + op := metrics.Begin(s.scope, "update", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + pathsJSON, err := json.Marshal(set.Paths) + if err != nil { + return fmt.Errorf("failed to marshal paths head=%s for Update speculation path set entity: %w", set.Head, err) + } + + result, err := s.db.ExecContext(ctx, + "UPDATE speculation_path_set SET paths = ?, version = ? WHERE head = ? AND version = ?", + pathsJSON, newVersion, set.Head, oldVersion, + ) + if err != nil { + return fmt.Errorf( + "failed to update speculation path set for head=%q oldVersion=%d newVersion=%d: %w", + set.Head, oldVersion, newVersion, err, + ) + } + + rowsAffected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf( + "failed to get rows affected from update for head=%q oldVersion=%d newVersion=%d: %w", + set.Head, oldVersion, newVersion, err, + ) + } + + if rowsAffected != 1 { + return fmt.Errorf( + "version mismatch for speculation path set update: head=%q expected_version=%d: %w", + set.Head, oldVersion, storage.ErrVersionMismatch, + ) + } + + return nil +} diff --git a/submitqueue/extension/storage/mysql/speculation_path_set_store_test.go b/submitqueue/extension/storage/mysql/speculation_path_set_store_test.go new file mode 100644 index 00000000..9b39ace0 --- /dev/null +++ b/submitqueue/extension/storage/mysql/speculation_path_set_store_test.go @@ -0,0 +1,310 @@ +// 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" + "encoding/json" + "fmt" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/go-sql-driver/mysql" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +func setupSpeculationPathSetStoreTest(t *testing.T) (*sql.DB, sqlmock.Sqlmock, storage.SpeculationPathSetStore) { + t.Helper() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + + store := NewSpeculationPathSetStore(db, testMetrics()) + + return db, mock, store +} + +// testPathSet returns a two-path set for the given head: one assuming its +// dependency succeeds, one assuming it fails. +func testPathSet(head string) entity.SpeculationPathSet { + const dep = "monorepo/batch/1" + succeeds := entity.SpeculationPath{ + Head: head, + Dependencies: []entity.PathDependency{{Batch: dep, Assumption: entity.DependencyAssumptionSucceeds}}, + } + fails := entity.SpeculationPath{ + Head: head, + Dependencies: []entity.PathDependency{{Batch: dep, Assumption: entity.DependencyAssumptionFails}}, + } + return entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{ + { + ID: succeeds.ID(), + Path: succeeds, + Status: entity.SpeculationPathStatusBuilding, + Attempt: 1, + Version: 1, + CreatedAtMs: 1000, + UpdatedAtMs: 2000, + }, + { + ID: fails.ID(), + Path: fails, + Status: entity.SpeculationPathStatusPending, + Attempt: 1, + Version: 1, + CreatedAtMs: 1000, + UpdatedAtMs: 1000, + }, + }, + Version: 3, + } +} + +func TestSpeculationPathSetStore_Get(t *testing.T) { + want := testPathSet("monorepo/batch/2") + pathsJSON, err := json.Marshal(want.Paths) + require.NoError(t, err) + + tests := []struct { + name string + head string + setup func(mock sqlmock.Sqlmock) + want entity.SpeculationPathSet + wantErr bool + wantErrIs error + }{ + { + name: "found", + head: want.Head, + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"head", "paths", "version"}). + AddRow(want.Head, pathsJSON, want.Version) + mock.ExpectQuery("SELECT head, paths, version FROM speculation_path_set"). + WithArgs(want.Head). + WillReturnRows(rows) + }, + want: want, + }, + { + name: "not found", + head: "missing", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT head, paths, version FROM speculation_path_set"). + WithArgs("missing"). + WillReturnError(sql.ErrNoRows) + }, + wantErr: true, + wantErrIs: storage.ErrNotFound, + }, + { + name: "malformed paths json", + head: "corrupt", + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"head", "paths", "version"}). + AddRow("corrupt", []byte("{not json"), 1) + mock.ExpectQuery("SELECT head, paths, version FROM speculation_path_set"). + WithArgs("corrupt"). + WillReturnRows(rows) + }, + wantErr: true, + }, + { + name: "query error", + head: "bad", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT head, paths, version FROM speculation_path_set"). + WithArgs("bad"). + WillReturnError(fmt.Errorf("connection reset")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupSpeculationPathSetStoreTest(t) + defer db.Close() + + tt.setup(mock) + + got, err := store.Get(context.Background(), tt.head) + if tt.wantErr { + require.Error(t, err) + if tt.wantErrIs != nil { + assert.ErrorIs(t, err, tt.wantErrIs) + } + } else { + require.NoError(t, err) + assert.Equal(t, tt.want, got) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestSpeculationPathSetStore_Create(t *testing.T) { + set := testPathSet("monorepo/batch/2") + + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + wantErr bool + wantErrIs error + }{ + { + name: "success", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO speculation_path_set"). + WithArgs(set.Head, sqlmock.AnyArg(), set.Version). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "duplicate head returns ErrAlreadyExists", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO speculation_path_set"). + WithArgs(set.Head, sqlmock.AnyArg(), set.Version). + WillReturnError(&mysql.MySQLError{Number: mysqlErrDuplicateEntry}) + }, + wantErr: true, + wantErrIs: storage.ErrAlreadyExists, + }, + { + name: "other exec error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO speculation_path_set"). + WithArgs(set.Head, sqlmock.AnyArg(), set.Version). + WillReturnError(fmt.Errorf("connection reset")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupSpeculationPathSetStoreTest(t) + defer db.Close() + + tt.setup(mock) + + err := store.Create(context.Background(), set) + if tt.wantErr { + require.Error(t, err) + if tt.wantErrIs != nil { + assert.ErrorIs(t, err, tt.wantErrIs) + } + } else { + require.NoError(t, err) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestSpeculationPathSetStore_Update(t *testing.T) { + const oldVersion, newVersion = int32(3), int32(4) + set := testPathSet("monorepo/batch/2") + + tests := []struct { + name string + setup func(mock sqlmock.Sqlmock) + wantErr bool + wantErrIs error + }{ + { + name: "success", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE speculation_path_set"). + WithArgs(sqlmock.AnyArg(), newVersion, set.Head, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "version mismatch", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE speculation_path_set"). + WithArgs(sqlmock.AnyArg(), newVersion, set.Head, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 0)) + }, + wantErr: true, + wantErrIs: storage.ErrVersionMismatch, + }, + { + name: "exec error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE speculation_path_set"). + WithArgs(sqlmock.AnyArg(), newVersion, set.Head, oldVersion). + WillReturnError(fmt.Errorf("connection reset")) + }, + wantErr: true, + }, + { + name: "rows affected error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE speculation_path_set"). + WithArgs(sqlmock.AnyArg(), newVersion, set.Head, oldVersion). + WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error"))) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupSpeculationPathSetStoreTest(t) + defer db.Close() + + tt.setup(mock) + + err := store.Update(context.Background(), set, oldVersion, newVersion) + if tt.wantErr { + require.Error(t, err) + if tt.wantErrIs != nil { + assert.ErrorIs(t, err, tt.wantErrIs) + } + } else { + require.NoError(t, err) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +// TestSpeculationPathSetStore_UpdateIgnoresEntityVersion pins the documented +// contract that the guard and the written value come from the explicit +// arguments, not from the entity — a caller that forgot to refresh set.Version +// must not accidentally write it. +func TestSpeculationPathSetStore_UpdateIgnoresEntityVersion(t *testing.T) { + db, mock, store := setupSpeculationPathSetStoreTest(t) + defer db.Close() + + set := testPathSet("monorepo/batch/2") + set.Version = 99 // deliberately disagrees with both arguments + + const oldVersion, newVersion = int32(3), int32(4) + mock.ExpectExec("UPDATE speculation_path_set"). + WithArgs(sqlmock.AnyArg(), newVersion, set.Head, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + + require.NoError(t, store.Update(context.Background(), set, oldVersion, newVersion)) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/submitqueue/extension/storage/mysql/storage.go b/submitqueue/extension/storage/mysql/storage.go index 053c4686..2a9b1953 100644 --- a/submitqueue/extension/storage/mysql/storage.go +++ b/submitqueue/extension/storage/mysql/storage.go @@ -28,33 +28,35 @@ import ( const mysqlErrDuplicateEntry = 1062 type mysqlStorage struct { - db *sql.DB - requestStore storage.RequestStore - requestBatchStore storage.RequestBatchStore - changeStore storage.ChangeStore - batchStore storage.BatchStore - batchDependentStore storage.BatchDependentStore - buildStore storage.BuildStore - requestLogStore storage.RequestLogStore - requestSummaryStore storage.RequestSummaryStore - requestQueueStore storage.RequestQueueSummaryStore - requestURIStore storage.RequestURIStore + db *sql.DB + requestStore storage.RequestStore + requestBatchStore storage.RequestBatchStore + changeStore storage.ChangeStore + batchStore storage.BatchStore + batchDependentStore storage.BatchDependentStore + buildStore storage.BuildStore + speculationPathSetStore storage.SpeculationPathSetStore + requestLogStore storage.RequestLogStore + requestSummaryStore storage.RequestSummaryStore + requestQueueStore storage.RequestQueueSummaryStore + requestURIStore storage.RequestURIStore } // 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")), - requestBatchStore: NewRequestBatchStore(db, scope.SubScope("request_batch_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")), - requestLogStore: NewRequestLogStore(db, scope.SubScope("request_log_store")), - requestSummaryStore: NewRequestSummaryStore(db, scope.SubScope("request_summary_store")), - requestQueueStore: NewRequestQueueSummaryStore(db, scope.SubScope("request_queue_summary_store")), - requestURIStore: NewRequestURIStore(db, scope.SubScope("request_uri_store")), + db: db, + requestStore: NewRequestStore(db, scope.SubScope("request_store")), + requestBatchStore: NewRequestBatchStore(db, scope.SubScope("request_batch_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")), + speculationPathSetStore: NewSpeculationPathSetStore(db, scope.SubScope("speculation_path_set_store")), + requestLogStore: NewRequestLogStore(db, scope.SubScope("request_log_store")), + requestSummaryStore: NewRequestSummaryStore(db, scope.SubScope("request_summary_store")), + requestQueueStore: NewRequestQueueSummaryStore(db, scope.SubScope("request_queue_summary_store")), + requestURIStore: NewRequestURIStore(db, scope.SubScope("request_uri_store")), }, nil } @@ -88,6 +90,11 @@ func (f *mysqlStorage) GetBuildStore() storage.BuildStore { return f.buildStore } +// GetSpeculationPathSetStore returns the MySQL-backed SpeculationPathSetStore. +func (f *mysqlStorage) GetSpeculationPathSetStore() storage.SpeculationPathSetStore { + return f.speculationPathSetStore +} + // GetRequestLogStore returns the MySQL-backed RequestLogStore. func (f *mysqlStorage) GetRequestLogStore() storage.RequestLogStore { return f.requestLogStore diff --git a/submitqueue/extension/storage/speculation_path_set_store.go b/submitqueue/extension/storage/speculation_path_set_store.go new file mode 100644 index 00000000..7746cb2c --- /dev/null +++ b/submitqueue/extension/storage/speculation_path_set_store.go @@ -0,0 +1,56 @@ +// 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_set_store.go -destination=mock/speculation_path_set_store_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// SpeculationPathSetStore persists one head batch's chosen speculation paths. +// +// A set is keyed by its head batch ID and versioned as a whole: every path in +// it shares that head, and the set is the unit of both replacement and +// optimistic locking. There is no lookup by anything but the head — callers +// that need a queue's live sets enumerate the heads from the batch listing they +// already hold and read each set by key. +type SpeculationPathSetStore interface { + // Get retrieves a head's path set, where head is the head batch's ID. + // Returns ErrNotFound if the head has no set yet, which is the normal state + // for a batch nothing has speculated on. + Get(ctx context.Context, head string) (entity.SpeculationPathSet, error) + + // Create stores a head's first path set. + // Returns ErrAlreadyExists if the head already has one. + Create(ctx context.Context, set entity.SpeculationPathSet) error + + // Update replaces the stored set with set and writes newVersion, but only if + // the persisted version still matches oldVersion. If it does not, returns + // ErrVersionMismatch and writes nothing. + // + // The whole entity goes in rather than the fields being changed: a set is + // replaced wholesale, so this is a conditional put on a key — the primitive + // every backend offers directly, instead of a field-level update each + // non-SQL backend would have to emulate with a read-modify-write. + // + // set.Version is ignored. oldVersion is the guard and newVersion is the + // value written, so version arithmetic stays with the caller: compute + // newVersion, call, and assign it to the in-memory set only once this + // returns nil. + Update(ctx context.Context, set entity.SpeculationPathSet, oldVersion, newVersion int32) error +} diff --git a/submitqueue/extension/storage/storage.go b/submitqueue/extension/storage/storage.go index c2233298..fbdaee24 100644 --- a/submitqueue/extension/storage/storage.go +++ b/submitqueue/extension/storage/storage.go @@ -64,6 +64,9 @@ type Storage interface { // GetBuildStore returns the BuildStore instance. GetBuildStore() BuildStore + // GetSpeculationPathSetStore returns the SpeculationPathSetStore instance. + GetSpeculationPathSetStore() SpeculationPathSetStore + // GetRequestLogStore returns the RequestLogStore instance. GetRequestLogStore() RequestLogStore diff --git a/test/integration/submitqueue/extension/storage/suite.go b/test/integration/submitqueue/extension/storage/suite.go index 2bcd3dce..9b8c1eb2 100644 --- a/test/integration/submitqueue/extension/storage/suite.go +++ b/test/integration/submitqueue/extension/storage/suite.go @@ -749,3 +749,109 @@ func (s *StorageContractSuite) TestStorage_RequestURIListIsBoundedAndOrdered() { require.NoError(t, err) assert.Empty(t, empty) } + +// speculationPathSet builds a two-path set for head over one dependency: one +// path assuming the dependency succeeds, one assuming it fails. +func speculationPathSet(head, dep string) entity.SpeculationPathSet { + succeeds := entity.SpeculationPath{ + Head: head, + Dependencies: []entity.PathDependency{{Batch: dep, Assumption: entity.DependencyAssumptionSucceeds}}, + } + fails := entity.SpeculationPath{ + Head: head, + Dependencies: []entity.PathDependency{{Batch: dep, Assumption: entity.DependencyAssumptionFails}}, + } + return entity.SpeculationPathSet{ + Head: head, + Paths: []entity.SpeculationPathEntry{ + { + ID: succeeds.ID(), + Path: succeeds, + Status: entity.SpeculationPathStatusPending, + Attempt: 1, + Version: 1, + CreatedAtMs: 1000, + UpdatedAtMs: 1000, + }, + { + ID: fails.ID(), + Path: fails, + Status: entity.SpeculationPathStatusBuilding, + Attempt: 2, + Version: 3, + CreatedAtMs: 1000, + UpdatedAtMs: 2000, + }, + }, + Version: 1, + } +} + +// TestStorage_SpeculationPathSetCreateAndGet tests that a set round-trips whole, +// including each path's assumptions — a path's identity hashes them, so an +// encoding that dropped or reordered them would silently change every path ID. +func (s *StorageContractSuite) TestStorage_SpeculationPathSetCreateAndGet() { + t := s.T() + ctx := s.ctx + store := s.storage.GetSpeculationPathSetStore() + + want := speculationPathSet("sps/head/1", "sps/dep/1") + require.NoError(t, store.Create(ctx, want)) + + got, err := store.Get(ctx, want.Head) + require.NoError(t, err) + assert.Equal(t, want, got) + + for _, p := range got.Paths { + assert.Equal(t, p.Path.ID(), p.ID, "stored ID must still equal the hash of the stored path") + } +} + +// TestStorage_SpeculationPathSetNotFound tests reading a head nothing has +// speculated on yet — the normal state for a freshly created batch. +func (s *StorageContractSuite) TestStorage_SpeculationPathSetNotFound() { + t := s.T() + + _, err := s.storage.GetSpeculationPathSetStore().Get(s.ctx, "sps/head/nonexistent") + assert.ErrorIs(t, err, storage.ErrNotFound) +} + +// TestStorage_SpeculationPathSetCreateDuplicate tests that a head gets at most one set. +func (s *StorageContractSuite) TestStorage_SpeculationPathSetCreateDuplicate() { + t := s.T() + ctx := s.ctx + store := s.storage.GetSpeculationPathSetStore() + + set := speculationPathSet("sps/head/duplicate", "sps/dep/1") + require.NoError(t, store.Create(ctx, set)) + assert.ErrorIs(t, store.Create(ctx, set), storage.ErrAlreadyExists) +} + +// TestStorage_SpeculationPathSetOptimisticLocking tests the compare-and-swap +// contract. Speculate, build, and buildsignal all write the same row, so a +// loser must be rejected outright rather than clobbering the winner's paths. +func (s *StorageContractSuite) TestStorage_SpeculationPathSetOptimisticLocking() { + t := s.T() + ctx := s.ctx + store := s.storage.GetSpeculationPathSetStore() + + set := speculationPathSet("sps/head/cas", "sps/dep/1") + require.NoError(t, store.Create(ctx, set)) + + // Winner: replaces the set under the version it read. + winner := set + winner.Paths = winner.Paths[:1] + winner.Paths[0].Status = entity.SpeculationPathStatusPassed + require.NoError(t, store.Update(ctx, winner, 1, 2)) + + // Loser: still holds version 1, so its write is rejected. + loser := set + loser.Paths[0].Status = entity.SpeculationPathStatusFailed + assert.ErrorIs(t, store.Update(ctx, loser, 1, 2), storage.ErrVersionMismatch) + + got, err := store.Get(ctx, set.Head) + require.NoError(t, err) + assert.Equal(t, int32(2), got.Version) + require.Len(t, got.Paths, 1, "the losing write must not have restored the dropped path") + assert.Equal(t, entity.SpeculationPathStatusPassed, got.Paths[0].Status) +}