feat(orchestrator): adopt QueueBatchStateStore for batch state reads and writes

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). `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.
+**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 replaced `BatchStore.GetByQueueAndStates`, which was the contract's one query-by-attribute. Unlike a `KEY idx_*`, the relationship is visible in the contract and portable to any backend.
 
 ### Decision path
 

submitqueue/extension/storage/batch_store.go

@@ -35,7 +35,4 @@ type BatchStore interface {
 	// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
 	// Version arithmetic is owned by the caller; the store performs a pure conditional write.
 	Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) error
-
-	// GetByQueueAndStates retrieves all batches that belong to the given queue and are in the given states.
-	GetByQueueAndStates(ctx context.Context, queue string, states []entity.BatchState) ([]entity.Batch, error)
 }

submitqueue/extension/storage/mysql/batch_store.go

@@ -20,7 +20,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"strings"
 
 	"github.com/go-sql-driver/mysql"
 	"github.com/uber-go/tally"
@@ -147,53 +146,3 @@ func (s *batchStore) Update(ctx context.Context, batch entity.Batch, oldVersion,
 
 	return nil
 }
-
-// GetByQueueAndStates retrieves all batches that belong to the given queue and are in the given states.
-func (s *batchStore) GetByQueueAndStates(ctx context.Context, queue string, states []entity.BatchState) (ret []entity.Batch, retErr error) {
-	op := metrics.Begin(s.scope, "get_by_queue_and_states", metrics.StorageLatencyBuckets)
-	defer func() { op.Complete(retErr) }()
-
-	if len(states) == 0 {
-		return nil, nil
-	}
-
-	query := "SELECT id, queue, contains, dependencies, state, version FROM batch WHERE queue = ? AND state IN (?" + strings.Repeat(", ?", len(states)-1) + ")"
-
-	args := make([]any, 1+len(states))
-	args[0] = queue
-	for i, state := range states {
-		args[i+1] = state
-	}
-
-	rows, err := s.db.QueryContext(ctx, query, args...)
-	if err != nil {
-		return nil, fmt.Errorf("failed to query batches by queue=%q states=%v from the database: %w", queue, states, err)
-	}
-	defer rows.Close()
-
-	var results []entity.Batch
-	for rows.Next() {
-		var batch entity.Batch
-		var containsJSON []byte
-		var dependenciesJSON []byte
-
-		if err := rows.Scan(&batch.ID, &batch.Queue, &containsJSON, &dependenciesJSON, &batch.State, &batch.Version); err != nil {
-			return nil, fmt.Errorf("failed to scan batch entity by queue=%q states=%v from the database: %w", queue, states, err)
-		}
-
-		if err := json.Unmarshal(containsJSON, &batch.Contains); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal contains for batch entity id=%s from the database: %w", batch.ID, err)
-		}
-
-		if err := json.Unmarshal(dependenciesJSON, &batch.Dependencies); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal dependencies for batch entity id=%s from the database: %w", batch.ID, err)
-		}
-
-		results = append(results, batch)
-	}
-	if err := rows.Err(); err != nil {
-		return nil, fmt.Errorf("failed to iterate batches by queue=%q states=%v from the database: %w", queue, states, err)
-	}
-
-	return results, nil
-}

submitqueue/extension/storage/mysql/batch_store_test.go

@@ -312,50 +312,3 @@ func TestBatchStore_Update(t *testing.T) {
 		})
 	}
 }
-
-func TestBatchStore_GetByQueueAndStates(t *testing.T) {
-	t.Run("empty states returns nil without querying", func(t *testing.T) {
-		db, mock, store := setupBatchStoreTest(t)
-		defer db.Close()
-
-		got, err := store.GetByQueueAndStates(context.Background(), "monorepo", nil)
-		require.NoError(t, err)
-		assert.Nil(t, got)
-		require.NoError(t, mock.ExpectationsWereMet())
-	})
-
-	t.Run("found", func(t *testing.T) {
-		db, mock, store := setupBatchStoreTest(t)
-		defer db.Close()
-
-		batch := entity.Batch{ID: "monorepo/batch/1", Queue: "monorepo", State: entity.BatchStateCreated, Version: 1}
-		containsJSON, err := json.Marshal(batch.Contains)
-		require.NoError(t, err)
-		dependenciesJSON, err := json.Marshal(batch.Dependencies)
-		require.NoError(t, err)
-
-		rows := sqlmock.NewRows([]string{"id", "queue", "contains", "dependencies", "state", "version"}).
-			AddRow(batch.ID, batch.Queue, containsJSON, dependenciesJSON, string(batch.State), batch.Version)
-		mock.ExpectQuery("SELECT id, queue, contains, dependencies, state, version FROM batch").
-			WithArgs("monorepo", entity.BatchStateCreated, entity.BatchStateMerging).
-			WillReturnRows(rows)
-
-		got, err := store.GetByQueueAndStates(context.Background(), "monorepo", []entity.BatchState{entity.BatchStateCreated, entity.BatchStateMerging})
-		require.NoError(t, err)
-		assert.Equal(t, []entity.Batch{batch}, got)
-		require.NoError(t, mock.ExpectationsWereMet())
-	})
-
-	t.Run("query error", func(t *testing.T) {
-		db, mock, store := setupBatchStoreTest(t)
-		defer db.Close()
-
-		mock.ExpectQuery("SELECT id, queue, contains, dependencies, state, version FROM batch").
-			WithArgs("monorepo", entity.BatchStateCreated).
-			WillReturnError(fmt.Errorf("connection reset"))
-
-		_, err := store.GetByQueueAndStates(context.Background(), "monorepo", []entity.BatchState{entity.BatchStateCreated})
-		require.Error(t, err)
-		require.NoError(t, mock.ExpectationsWereMet())
-	})
-}

submitqueue/extension/storage/mysql/schema/README.md

@@ -2,19 +2,17 @@
 
 ## batch table
 
-### Secondary index: `idx_queue_state (queue, state)`
+The `batch` table is keyed by `id` alone and carries no secondary index. Listing a queue's batches by state goes through the `queue_batch_state` table instead, so batch reads and writes stay pure primary-key operations.
 
-The `batch` table has a composite secondary index on `(queue, state)`. This index supports the `GetByQueueAndStates` query, which retrieves batches filtered by queue and one or more states. Without this index, the query would require a full table scan.
+## queue_batch_state table
 
-#### Trade-offs
+### Composite primary key: `(queue, state, batch_id)`
 
-- **Write overhead**: Every `INSERT` and `UPDATE` to the `batch` table must also update the secondary index, adding latency to write operations.
-- **Storage cost**: The index consumes additional disk space proportional to the number of rows in the table.
-- **Lock contention**: Under high write concurrency, index maintenance can increase lock contention on the affected index pages.
+`queue_batch_state` holds the queue's advisory per-state membership records (see `entity.QueueBatchState`): one row per batch per state bucket, no payload and no version column. The key leads with `queue` so a state-bucket listing is a primary-key-prefix scan and the table is shardable by queue. Rows are moved between buckets by the shared transition protocol in `submitqueue/core/batch`; writes are idempotent (`INSERT IGNORE`, keyed `DELETE`). The `batch` row remains authoritative — readers hydrate each candidate and classify by the batch's own state.
 
 #### Future: Prune job
 
-As the `batch` table grows, the secondary index will grow with it, increasing storage costs and degrading write performance. To mitigate this, a prune job should be introduced to periodically delete batches in terminal states (`succeeded`, `failed`, `cancelled`) that are older than a configurable retention period. This keeps the table and its indexes bounded in size, ensuring consistent query and write performance over time.
+Terminal-state records (and their batches) accumulate as the queue processes work. A prune job should periodically delete records and batches in terminal states (`succeeded`, `failed`, `cancelled`) older than a configurable retention period, keeping both tables bounded so query and write performance stay consistent over time.
 
 ## change table
 

submitqueue/extension/storage/mysql/schema/batch.sql

@@ -5,6 +5,5 @@ CREATE TABLE IF NOT EXISTS batch (
     dependencies JSON NOT NULL,
     state VARCHAR(255) NOT NUll,
     version INT NOT NULL,
-    PRIMARY KEY (id),
-    INDEX idx_queue_state (queue, state)
+    PRIMARY KEY (id)
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;

submitqueue/orchestrator/controller/batch/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//platform/consumer:go_default_library",
         "//platform/extension/counter:go_default_library",
         "//platform/metrics:go_default_library",
+        "//submitqueue/core/batch:go_default_library",
         "//submitqueue/core/request:go_default_library",
         "//submitqueue/core/topickey:go_default_library",
         "//submitqueue/entity:go_default_library",

submitqueue/orchestrator/controller/batch/batch.go

@@ -24,6 +24,7 @@ import (
 	"github.com/uber/submitqueue/platform/consumer"
 	"github.com/uber/submitqueue/platform/extension/counter"
 	"github.com/uber/submitqueue/platform/metrics"
+	corebatch "github.com/uber/submitqueue/submitqueue/core/batch"
 	corerequest "github.com/uber/submitqueue/submitqueue/core/request"
 	"github.com/uber/submitqueue/submitqueue/core/topickey"
 	"github.com/uber/submitqueue/submitqueue/entity"
@@ -135,8 +136,10 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er
 
 	// Get active batches for this queue and ask the conflict analyzer which
 	// of them the new batch must serialize behind. The dependency set drives
-	// the speculation graph downstream.
-	activeBatches, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, request.Queue, entity.DependencyBatchStates())
+	// the speculation graph downstream. The read goes through the queue's
+	// per-state membership records; classification uses each batch's own
+	// hydrated state, so a stale record can never misreport a batch.
+	activeBatches, err := corebatch.ListByStates(ctx, c.store, request.Queue, entity.DependencyBatchStates())
 	if err != nil {
 		metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1)
 		return fmt.Errorf("failed to get active batches for queue=%s: %w", request.Queue, err)
@@ -246,6 +249,13 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er
 		return fmt.Errorf("failed to create batch in batch store: %w", err)
 	}
 
+	// File the queue's membership record for the new batch so it is
+	// discoverable by state from its first moment in the queue.
+	if err := corebatch.EnsureRecord(ctx, c.store, batch); err != nil {
+		metrics.NamedCounter(c.metricsScope, opName, "queue_batch_state_errors", 1)
+		return err
+	}
+
 	for _, requestID := range batch.Contains {
 		association := entity.RequestBatch{
 			RequestID: requestID,
@@ -337,13 +347,11 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent
 
 	// The batch's own reverse-index row now exists and every dependency lists this batch as a dependent.
 	// Structural initialization is complete, so transition Creating → Created to make the batch ready for processing once published to speculate.
-	newVersion := batch.Version + 1
-	batch.State = entity.BatchStateCreated
-	if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil {
+	batch, err := corebatch.Transition(ctx, c.store, batch, entity.BatchStateCreated)
+	if err != nil {
 		metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1)
 		return entity.Batch{}, fmt.Errorf("failed to mark batch %s created: %w", batch.ID, err)
 	}
-	batch.Version = newVersion
 	return batch, nil
 }