From 16081705fbca804a38ce9f5e0d31541110039f22 Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Wed, 5 Aug 2026 18:06:36 -0700 Subject: [PATCH] feat(storage): add QueueBatchStateStore and shared batch state helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Why? `BatchStore.GetByQueueAndStates` is the storage contract's only query-by-attribute, backed by `idx_queue_state` — the only secondary index in the entire schema. The storage README's key-value contract flags exactly this shape: a genuinely needed reverse lookup should be a first-class mapping store keyed by the lookup attribute, so any backend (SQL, DynamoDB, Bigtable) can serve it as a primary-key read. This PR adds that store and the shared helpers controllers will migrate onto; `speculate/run.go` already carries a TODO for this replacement. ### What? - `entity.QueueBatchState` — an advisory membership record filing an in-queue batch under a (queue, state) bucket, existing from batch creation until the batch exits the queue (through terminal states, hence not "active"). Also adds `entity.AllBatchStates()` for the future conclude-time sweep. - `storage.QueueBatchStateStore` — `List(queue, state)` / `Put` / `Delete`, all idempotent. The MySQL impl is backed by a new `queue_batch_state` table whose PK (queue, state, batch_id) *is* the lookup: listing a state bucket is a PK-prefix scan, no secondary index. - `submitqueue/core/batch` — the shared protocol primitives: `Transition` (batch CAS, then Put the new-bucket record before Deleting the old one, so a batch always has at least one record), `EnsureRecord` (idempotent repair for redelivery skip branches), and `ListByStates` (scan requested buckets, dedupe, hydrate by key with bounded concurrency, classify by the hydrated authoritative state — a stale record can misplace a batch but never misreport it). - Contract-suite coverage for the new store and a README pointer beside the `ChangeRecord` mapping-store example. Scoped to the foundation only: the store has no callers yet, so there is no runtime behavior change. Follow-ups migrate the six batch-transition sites and the two `GetByQueueAndStates` readers onto the helpers, add conclude-time record deletion, and then remove `GetByQueueAndStates` + `idx_queue_state`. ## Test Plan - ✅ `bazel test //submitqueue/entity:all //submitqueue/core/batch:all //submitqueue/extension/storage/...` - ✅ `bazel test //test/integration/submitqueue/extension/storage/mysql:go_default_test` (real MySQL; includes the new `TestStorage_QueueBatchStateRecordLifecycle` contract case) - ✅ `make tidy` / `make gazelle` / `make fmt` / `make mocks` — all stable (re-running produces no diffs) - Note: `make test` failures in `orchestrator/controller/batch` and `controller/cancel` are pre-existing — they reproduce at clean HEAD db51df4d in a fresh worktree. --- MODULE.bazel | 1 + go.mod | 2 +- submitqueue/core/batch/BUILD.bazel | 33 +++ submitqueue/core/batch/list.go | 89 ++++++++ submitqueue/core/batch/list_test.go | 150 ++++++++++++ submitqueue/core/batch/transition.go | 87 +++++++ submitqueue/core/batch/transition_test.go | 187 +++++++++++++++ submitqueue/entity/BUILD.bazel | 1 + submitqueue/entity/batch.go | 15 ++ submitqueue/entity/batch_test.go | 8 + submitqueue/entity/queue_batch_state.go | 40 ++++ submitqueue/extension/storage/BUILD.bazel | 1 + submitqueue/extension/storage/README.md | 2 +- .../extension/storage/mock/BUILD.bazel | 1 + .../mock/queue_batch_state_store_mock.go | 85 +++++++ .../extension/storage/mock/storage_mock.go | 14 ++ .../extension/storage/mysql/BUILD.bazel | 2 + .../storage/mysql/queue_batch_state_store.go | 91 ++++++++ .../mysql/queue_batch_state_store_test.go | 215 ++++++++++++++++++ .../mysql/schema/queue_batch_state.sql | 11 + .../extension/storage/mysql/storage.go | 51 +++-- .../storage/queue_batch_state_store.go | 52 +++++ submitqueue/extension/storage/storage.go | 3 + .../submitqueue/extension/storage/suite.go | 57 +++++ 24 files changed, 1174 insertions(+), 24 deletions(-) create mode 100644 submitqueue/core/batch/BUILD.bazel create mode 100644 submitqueue/core/batch/list.go create mode 100644 submitqueue/core/batch/list_test.go create mode 100644 submitqueue/core/batch/transition.go create mode 100644 submitqueue/core/batch/transition_test.go create mode 100644 submitqueue/entity/queue_batch_state.go create mode 100644 submitqueue/extension/storage/mock/queue_batch_state_store_mock.go create mode 100644 submitqueue/extension/storage/mysql/queue_batch_state_store.go create mode 100644 submitqueue/extension/storage/mysql/queue_batch_state_store_test.go create mode 100644 submitqueue/extension/storage/mysql/schema/queue_batch_state.sql create mode 100644 submitqueue/extension/storage/queue_batch_state_store.go diff --git a/MODULE.bazel b/MODULE.bazel index 8202f504..e25908eb 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -65,6 +65,7 @@ use_repo( "org_golang_google_grpc_cmd_protoc_gen_go_grpc", "org_golang_google_protobuf", "org_golang_x_oauth2", + "org_golang_x_sync", "org_uber_go_fx", "org_uber_go_mock", "org_uber_go_yarpc", diff --git a/go.mod b/go.mod index ba21f2e4..644ce7e9 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( go.uber.org/yarpc v1.81.0 go.uber.org/zap v1.27.1 golang.org/x/oauth2 v0.34.0 + golang.org/x/sync v0.19.0 google.golang.org/grpc v1.68.1 google.golang.org/protobuf v1.36.10 gopkg.in/yaml.v3 v3.0.1 @@ -46,7 +47,6 @@ require ( golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 // indirect golang.org/x/mod v0.32.0 // indirect golang.org/x/net v0.49.0 // indirect - golang.org/x/sync v0.19.0 // indirect golang.org/x/sys v0.40.0 // indirect golang.org/x/text v0.34.0 // indirect golang.org/x/tools v0.41.0 // indirect diff --git a/submitqueue/core/batch/BUILD.bazel b/submitqueue/core/batch/BUILD.bazel new file mode 100644 index 00000000..ef82b20a --- /dev/null +++ b/submitqueue/core/batch/BUILD.bazel @@ -0,0 +1,33 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = [ + "list.go", + "transition.go", + ], + importpath = "github.com/uber/submitqueue/submitqueue/core/batch", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/storage:go_default_library", + "@org_golang_x_sync//errgroup:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = [ + "list_test.go", + "transition_test.go", + ], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/storage:go_default_library", + "//submitqueue/extension/storage/mock:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/submitqueue/core/batch/list.go b/submitqueue/core/batch/list.go new file mode 100644 index 00000000..edbc0f64 --- /dev/null +++ b/submitqueue/core/batch/list.go @@ -0,0 +1,89 @@ +// Copyright (c) 2026 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 batch + +import ( + "context" + "fmt" + + "golang.org/x/sync/errgroup" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +// hydrateConcurrency bounds the parallel per-key batch reads a single +// ListByStates call issues while hydrating candidate IDs. +const hydrateConcurrency = 16 + +// ListByStates returns the queue's batches whose current state is one of the given +// states, read through the queue's membership records: each requested state bucket +// is listed, candidate IDs are deduplicated across buckets, every candidate is +// hydrated by key with bounded concurrency, and the result keeps only batches whose +// hydrated State is in states. Classification always uses the hydrated state — a +// record found in a stale bucket can therefore never misreport a batch, only route +// an extra read. Result order is unspecified. +// +// A candidate ID whose batch does not exist is returned as an error rather than +// skipped: batch rows are never deleted, so a dangling record means the store is +// inconsistent, not that the batch concluded. +func ListByStates(ctx context.Context, store storage.Storage, queue string, states []entity.BatchState) ([]entity.Batch, error) { + wanted := make(map[entity.BatchState]bool, len(states)) + seen := make(map[string]bool) + var ids []string + for _, state := range states { + if wanted[state] { + continue + } + wanted[state] = true + + records, err := store.GetQueueBatchStateStore().List(ctx, queue, state) + if err != nil { + return nil, fmt.Errorf("failed to list queue batch state records for queue %s state %s: %w", queue, state, err) + } + for _, record := range records { + if seen[record.BatchID] { + continue + } + seen[record.BatchID] = true + ids = append(ids, record.BatchID) + } + } + + hydrated := make([]entity.Batch, len(ids)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(hydrateConcurrency) + for i, id := range ids { + g.Go(func() error { + batch, err := store.GetBatchStore().Get(gctx, id) + if err != nil { + return fmt.Errorf("failed to get batch %s of queue %s: %w", id, queue, err) + } + hydrated[i] = batch + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, err + } + + var result []entity.Batch + for _, batch := range hydrated { + if wanted[batch.State] { + result = append(result, batch) + } + } + return result, nil +} diff --git a/submitqueue/core/batch/list_test.go b/submitqueue/core/batch/list_test.go new file mode 100644 index 00000000..6a632138 --- /dev/null +++ b/submitqueue/core/batch/list_test.go @@ -0,0 +1,150 @@ +// Copyright (c) 2026 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 batch + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" + storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" +) + +const testQueue = "monorepo" + +// record builds a QueueBatchState for testQueue. +func record(state entity.BatchState, batchID string) entity.QueueBatchState { + return entity.QueueBatchState{Queue: testQueue, State: state, BatchID: batchID} +} + +// batchIn builds a hydrated Batch for testQueue in the given state. +func batchIn(id string, state entity.BatchState) entity.Batch { + return entity.Batch{ID: id, Queue: testQueue, State: state, Version: 1} +} + +func TestListByStates(t *testing.T) { + storeErr := errors.New("storage failed") + + tests := map[string]struct { + states []entity.BatchState + setup func(*storagemock.MockBatchStore, *storagemock.MockQueueBatchStateStore) + want []entity.Batch + wantErr error + }{ + "empty states lists nothing": { + states: nil, + setup: func(*storagemock.MockBatchStore, *storagemock.MockQueueBatchStateStore) {}, + }, + "hydrates every bucket and dedupes across them": { + states: []entity.BatchState{entity.BatchStateCreated, entity.BatchStateSpeculating}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + // b2 appears in both buckets (mid-move duplicate): it must be hydrated + // and returned exactly once. + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1"), record(entity.BatchStateCreated, "b2")}, nil) + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateSpeculating). + Return([]entity.QueueBatchState{record(entity.BatchStateSpeculating, "b2"), record(entity.BatchStateSpeculating, "b3")}, nil) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(batchIn("b1", entity.BatchStateCreated), nil) + batchStore.EXPECT().Get(gomock.Any(), "b2").Return(batchIn("b2", entity.BatchStateSpeculating), nil) + batchStore.EXPECT().Get(gomock.Any(), "b3").Return(batchIn("b3", entity.BatchStateSpeculating), nil) + }, + want: []entity.Batch{ + batchIn("b1", entity.BatchStateCreated), + batchIn("b2", entity.BatchStateSpeculating), + batchIn("b3", entity.BatchStateSpeculating), + }, + }, + "classifies by hydrated state, not by bucket": { + states: []entity.BatchState{entity.BatchStateCreated}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + // A stale record files b1 under created, but the batch has moved on to + // speculating — a state outside the requested set, so it is dropped. + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1")}, nil) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(batchIn("b1", entity.BatchStateSpeculating), nil) + }, + }, + "stale bucket still surfaces a batch whose true state is requested": { + states: []entity.BatchState{entity.BatchStateCreated, entity.BatchStateSpeculating}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + // Only a stale created record exists for b1, but its hydrated state is + // speculating — requested, so the batch is returned under its true state. + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1")}, nil) + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateSpeculating). + Return(nil, nil) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(batchIn("b1", entity.BatchStateSpeculating), nil) + }, + want: []entity.Batch{batchIn("b1", entity.BatchStateSpeculating)}, + }, + "duplicate input states are listed once": { + states: []entity.BatchState{entity.BatchStateCreated, entity.BatchStateCreated}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1")}, nil). + Times(1) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(batchIn("b1", entity.BatchStateCreated), nil) + }, + want: []entity.Batch{batchIn("b1", entity.BatchStateCreated)}, + }, + "list failure surfaces": { + states: []entity.BatchState{entity.BatchStateCreated}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated).Return(nil, storeErr) + }, + wantErr: storeErr, + }, + "hydrate failure surfaces": { + states: []entity.BatchState{entity.BatchStateCreated}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1")}, nil) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(entity.Batch{}, storeErr) + }, + wantErr: storeErr, + }, + "dangling record is an error, not a skip": { + states: []entity.BatchState{entity.BatchStateCreated}, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().List(gomock.Any(), testQueue, entity.BatchStateCreated). + Return([]entity.QueueBatchState{record(entity.BatchStateCreated, "b1")}, nil) + batchStore.EXPECT().Get(gomock.Any(), "b1").Return(entity.Batch{}, storage.WrapNotFound(errors.New("no rows"))) + }, + wantErr: storage.ErrNotFound, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + mockStorage, mockBatchStore, mockRecordStore := testStores(t) + tt.setup(mockBatchStore, mockRecordStore) + + got, err := ListByStates(context.Background(), mockStorage, testQueue, tt.states) + if tt.wantErr != nil { + require.Error(t, err) + assert.ErrorIs(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.ElementsMatch(t, tt.want, got) + }) + } +} diff --git a/submitqueue/core/batch/transition.go b/submitqueue/core/batch/transition.go new file mode 100644 index 00000000..2f442501 --- /dev/null +++ b/submitqueue/core/batch/transition.go @@ -0,0 +1,87 @@ +// Copyright (c) 2026 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 batch provides the shared primitives for moving a batch through its +// lifecycle states while keeping the queue's per-state membership records +// (entity.QueueBatchState) in step. +// +// The records are advisory and the Batch entity is authoritative, so the +// primitives follow one protocol: +// +// - A transition CASes the batch first, then files the record under the new +// state before removing the one under the old state, so a batch always has +// at least one record while it is in the queue. +// - A crash between the CAS and the record move is repaired by the pipeline's +// at-least-once redelivery: the retry's "already in target state" branch +// calls EnsureRecord, and every record write is idempotent. +// - Readers treat records as candidate batch IDs only: they hydrate each +// batch by key and classify it by its own State, never by the bucket the +// record was found in, so a stale record can misplace a batch but never +// misreport it. +package batch + +import ( + "context" + "fmt" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +// Transition moves a batch to newState: it performs the optimistic-locking CAS on +// the batch (newVersion = Version+1, assigned in memory only after the store write +// succeeds), then re-files the queue's membership record — Put under newState first, +// Delete under the prior state after, so the batch is never without a record. The +// Delete is skipped when the state is unchanged. It returns the batch as last +// successfully written. +// +// A storage.ErrVersionMismatch from the CAS is returned wrapped (errors.Is works), +// with no record writes attempted, so callers keep their existing lost-race +// semantics. Any other non-nil error means the transition may have partially +// applied — the CAS may have committed with the record move incomplete — and the +// caller is expected to let redelivery retry; the retry's already-in-target-state +// branch repairs the record via EnsureRecord. +func Transition(ctx context.Context, store storage.Storage, batch entity.Batch, newState entity.BatchState) (entity.Batch, error) { + oldState := batch.State + newVersion := batch.Version + 1 + updated := batch + updated.State = newState + if err := store.GetBatchStore().Update(ctx, updated, batch.Version, newVersion); err != nil { + return batch, fmt.Errorf("failed to update batch %s state to %s: %w", batch.ID, newState, err) + } + updated.Version = newVersion + + record := entity.QueueBatchState{Queue: updated.Queue, State: newState, BatchID: updated.ID} + if err := store.GetQueueBatchStateStore().Put(ctx, record); err != nil { + return updated, fmt.Errorf("failed to put queue batch state record for batch %s under state %s: %w", updated.ID, newState, err) + } + if oldState != newState { + if err := store.GetQueueBatchStateStore().Delete(ctx, updated.Queue, oldState, updated.ID); err != nil { + return updated, fmt.Errorf("failed to delete queue batch state record for batch %s under state %s: %w", updated.ID, oldState, err) + } + } + return updated, nil +} + +// EnsureRecord idempotently files the batch under its current state bucket. It is +// the repair half of the transition protocol: idempotent redelivery branches that +// skip the CAS because the batch is already in the target state call this instead, +// covering a prior attempt that crashed between the CAS and the record move. +func EnsureRecord(ctx context.Context, store storage.Storage, batch entity.Batch) error { + record := entity.QueueBatchState{Queue: batch.Queue, State: batch.State, BatchID: batch.ID} + if err := store.GetQueueBatchStateStore().Put(ctx, record); err != nil { + return fmt.Errorf("failed to put queue batch state record for batch %s under state %s: %w", batch.ID, batch.State, err) + } + return nil +} diff --git a/submitqueue/core/batch/transition_test.go b/submitqueue/core/batch/transition_test.go new file mode 100644 index 00000000..3aba102e --- /dev/null +++ b/submitqueue/core/batch/transition_test.go @@ -0,0 +1,187 @@ +// Copyright (c) 2026 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 batch + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" + storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" +) + +// testStores wires a MockStorage whose batch and queue-batch-state accessors +// return the two mocks the tests set expectations on. +func testStores(t *testing.T) (*storagemock.MockStorage, *storagemock.MockBatchStore, *storagemock.MockQueueBatchStateStore) { + t.Helper() + + ctrl := gomock.NewController(t) + mockStorage := storagemock.NewMockStorage(ctrl) + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockRecordStore := storagemock.NewMockQueueBatchStateStore(ctrl) + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetQueueBatchStateStore().Return(mockRecordStore).AnyTimes() + return mockStorage, mockBatchStore, mockRecordStore +} + +func TestTransition(t *testing.T) { + base := entity.Batch{ + ID: "monorepo/batch/7", + Queue: "monorepo", + Contains: []string{"monorepo/1"}, + State: entity.BatchStateCreated, + Version: 3, + } + casTarget := base + casTarget.State = entity.BatchStateSpeculating + storeErr := errors.New("storage failed") + + tests := map[string]struct { + newState entity.BatchState + setup func(*storagemock.MockBatchStore, *storagemock.MockQueueBatchStateStore) + want entity.Batch + wantErr error + }{ + "success moves the record to the new bucket": { + newState: entity.BatchStateSpeculating, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + batchStore.EXPECT().Update(gomock.Any(), casTarget, int32(3), int32(4)).Return(nil) + recordStore.EXPECT().Put(gomock.Any(), entity.QueueBatchState{ + Queue: base.Queue, State: entity.BatchStateSpeculating, BatchID: base.ID, + }).Return(nil) + recordStore.EXPECT().Delete(gomock.Any(), base.Queue, entity.BatchStateCreated, base.ID).Return(nil) + }, + want: func() entity.Batch { + b := casTarget + b.Version = 4 + return b + }(), + }, + "same state skips the delete": { + newState: entity.BatchStateCreated, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + batchStore.EXPECT().Update(gomock.Any(), base, int32(3), int32(4)).Return(nil) + recordStore.EXPECT().Put(gomock.Any(), entity.QueueBatchState{ + Queue: base.Queue, State: entity.BatchStateCreated, BatchID: base.ID, + }).Return(nil) + }, + want: func() entity.Batch { + b := base + b.Version = 4 + return b + }(), + }, + "lost CAS returns version mismatch and writes no records": { + newState: entity.BatchStateSpeculating, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + batchStore.EXPECT().Update(gomock.Any(), casTarget, int32(3), int32(4)).Return(storage.ErrVersionMismatch) + }, + want: base, + wantErr: storage.ErrVersionMismatch, + }, + "put failure surfaces after a committed CAS": { + newState: entity.BatchStateSpeculating, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + batchStore.EXPECT().Update(gomock.Any(), casTarget, int32(3), int32(4)).Return(nil) + recordStore.EXPECT().Put(gomock.Any(), gomock.Any()).Return(storeErr) + }, + want: func() entity.Batch { + b := casTarget + b.Version = 4 + return b + }(), + wantErr: storeErr, + }, + "delete failure surfaces after the put": { + newState: entity.BatchStateSpeculating, + setup: func(batchStore *storagemock.MockBatchStore, recordStore *storagemock.MockQueueBatchStateStore) { + batchStore.EXPECT().Update(gomock.Any(), casTarget, int32(3), int32(4)).Return(nil) + recordStore.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil) + recordStore.EXPECT().Delete(gomock.Any(), base.Queue, entity.BatchStateCreated, base.ID).Return(storeErr) + }, + want: func() entity.Batch { + b := casTarget + b.Version = 4 + return b + }(), + wantErr: storeErr, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + mockStorage, mockBatchStore, mockRecordStore := testStores(t) + tt.setup(mockBatchStore, mockRecordStore) + + got, err := Transition(context.Background(), mockStorage, base, tt.newState) + if tt.wantErr != nil { + require.Error(t, err) + assert.ErrorIs(t, err, tt.wantErr) + } else { + require.NoError(t, err) + } + assert.Equal(t, tt.want, got) + }) + } +} + +func TestEnsureRecord(t *testing.T) { + batch := entity.Batch{ + ID: "monorepo/batch/7", + Queue: "monorepo", + State: entity.BatchStateMerging, + Version: 5, + } + storeErr := errors.New("storage failed") + + tests := map[string]struct { + setup func(*storagemock.MockQueueBatchStateStore) + wantErr error + }{ + "files the batch under its current state": { + setup: func(recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().Put(gomock.Any(), entity.QueueBatchState{ + Queue: batch.Queue, State: batch.State, BatchID: batch.ID, + }).Return(nil) + }, + }, + "put failure surfaces": { + setup: func(recordStore *storagemock.MockQueueBatchStateStore) { + recordStore.EXPECT().Put(gomock.Any(), gomock.Any()).Return(storeErr) + }, + wantErr: storeErr, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + mockStorage, _, mockRecordStore := testStores(t) + tt.setup(mockRecordStore) + + err := EnsureRecord(context.Background(), mockStorage, batch) + if tt.wantErr != nil { + assert.ErrorIs(t, err, tt.wantErr) + } else { + assert.NoError(t, err) + } + }) + } +} diff --git a/submitqueue/entity/BUILD.bazel b/submitqueue/entity/BUILD.bazel index c31bf21f..ae7cfd03 100644 --- a/submitqueue/entity/BUILD.bazel +++ b/submitqueue/entity/BUILD.bazel @@ -15,6 +15,7 @@ go_library( "list.go", "merge_result.go", "push_result.go", + "queue_batch_state.go", "queue_config.go", "request.go", "request_batch.go", diff --git a/submitqueue/entity/batch.go b/submitqueue/entity/batch.go index b360e329..e1b4431f 100644 --- a/submitqueue/entity/batch.go +++ b/submitqueue/entity/batch.go @@ -85,6 +85,21 @@ func IsBatchStateHalted(s BatchState) bool { return s.IsTerminal() || s == BatchStateCancelling } +// AllBatchStates returns every named batch state, in lifecycle order. +// BatchStateUnknown is excluded: it is the zero-value sentinel, not a state a batch can occupy. +func AllBatchStates() []BatchState { + return []BatchState{ + BatchStateCreating, + BatchStateCreated, + BatchStateSpeculating, + BatchStateMerging, + BatchStateSucceeded, + BatchStateFailed, + BatchStateCancelling, + BatchStateCancelled, + } +} + // ActiveBatchStates returns batch states eligible for active pipeline and cancellation lookups. // Creating is excluded because its reverse-index structure may still be incomplete. func ActiveBatchStates() []BatchState { diff --git a/submitqueue/entity/batch_test.go b/submitqueue/entity/batch_test.go index 0afc38f2..558d9514 100644 --- a/submitqueue/entity/batch_test.go +++ b/submitqueue/entity/batch_test.go @@ -62,6 +62,14 @@ func TestActiveBatchStates_ExcludesCreating(t *testing.T) { assert.NotContains(t, ActiveBatchStates(), BatchStateCreating) } +func TestAllBatchStates_SupersetOfStateSubsets(t *testing.T) { + all := AllBatchStates() + assert.NotContains(t, all, BatchStateUnknown) + assert.Subset(t, all, ActiveBatchStates()) + assert.Subset(t, all, DependencyBatchStates()) + assert.Subset(t, all, []BatchState{BatchStateSucceeded, BatchStateFailed, BatchStateCancelled}) +} + func TestDependencyBatchStates_ExcludesCreating(t *testing.T) { assert.NotContains(t, DependencyBatchStates(), BatchStateCreating) } diff --git a/submitqueue/entity/queue_batch_state.go b/submitqueue/entity/queue_batch_state.go new file mode 100644 index 00000000..1aa7ab1c --- /dev/null +++ b/submitqueue/entity/queue_batch_state.go @@ -0,0 +1,40 @@ +// Copyright (c) 2026 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 + +// QueueBatchState is a membership record filing one in-queue batch under one lifecycle +// state bucket of its queue. A record exists for every batch from its creation until it +// exits the queue — through terminal states, not just while in flight. +// +// The (Queue, State, BatchID) triple is the record's identity; records carry no other +// data and are never updated in place — a batch changes buckets by a record appearing +// under the new state and the old record disappearing. +// +// Records are advisory. The authoritative state is the State field of the Batch +// identified by BatchID; a record may transiently file a batch under a bucket the batch +// has already left, and a batch may transiently have records in more than one bucket. A +// batch is never without at least one record while it is in the queue. +type QueueBatchState struct { + // Queue is the name of the queue the batch belongs to. Queue name is defined in the + // configuration and should be unique within the system. + Queue string + + // State is the lifecycle state bucket this record files the batch under. Advisory: + // the batch's authoritative state lives on the Batch entity and may differ transiently. + State BatchState + + // BatchID is the globally unique identifier of the batch. Format: "/batch/". + BatchID string +} diff --git a/submitqueue/extension/storage/BUILD.bazel b/submitqueue/extension/storage/BUILD.bazel index 2c1b3bb9..2c134e5d 100644 --- a/submitqueue/extension/storage/BUILD.bazel +++ b/submitqueue/extension/storage/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "queue_batch_state_store.go", "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", diff --git a/submitqueue/extension/storage/README.md b/submitqueue/extension/storage/README.md index 2b3a3244..e9a6c366 100644 --- a/submitqueue/extension/storage/README.md +++ b/submitqueue/extension/storage/README.md @@ -48,7 +48,7 @@ Store interfaces are designed for the storage technology *space*, not for SQL (s **Domain state is often already the index.** Before adding any lookup, check whether an entity the caller already loads enumerates the children — an aggregate that references its parts by ID (e.g. a tree whose paths record their build identities) is the batch→children index, persisted and versioned as domain state. Duplicating that relationship as a database index adds a second source of truth for something the domain already owns. -**When neither applies, the reverse lookup is real — give it its own mapping store.** In the KV space there is no third mechanism: the only way to look up by an attribute is to make that attribute a primary key somewhere. So promote the relationship to a first-class mapping entity — keyed by the lookup attribute, written by the same flow that creates the source entity with idempotent puts, and rebuildable as a projection if it drifts. `ChangeRecord` is the in-repo example: it exists so "which requests claimed this change URI" is a by-key read on (queue, URI). Unlike a `KEY idx_*`, the relationship is visible in the contract and portable to any backend. +**When neither applies, the reverse lookup is real — give it its own mapping store.** In the KV space there is no third mechanism: the only way to look up by an attribute is to make that attribute a primary key somewhere. So promote the relationship to a first-class mapping entity — keyed by the lookup attribute, written by the same flow that creates the source entity with idempotent puts, and rebuildable as a projection if it drifts. `ChangeRecord` is the in-repo example: it exists so "which requests claimed this change URI" is a by-key read on (queue, URI). `QueueBatchState` is the same pattern for a mutable attribute: "which batches of this queue are in this state" is a by-key read on (queue, state), maintained as advisory records that move buckets alongside the batch's own state CAS (the shared primitives in `submitqueue/core/batch` own that protocol) — it exists to replace `BatchStore.GetByQueueAndStates`, the contract's one remaining query-by-attribute. Unlike a `KEY idx_*`, the relationship is visible in the contract and portable to any backend. ### Decision path diff --git a/submitqueue/extension/storage/mock/BUILD.bazel b/submitqueue/extension/storage/mock/BUILD.bazel index 8f7c1f8b..01aba47f 100644 --- a/submitqueue/extension/storage/mock/BUILD.bazel +++ b/submitqueue/extension/storage/mock/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store_mock.go", "build_store_mock.go", "change_store_mock.go", + "queue_batch_state_store_mock.go", "request_batch_store_mock.go", "request_log_store_mock.go", "request_queue_summary_store_mock.go", diff --git a/submitqueue/extension/storage/mock/queue_batch_state_store_mock.go b/submitqueue/extension/storage/mock/queue_batch_state_store_mock.go new file mode 100644 index 00000000..2b0c2e0f --- /dev/null +++ b/submitqueue/extension/storage/mock/queue_batch_state_store_mock.go @@ -0,0 +1,85 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: queue_batch_state_store.go +// +// Generated by this command: +// +// mockgen -source=queue_batch_state_store.go -destination=mock/queue_batch_state_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" +) + +// MockQueueBatchStateStore is a mock of QueueBatchStateStore interface. +type MockQueueBatchStateStore struct { + ctrl *gomock.Controller + recorder *MockQueueBatchStateStoreMockRecorder + isgomock struct{} +} + +// MockQueueBatchStateStoreMockRecorder is the mock recorder for MockQueueBatchStateStore. +type MockQueueBatchStateStoreMockRecorder struct { + mock *MockQueueBatchStateStore +} + +// NewMockQueueBatchStateStore creates a new mock instance. +func NewMockQueueBatchStateStore(ctrl *gomock.Controller) *MockQueueBatchStateStore { + mock := &MockQueueBatchStateStore{ctrl: ctrl} + mock.recorder = &MockQueueBatchStateStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockQueueBatchStateStore) EXPECT() *MockQueueBatchStateStoreMockRecorder { + return m.recorder +} + +// Delete mocks base method. +func (m *MockQueueBatchStateStore) Delete(ctx context.Context, queue string, state entity.BatchState, batchID string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Delete", ctx, queue, state, batchID) + ret0, _ := ret[0].(error) + return ret0 +} + +// Delete indicates an expected call of Delete. +func (mr *MockQueueBatchStateStoreMockRecorder) Delete(ctx, queue, state, batchID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockQueueBatchStateStore)(nil).Delete), ctx, queue, state, batchID) +} + +// List mocks base method. +func (m *MockQueueBatchStateStore) List(ctx context.Context, queue string, state entity.BatchState) ([]entity.QueueBatchState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "List", ctx, queue, state) + ret0, _ := ret[0].([]entity.QueueBatchState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// List indicates an expected call of List. +func (mr *MockQueueBatchStateStoreMockRecorder) List(ctx, queue, state any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "List", reflect.TypeOf((*MockQueueBatchStateStore)(nil).List), ctx, queue, state) +} + +// Put mocks base method. +func (m *MockQueueBatchStateStore) Put(ctx context.Context, record entity.QueueBatchState) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Put", ctx, record) + ret0, _ := ret[0].(error) + return ret0 +} + +// Put indicates an expected call of Put. +func (mr *MockQueueBatchStateStoreMockRecorder) Put(ctx, record any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockQueueBatchStateStore)(nil).Put), ctx, record) +} diff --git a/submitqueue/extension/storage/mock/storage_mock.go b/submitqueue/extension/storage/mock/storage_mock.go index d32eb70a..181c4cb6 100644 --- a/submitqueue/extension/storage/mock/storage_mock.go +++ b/submitqueue/extension/storage/mock/storage_mock.go @@ -110,6 +110,20 @@ func (mr *MockStorageMockRecorder) GetChangeStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChangeStore", reflect.TypeOf((*MockStorage)(nil).GetChangeStore)) } +// GetQueueBatchStateStore mocks base method. +func (m *MockStorage) GetQueueBatchStateStore() storage.QueueBatchStateStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetQueueBatchStateStore") + ret0, _ := ret[0].(storage.QueueBatchStateStore) + return ret0 +} + +// GetQueueBatchStateStore indicates an expected call of GetQueueBatchStateStore. +func (mr *MockStorageMockRecorder) GetQueueBatchStateStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetQueueBatchStateStore", reflect.TypeOf((*MockStorage)(nil).GetQueueBatchStateStore)) +} + // GetRequestBatchStore mocks base method. func (m *MockStorage) GetRequestBatchStore() storage.RequestBatchStore { m.ctrl.T.Helper() diff --git a/submitqueue/extension/storage/mysql/BUILD.bazel b/submitqueue/extension/storage/mysql/BUILD.bazel index af189d37..ff9475a0 100644 --- a/submitqueue/extension/storage/mysql/BUILD.bazel +++ b/submitqueue/extension/storage/mysql/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "queue_batch_state_store.go", "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", @@ -33,6 +34,7 @@ go_test( "batch_store_test.go", "build_store_test.go", "change_store_test.go", + "queue_batch_state_store_test.go", "request_batch_store_test.go", "request_log_store_test.go", "request_queue_summary_store_test.go", diff --git a/submitqueue/extension/storage/mysql/queue_batch_state_store.go b/submitqueue/extension/storage/mysql/queue_batch_state_store.go new file mode 100644 index 00000000..1e43c0d1 --- /dev/null +++ b/submitqueue/extension/storage/mysql/queue_batch_state_store.go @@ -0,0 +1,91 @@ +// Copyright (c) 2026 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" + "fmt" + + "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 queueBatchStateStore struct { + db *sql.DB + scope tally.Scope +} + +// NewQueueBatchStateStore creates a new MySQL-backed QueueBatchStateStore. +func NewQueueBatchStateStore(db *sql.DB, scope tally.Scope) storage.QueueBatchStateStore { + return &queueBatchStateStore{db: db, scope: scope} +} + +// List returns every record filed under (queue, state). The WHERE clause is a +// prefix of the (queue, state, batch_id) PK, so this is a PK-prefix scan. +func (s *queueBatchStateStore) List(ctx context.Context, queue string, state entity.BatchState) (ret []entity.QueueBatchState, retErr error) { + op := metrics.Begin(s.scope, "list", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + const query = "SELECT queue, state, batch_id FROM queue_batch_state WHERE queue = ? AND state = ?" + rows, err := s.db.QueryContext(ctx, query, queue, string(state)) + if err != nil { + return nil, fmt.Errorf("failed to query queue batch state records for queue=%s state=%s: %w", queue, state, err) + } + defer rows.Close() + + var results []entity.QueueBatchState + for rows.Next() { + var rec entity.QueueBatchState + if err := rows.Scan(&rec.Queue, &rec.State, &rec.BatchID); err != nil { + return nil, fmt.Errorf("failed to scan queue batch state record for queue=%s state=%s: %w", queue, state, err) + } + results = append(results, rec) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("failed to iterate queue batch state records for queue=%s state=%s: %w", queue, state, err) + } + return results, nil +} + +// Put inserts a record. A primary-key conflict on (queue, state, batch_id) is +// silently ignored via INSERT IGNORE — the record carries no data beyond its +// identity, so re-putting it is a no-op success. +func (s *queueBatchStateStore) Put(ctx context.Context, record entity.QueueBatchState) (retErr error) { + op := metrics.Begin(s.scope, "put", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + const query = "INSERT IGNORE INTO queue_batch_state (queue, state, batch_id) VALUES (?, ?, ?)" + if _, err := s.db.ExecContext(ctx, query, record.Queue, string(record.State), record.BatchID); err != nil { + return fmt.Errorf("failed to put queue batch state record queue=%s state=%s batch_id=%s: %w", record.Queue, record.State, record.BatchID, err) + } + return nil +} + +// Delete removes the record identified by (queue, state, batchID). Deleting an +// absent record is a no-op success — rows-affected is intentionally not checked. +func (s *queueBatchStateStore) Delete(ctx context.Context, queue string, state entity.BatchState, batchID string) (retErr error) { + op := metrics.Begin(s.scope, "delete", metrics.StorageLatencyBuckets) + defer func() { op.Complete(retErr) }() + + const query = "DELETE FROM queue_batch_state WHERE queue = ? AND state = ? AND batch_id = ?" + if _, err := s.db.ExecContext(ctx, query, queue, string(state), batchID); err != nil { + return fmt.Errorf("failed to delete queue batch state record queue=%s state=%s batch_id=%s: %w", queue, state, batchID, err) + } + return nil +} diff --git a/submitqueue/extension/storage/mysql/queue_batch_state_store_test.go b/submitqueue/extension/storage/mysql/queue_batch_state_store_test.go new file mode 100644 index 00000000..652c15f3 --- /dev/null +++ b/submitqueue/extension/storage/mysql/queue_batch_state_store_test.go @@ -0,0 +1,215 @@ +// Copyright (c) 2026 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" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +func setupQueueBatchStateStoreTest(t *testing.T) (*sql.DB, sqlmock.Sqlmock, storage.QueueBatchStateStore) { + t.Helper() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + return db, mock, NewQueueBatchStateStore(db, testMetrics()) +} + +func TestQueueBatchStateStore_List(t *testing.T) { + record1 := entity.QueueBatchState{ + Queue: "monorepo", + State: entity.BatchStateSpeculating, + BatchID: "monorepo/batch/1", + } + record2 := entity.QueueBatchState{ + Queue: record1.Queue, + State: record1.State, + BatchID: "monorepo/batch/2", + } + storeErr := errors.New("storage failed") + tests := map[string]struct { + setup func(sqlmock.Sqlmock) + want []entity.QueueBatchState + errMsg string + }{ + "query fails": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT queue, state, batch_id FROM queue_batch_state"). + WithArgs(record1.Queue, string(record1.State)). + WillReturnError(fmt.Errorf("connection reset")) + }, + errMsg: "connection reset", + }, + "empty bucket": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT queue, state, batch_id FROM queue_batch_state"). + WithArgs(record1.Queue, string(record1.State)). + WillReturnRows(sqlmock.NewRows([]string{"queue", "state", "batch_id"})) + }, + }, + "row iteration fails": { + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"queue", "state", "batch_id"}). + AddRow(record1.Queue, string(record1.State), record1.BatchID). + RowError(0, storeErr) + mock.ExpectQuery("SELECT queue, state, batch_id FROM queue_batch_state"). + WithArgs(record1.Queue, string(record1.State)). + WillReturnRows(rows) + }, + errMsg: storeErr.Error(), + }, + "success": { + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"queue", "state", "batch_id"}). + AddRow(record1.Queue, string(record1.State), record1.BatchID). + AddRow(record2.Queue, string(record2.State), record2.BatchID) + mock.ExpectQuery("SELECT queue, state, batch_id FROM queue_batch_state"). + WithArgs(record1.Queue, string(record1.State)). + WillReturnRows(rows) + }, + want: []entity.QueueBatchState{record1, record2}, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + db, mock, store := setupQueueBatchStateStoreTest(t) + defer db.Close() + tt.setup(mock) + + got, err := store.List(context.Background(), record1.Queue, record1.State) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.want, got) + } + assert.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestQueueBatchStateStore_Put(t *testing.T) { + record := entity.QueueBatchState{ + Queue: "monorepo", + State: entity.BatchStateCreated, + BatchID: "monorepo/batch/1", + } + tests := map[string]struct { + setup func(sqlmock.Sqlmock) + errMsg string + }{ + "insert fails": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT IGNORE INTO queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnError(fmt.Errorf("connection reset")) + }, + errMsg: "connection reset", + }, + "existing record is a no-op success": { + setup: func(mock sqlmock.Sqlmock) { + // INSERT IGNORE reports zero affected rows on a PK conflict. + mock.ExpectExec("INSERT IGNORE INTO queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnResult(sqlmock.NewResult(0, 0)) + }, + }, + "success": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT IGNORE INTO queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + db, mock, store := setupQueueBatchStateStoreTest(t) + defer db.Close() + tt.setup(mock) + + err := store.Put(context.Background(), record) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + } else { + assert.NoError(t, err) + } + assert.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestQueueBatchStateStore_Delete(t *testing.T) { + record := entity.QueueBatchState{ + Queue: "monorepo", + State: entity.BatchStateSucceeded, + BatchID: "monorepo/batch/1", + } + tests := map[string]struct { + setup func(sqlmock.Sqlmock) + errMsg string + }{ + "delete fails": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnError(fmt.Errorf("connection reset")) + }, + errMsg: "connection reset", + }, + "absent record is a no-op success": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnResult(sqlmock.NewResult(0, 0)) + }, + }, + "success": { + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("DELETE FROM queue_batch_state"). + WithArgs(record.Queue, string(record.State), record.BatchID). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + db, mock, store := setupQueueBatchStateStoreTest(t) + defer db.Close() + tt.setup(mock) + + err := store.Delete(context.Background(), record.Queue, record.State, record.BatchID) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + } else { + assert.NoError(t, err) + } + assert.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/submitqueue/extension/storage/mysql/schema/queue_batch_state.sql b/submitqueue/extension/storage/mysql/schema/queue_batch_state.sql new file mode 100644 index 00000000..c9c3ba92 --- /dev/null +++ b/submitqueue/extension/storage/mysql/schema/queue_batch_state.sql @@ -0,0 +1,11 @@ +-- Membership records filing each in-queue batch under a lifecycle state bucket. +-- The (queue, state, batch_id) PK *is* the lookup: "batches of queue Q in state S" +-- is a PK-prefix scan, so no secondary index is needed. Rows are advisory — the +-- authoritative state is batch.state — and a batch's rows are removed when it +-- exits the queue, so the table only ever holds in-queue batches. +CREATE TABLE IF NOT EXISTS queue_batch_state ( + queue VARCHAR(255) NOT NULL, + state VARCHAR(255) NOT NULL, + batch_id VARCHAR(255) NOT NULL, + PRIMARY KEY (queue, state, batch_id) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/submitqueue/extension/storage/mysql/storage.go b/submitqueue/extension/storage/mysql/storage.go index 053c4686..e6f174bd 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 + queueBatchStateStore storage.QueueBatchStateStore + buildStore storage.BuildStore + 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")), + queueBatchStateStore: NewQueueBatchStateStore(db, scope.SubScope("queue_batch_state_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")), }, nil } @@ -83,6 +85,11 @@ func (f *mysqlStorage) GetBatchDependentStore() storage.BatchDependentStore { return f.batchDependentStore } +// GetQueueBatchStateStore returns the MySQL-backed QueueBatchStateStore. +func (f *mysqlStorage) GetQueueBatchStateStore() storage.QueueBatchStateStore { + return f.queueBatchStateStore +} + // GetBuildStore returns the MySQL-backed BuildStore. func (f *mysqlStorage) GetBuildStore() storage.BuildStore { return f.buildStore diff --git a/submitqueue/extension/storage/queue_batch_state_store.go b/submitqueue/extension/storage/queue_batch_state_store.go new file mode 100644 index 00000000..e59573ce --- /dev/null +++ b/submitqueue/extension/storage/queue_batch_state_store.go @@ -0,0 +1,52 @@ +// Copyright (c) 2026 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=queue_batch_state_store.go -destination=mock/queue_batch_state_store_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// QueueBatchStateStore manages the per-queue membership records that file each in-queue +// batch under a lifecycle state bucket, so "the batches of queue Q filed under state S" +// is a single read keyed by (queue, state) — a primary-key prefix on any backend, with +// no secondary index or server-side filtering required. +// +// Records are advisory: the authoritative state is on the Batch entity, and a record may +// transiently file a batch under a bucket it has already left. Readers therefore treat a +// listing as a set of candidate batch IDs — they load each Batch by key and classify it +// by its own State, never by the bucket the record was found in. +// +// The interface is intentionally per-state and per-record so that any backend (SQL, +// DynamoDB, Bigtable, …) can implement it without multi-key queries or batch atomicity. +// Callers loop over the states they care about; a batch moves buckets via Put of the new +// record followed by Delete of the old one, which keeps at least one record visible +// throughout. All writes are idempotent so queue redeliveries can safely repeat them. +type QueueBatchStateStore interface { + // List returns every record filed under (queue, state). An empty slice means the + // bucket is empty. Order is unspecified. + List(ctx context.Context, queue string, state entity.BatchState) ([]entity.QueueBatchState, error) + + // Put persists a record. Writing an already-existing (queue, state, batchID) record + // is a no-op success, so the call is idempotent under redeliveries. + Put(ctx context.Context, record entity.QueueBatchState) error + + // Delete removes the record identified by (queue, state, batchID). Deleting an + // absent record is a no-op success, so the call is idempotent under redeliveries. + Delete(ctx context.Context, queue string, state entity.BatchState, batchID string) error +} diff --git a/submitqueue/extension/storage/storage.go b/submitqueue/extension/storage/storage.go index c2233298..7b03806e 100644 --- a/submitqueue/extension/storage/storage.go +++ b/submitqueue/extension/storage/storage.go @@ -61,6 +61,9 @@ type Storage interface { // GetBatchDependentStore returns the BatchDependentStore instance. GetBatchDependentStore() BatchDependentStore + // GetQueueBatchStateStore returns the QueueBatchStateStore instance. + GetQueueBatchStateStore() QueueBatchStateStore + // GetBuildStore returns the BuildStore instance. GetBuildStore() BuildStore diff --git a/test/integration/submitqueue/extension/storage/suite.go b/test/integration/submitqueue/extension/storage/suite.go index 2bcd3dce..fa34de05 100644 --- a/test/integration/submitqueue/extension/storage/suite.go +++ b/test/integration/submitqueue/extension/storage/suite.go @@ -339,6 +339,63 @@ func (s *StorageContractSuite) TestStorage_BatchUpdateReplacesAllNonKeyFields() assert.Equal(t, got, unchanged) } +// TestStorage_QueueBatchStateRecordLifecycle exercises the QueueBatchStateStore +// contract: Put idempotency, List isolation between (queue, state) buckets, the +// Put-then-Delete record move, and Delete idempotency. +func (s *StorageContractSuite) TestStorage_QueueBatchStateRecordLifecycle() { + t := s.T() + ctx := s.ctx + store := s.storage.GetQueueBatchStateStore() + + created1 := entity.QueueBatchState{Queue: "qbs-queue-a", State: entity.BatchStateCreated, BatchID: "qbs-queue-a/batch/1"} + created2 := entity.QueueBatchState{Queue: "qbs-queue-a", State: entity.BatchStateCreated, BatchID: "qbs-queue-a/batch/2"} + speculating := entity.QueueBatchState{Queue: "qbs-queue-a", State: entity.BatchStateSpeculating, BatchID: "qbs-queue-a/batch/3"} + otherQueue := entity.QueueBatchState{Queue: "qbs-queue-b", State: entity.BatchStateCreated, BatchID: "qbs-queue-b/batch/1"} + + require.NoError(t, store.Put(ctx, created1)) + require.NoError(t, store.Put(ctx, created2)) + require.NoError(t, store.Put(ctx, speculating)) + require.NoError(t, store.Put(ctx, otherQueue)) + + // Re-putting an existing record is a no-op success. + require.NoError(t, store.Put(ctx, created1)) + + // List returns exactly one (queue, state) bucket: no other states, no other queues, no duplicates. + got, err := store.List(ctx, "qbs-queue-a", entity.BatchStateCreated) + require.NoError(t, err) + assert.ElementsMatch(t, []entity.QueueBatchState{created1, created2}, got) + + got, err = store.List(ctx, "qbs-queue-a", entity.BatchStateSpeculating) + require.NoError(t, err) + assert.ElementsMatch(t, []entity.QueueBatchState{speculating}, got) + + // An empty bucket lists empty, not an error. + got, err = store.List(ctx, "qbs-queue-a", entity.BatchStateMerging) + require.NoError(t, err) + assert.Empty(t, got) + + // A record move: file under the new state, then remove the old bucket's record. + moved := entity.QueueBatchState{Queue: created1.Queue, State: entity.BatchStateSpeculating, BatchID: created1.BatchID} + require.NoError(t, store.Put(ctx, moved)) + require.NoError(t, store.Delete(ctx, created1.Queue, created1.State, created1.BatchID)) + + got, err = store.List(ctx, "qbs-queue-a", entity.BatchStateCreated) + require.NoError(t, err) + assert.ElementsMatch(t, []entity.QueueBatchState{created2}, got) + + got, err = store.List(ctx, "qbs-queue-a", entity.BatchStateSpeculating) + require.NoError(t, err) + assert.ElementsMatch(t, []entity.QueueBatchState{speculating, moved}, got) + + // Deleting an absent record is a no-op success. + require.NoError(t, store.Delete(ctx, created1.Queue, created1.State, created1.BatchID)) + + // The other queue is untouched by all of the above. + got, err = store.List(ctx, "qbs-queue-b", entity.BatchStateCreated) + require.NoError(t, err) + assert.ElementsMatch(t, []entity.QueueBatchState{otherQueue}, got) +} + // TestStorage_NotFound tests getting a non-existent request func (s *StorageContractSuite) TestStorage_NotFound() { t := s.T()