feat(storage): add QueueBatchStateStore and shared batch state helpers

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",

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

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",
+    ],
+)

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
+}

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)
+		})
+	}
+}

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
+}