diff --git a/submitqueue/extension/storage/README.md b/submitqueue/extension/storage/README.md index e9a6c366..eeab72c5 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). `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 diff --git a/submitqueue/extension/storage/batch_store.go b/submitqueue/extension/storage/batch_store.go index 5d12a7ac..defe08aa 100644 --- a/submitqueue/extension/storage/batch_store.go +++ b/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) } diff --git a/submitqueue/extension/storage/mock/batch_store_mock.go b/submitqueue/extension/storage/mock/batch_store_mock.go index f4a589de..3fd698e9 100644 --- a/submitqueue/extension/storage/mock/batch_store_mock.go +++ b/submitqueue/extension/storage/mock/batch_store_mock.go @@ -70,21 +70,6 @@ func (mr *MockBatchStoreMockRecorder) Get(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockBatchStore)(nil).Get), ctx, id) } -// GetByQueueAndStates mocks base method. -func (m *MockBatchStore) GetByQueueAndStates(ctx context.Context, queue string, states []entity.BatchState) ([]entity.Batch, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetByQueueAndStates", ctx, queue, states) - ret0, _ := ret[0].([]entity.Batch) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetByQueueAndStates indicates an expected call of GetByQueueAndStates. -func (mr *MockBatchStoreMockRecorder) GetByQueueAndStates(ctx, queue, states any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetByQueueAndStates", reflect.TypeOf((*MockBatchStore)(nil).GetByQueueAndStates), ctx, queue, states) -} - // Update mocks base method. func (m *MockBatchStore) Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) error { m.ctrl.T.Helper() diff --git a/submitqueue/extension/storage/mysql/batch_store.go b/submitqueue/extension/storage/mysql/batch_store.go index 3aa70584..dbfa51c9 100644 --- a/submitqueue/extension/storage/mysql/batch_store.go +++ b/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 -} diff --git a/submitqueue/extension/storage/mysql/batch_store_test.go b/submitqueue/extension/storage/mysql/batch_store_test.go index 28b0b4f6..34c833dd 100644 --- a/submitqueue/extension/storage/mysql/batch_store_test.go +++ b/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()) - }) -} diff --git a/submitqueue/extension/storage/mysql/schema/README.md b/submitqueue/extension/storage/mysql/schema/README.md index b1445624..2e7ee204 100644 --- a/submitqueue/extension/storage/mysql/schema/README.md +++ b/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 diff --git a/submitqueue/extension/storage/mysql/schema/batch.sql b/submitqueue/extension/storage/mysql/schema/batch.sql index 398d17ca..0b12e792 100644 --- a/submitqueue/extension/storage/mysql/schema/batch.sql +++ b/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; diff --git a/submitqueue/orchestrator/controller/batch/BUILD.bazel b/submitqueue/orchestrator/controller/batch/BUILD.bazel index 9394c33d..048125fa 100644 --- a/submitqueue/orchestrator/controller/batch/BUILD.bazel +++ b/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", diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 44472a50..74a1d6ec 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/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 } diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index a1c1afed..524d2d6f 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -71,6 +71,28 @@ func newSequentialCounter(ctrl *gomock.Controller) *countermock.MockCounter { return cnt } +// newQueueBatchStateStore returns a QueueBatchStateStore mock that accepts any +// record write and lists the given batches as membership records under their +// current state. Callers hydrating candidates must set up the corresponding +// BatchStore.Get expectations themselves. +func newQueueBatchStateStore(ctrl *gomock.Controller, active ...entity.Batch) *storagemock.MockQueueBatchStateStore { + s := storagemock.NewMockQueueBatchStateStore(ctrl) + s.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().Delete(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().List(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, queue string, state entity.BatchState) ([]entity.QueueBatchState, error) { + var records []entity.QueueBatchState + for _, b := range active { + if b.Queue == queue && b.State == state { + records = append(records, entity.QueueBatchState{Queue: queue, State: state, BatchID: b.ID}) + } + } + return records, nil + }, + ).AnyTimes() + return s +} + // testRequest returns a standard test request for batch tests. func testRequest() entity.Request { return entity.Request{ @@ -95,7 +117,6 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M if mockStorage == nil { mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mockBatchStore.EXPECT().Update(gomock.Any(), gomock.Any(), int32(1), int32(2)).Return(nil).AnyTimes() @@ -111,6 +132,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M mockRequestBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mockStorage = storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() @@ -182,7 +204,6 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { request := testRequest() mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{ ID: "test-queue/batch/1", @@ -208,6 +229,7 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { }).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() @@ -265,6 +287,7 @@ func TestController_Process_StorageFailure(t *testing.T) { mockReqStore.EXPECT().Get(gomock.Any(), "test-queue/123").Return(entity.Request{}, fmt.Errorf("db connection lost")) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -285,7 +308,6 @@ func TestController_Process_RequestBatchStoreFailure(t *testing.T) { storeErr := errors.New("storage failed") batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, gomock.Any()).Return(nil, nil) batchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) @@ -300,6 +322,7 @@ func TestController_Process_RequestBatchStoreFailure(t *testing.T) { }).Return(storeErr) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -365,7 +388,8 @@ func TestController_Process_WithDependencies(t *testing.T) { } mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil) + mockBatchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(activeBatches[0], nil) + mockBatchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(activeBatches[1], nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{ ID: "test-queue/batch/1", @@ -413,6 +437,7 @@ func TestController_Process_WithDependencies(t *testing.T) { }).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl, activeBatches...)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() @@ -441,7 +466,8 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { } mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil) + mockBatchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(activeBatches[0], nil) + mockBatchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(activeBatches[1], nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{ ID: "test-queue/batch/1", @@ -477,6 +503,7 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { }).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl, activeBatches...)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() @@ -520,9 +547,11 @@ func TestController_Process_BatchDependentUpdateFailureDoesNotMutateFetchedDepen } mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, gomock.Any()).Return([]entity.Batch{activeBatch}, nil) + mockBatchStore.EXPECT().Get(gomock.Any(), activeBatch.ID).Return(activeBatch, nil) + mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchDependentStore.EXPECT().Get(gomock.Any(), activeBatch.ID).Return(existing, nil) mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ BatchID: activeBatch.ID, @@ -532,10 +561,16 @@ func TestController_Process_BatchDependentUpdateFailureDoesNotMutateFetchedDepen mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + mockReqStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateBatched), request.Version, request.Version+1).Return(nil) + + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + mockRequestBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl, activeBatch)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -557,12 +592,12 @@ func TestController_Process_AnalyzerFailure(t *testing.T) { request := testRequest() mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() @@ -619,6 +654,7 @@ func TestController_Process_HaltedShortCircuit(t *testing.T) { // No UpdateState expected — gomock fails if called. mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() @@ -653,7 +689,6 @@ func TestController_Process_CASLostToCancel(t *testing.T) { request := testRequest() mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) // Create must NOT be called — gomock fails if it is. mockReqStore := storagemock.NewMockRequestStore(ctrl) @@ -663,6 +698,7 @@ func TestController_Process_CASLostToCancel(t *testing.T) { ).Return(fmt.Errorf("cas: %w", storage.ErrVersionMismatch)) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() @@ -702,7 +738,6 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { request := testRequest() mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) // Create must NOT be called — gomock fails if it is. casErr := fmt.Errorf("db connection lost") @@ -713,6 +748,7 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { ).Return(casErr) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() @@ -746,7 +782,6 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { request.Version = 2 // prior attempt bumped from 1 → 2 mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{ ID: "test-queue/batch/1", @@ -774,6 +809,7 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { }).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() @@ -809,7 +845,6 @@ func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) { requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil) requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) @@ -833,6 +868,7 @@ func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) { ) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() @@ -881,7 +917,6 @@ func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) { var createdIDs []string batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), firstRequest.Queue, entity.DependencyBatchStates()).Return(nil, nil).Times(2) batchStore.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( func(_ context.Context, batch entity.Batch) error { createdIDs = append(createdIDs, batch.ID) @@ -922,6 +957,7 @@ func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) { }).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() @@ -947,7 +983,6 @@ func TestController_Process_InitializationFailure(t *testing.T) { requestStore.EXPECT().Update(gomock.Any(), requestWithState(request, entity.RequestStateBatched), int32(1), int32(2)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil) batchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) @@ -961,6 +996,7 @@ func TestController_Process_InitializationFailure(t *testing.T) { }).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() @@ -1049,6 +1085,7 @@ func TestController_PopulateBatch_Errors(t *testing.T) { tt.mockFunc(batchStore, batchDependentStore) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() diff --git a/submitqueue/orchestrator/controller/cancel/BUILD.bazel b/submitqueue/orchestrator/controller/cancel/BUILD.bazel index c3e5f24d..2f604f98 100644 --- a/submitqueue/orchestrator/controller/cancel/BUILD.bazel +++ b/submitqueue/orchestrator/controller/cancel/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "//platform/base/messagequeue:go_default_library", "//platform/consumer: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", diff --git a/submitqueue/orchestrator/controller/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 6af35525..69366193 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -60,6 +60,7 @@ import ( entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "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" @@ -311,9 +312,9 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error ) if batch.State != entity.BatchStateCancelling { - newVersion := batch.Version + 1 - batch.State = entity.BatchStateCancelling - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + var err error + batch, err = corebatch.Transition(ctx, c.store, batch, entity.BatchStateCancelling) + if err != nil { metrics.NamedCounter(c.metricsScope, opName, "batch_update_errors", 1) // storage.ErrVersionMismatch here means the batch advanced concurrently // (e.g. speculate / merge progressed). Returned as-is because the @@ -322,10 +323,15 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error // again. return fmt.Errorf("failed to mark batch %s as cancelling: %w", batch.ID, err) } - batch.Version = newVersion metrics.NamedCounter(c.metricsScope, opName, "batch_cancelling", 1) } else { metrics.NamedCounter(c.metricsScope, opName, "batch_already_cancelling", 1) + // A prior pass wrote the intent but may have crashed before completing + // the membership record move; repair before re-publishing. + if err := corebatch.EnsureRecord(ctx, c.store, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_update_errors", 1) + return err + } } if err := c.publishBatchID(ctx, topickey.TopicKeySpeculate, batch.ID, batch.Queue); err != nil { diff --git a/submitqueue/orchestrator/controller/cancel/cancel_test.go b/submitqueue/orchestrator/controller/cancel/cancel_test.go index 887b9acd..7d9ac361 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel_test.go +++ b/submitqueue/orchestrator/controller/cancel/cancel_test.go @@ -34,6 +34,15 @@ import ( "go.uber.org/zap/zaptest" ) +// newQueueBatchStateStore returns a QueueBatchStateStore mock that accepts any +// membership-record write; cancel never lists record buckets. +func newQueueBatchStateStore(ctrl *gomock.Controller) *storagemock.MockQueueBatchStateStore { + s := storagemock.NewMockQueueBatchStateStore(ctrl) + s.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().Delete(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return s +} + func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch { batch.State = state return batch @@ -110,6 +119,7 @@ func TestNewController(t *testing.T) { pub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() controller := newController(t, store, registry) require.NotNil(t, controller) @@ -132,6 +142,7 @@ func TestProcess_AlreadyTerminal_NoOp(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() controller := newController(t, store, registry) @@ -147,6 +158,7 @@ func TestProcess_RequestNotFound_Retryable(t *testing.T) { reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{}, storage.ErrNotFound) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() controller := newController(t, store, registry) @@ -181,6 +193,7 @@ func TestProcess_CancelsUnbatchedRequest(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, "q/1") @@ -213,6 +226,7 @@ func TestProcess_AlreadyCancelling_SkipsMarkCancelling(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, "q/1") @@ -241,6 +255,7 @@ func TestProcess_MarkCancellingVersionMismatch_Retryable(t *testing.T) { Return(storage.ErrVersionMismatch) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() controller := newController(t, store, registry) @@ -270,6 +285,7 @@ func TestProcess_UnbatchedVersionMismatch_Retryable(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, "q/1") @@ -295,11 +311,12 @@ func TestProcess_UnbatchedRequestDiverged_Acks(t *testing.T) { ) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return(nil, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, "q/1") controller := newController(t, store, registry) err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) @@ -320,11 +337,12 @@ func TestProcess_UnbatchedRequestDisappears_Retryable(t *testing.T) { ) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return(nil, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, "q/1") controller := newController(t, store, registry) err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) @@ -372,6 +390,7 @@ func TestProcess_BatchPath_HandsOffToSpeculate(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateCancelling), int32(3), int32(4)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, req.ID, batch) @@ -419,6 +438,7 @@ func TestProcess_CancelsEveryApplicableBatch(t *testing.T) { ).Times(2) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, request.ID, terminalBatch, batch2, batch1) @@ -457,6 +477,7 @@ func TestProcess_BatchFailureDoesNotPreventLaterCancellation(t *testing.T) { ) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, request.ID, batch2, batch1) @@ -492,6 +513,7 @@ func TestProcess_NonCancellableBatchSuppressesRequestCancellation(t *testing.T) batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, request.ID, batch) @@ -521,6 +543,7 @@ func TestProcess_BatchedWithoutMatchCancelsRequest(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, request.ID) @@ -555,6 +578,7 @@ func TestProcess_CreatingBatchDoesNotSuppressRequestCancellation(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, request.ID, batch) @@ -599,6 +623,7 @@ func TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate(t *testing.T) { // No batch Update — already in Cancelling. store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, req.ID, batch) @@ -630,6 +655,7 @@ func TestProcess_BatchIntentVersionMismatch_Retryable(t *testing.T) { Return(storage.ErrVersionMismatch) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() expectBatchLookup(ctrl, store, batchStore, req.ID, batch) @@ -645,6 +671,7 @@ func TestProcess_DeserializeError(t *testing.T) { registry, _ := newRegistry(t, ctrl) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() controller := newController(t, store, registry) err := controller.Process(context.Background(), newDelivery(t, ctrl, []byte("not json"), "q/1")) require.Error(t, err) @@ -658,6 +685,7 @@ func TestProcess_RequestStoreError(t *testing.T) { reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{}, fmt.Errorf("db down")) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() controller := newController(t, store, registry) @@ -729,6 +757,7 @@ func TestFindBatches(t *testing.T) { tt.mockFunc(requestBatchStore, batchStore) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() diff --git a/submitqueue/orchestrator/controller/dlq/BUILD.bazel b/submitqueue/orchestrator/controller/dlq/BUILD.bazel index 0dda0190..85768f22 100644 --- a/submitqueue/orchestrator/controller/dlq/BUILD.bazel +++ b/submitqueue/orchestrator/controller/dlq/BUILD.bazel @@ -17,6 +17,7 @@ go_library( "//api/runway/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", + "//submitqueue/core/batch:go_default_library", "//submitqueue/core/request:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/storage:go_default_library", diff --git a/submitqueue/orchestrator/controller/dlq/batch_test.go b/submitqueue/orchestrator/controller/dlq/batch_test.go index 5863b271..92c5c5d7 100644 --- a/submitqueue/orchestrator/controller/dlq/batch_test.go +++ b/submitqueue/orchestrator/controller/dlq/batch_test.go @@ -31,6 +31,7 @@ import ( func TestDLQBatchController_InterfaceAndAccessors(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") @@ -62,6 +63,7 @@ func TestDLQBatchController_Process_FailsAndFansOut(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -78,6 +80,7 @@ func TestDLQBatchController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") delivery := newMockDelivery(ctrl, []byte("garbage")) @@ -89,6 +92,7 @@ func TestDLQBatchController_Process_EmptyIDFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBatchController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyMerge), "orchestrator-merge-dlq") payload, err := entity.BatchID{ID: ""}.ToBytes() diff --git a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go index 10a3eae4..fa8c12e2 100644 --- a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go @@ -32,6 +32,7 @@ import ( func TestDLQBuildSignalController_InterfaceAndAccessors(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") @@ -68,6 +69,7 @@ func TestDLQBuildSignalController_Process_FansOutToBatch(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -88,6 +90,7 @@ func TestDLQBuildSignalController_Process_BuildNotFoundIsNoOp(t *testing.T) { buildStore.EXPECT().Get(gomock.Any(), "build-1").Return(entity.Build{}, storage.ErrNotFound) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") @@ -108,6 +111,7 @@ func TestDLQBuildSignalController_Process_BuildMissingBatchIsNoOp(t *testing.T) }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") @@ -123,6 +127,7 @@ func TestDLQBuildSignalController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQBuildSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(topickey.TopicKeyBuildSignal), "orchestrator-buildsignal-dlq") delivery := newMockDelivery(ctrl, []byte("garbage")) diff --git a/submitqueue/orchestrator/controller/dlq/dlq.go b/submitqueue/orchestrator/controller/dlq/dlq.go index d8114994..689360de 100644 --- a/submitqueue/orchestrator/controller/dlq/dlq.go +++ b/submitqueue/orchestrator/controller/dlq/dlq.go @@ -39,6 +39,7 @@ import ( "fmt" "github.com/uber/submitqueue/platform/consumer" + corebatch "github.com/uber/submitqueue/submitqueue/core/batch" requestcore "github.com/uber/submitqueue/submitqueue/core/request" "github.com/uber/submitqueue/submitqueue/entity" "github.com/uber/submitqueue/submitqueue/extension/storage" @@ -129,6 +130,11 @@ func failBatch(ctx context.Context, store storage.Storage, registry consumer.Top logger.Infow("dlq reconcile: batch already failed, repairing request fan-out", "batch_id", batchID, ) + // A prior attempt may have CAS'd to Failed without completing the + // membership record move; repair it alongside the fan-out. + if err := corebatch.EnsureRecord(ctx, store, batch); err != nil { + return err + } case entity.BatchStateSucceeded, entity.BatchStateCancelled: logger.Infow("dlq reconcile: batch has a different terminal outcome, skipping", "batch_id", batchID, @@ -136,13 +142,12 @@ func failBatch(ctx context.Context, store storage.Storage, registry consumer.Top ) return nil default: - newVersion := batch.Version + 1 previousState := batch.State - batch.State = entity.BatchStateFailed - if err := store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { - return fmt.Errorf("failed to update batch %s state to failed: %w", batchID, err) + updated, err := corebatch.Transition(ctx, store, batch, entity.BatchStateFailed) + if err != nil { + return err } - batch.Version = newVersion + batch = updated logger.Infow("dlq reconcile: batch marked failed", "batch_id", batchID, "previous_state", string(previousState), diff --git a/submitqueue/orchestrator/controller/dlq/dlq_test.go b/submitqueue/orchestrator/controller/dlq/dlq_test.go index c3e73998..11d0888a 100644 --- a/submitqueue/orchestrator/controller/dlq/dlq_test.go +++ b/submitqueue/orchestrator/controller/dlq/dlq_test.go @@ -31,6 +31,15 @@ import ( "go.uber.org/zap/zaptest" ) +// newQueueBatchStateStore returns a QueueBatchStateStore mock that accepts any +// membership-record write; these tests never list record buckets. +func newQueueBatchStateStore(ctrl *gomock.Controller) *storagemock.MockQueueBatchStateStore { + s := storagemock.NewMockQueueBatchStateStore(ctrl) + s.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().Delete(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return s +} + func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch { batch.State = state return batch @@ -63,6 +72,7 @@ func TestFailRequest_TerminalStates(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() registry := consumer.TopicRegistry{} if tt.wantLog { @@ -103,6 +113,7 @@ func TestFailRequest_CancellingTransitionsToError(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() err := failRequest(context.Background(), store, registry, zaptest.NewLogger(t).Sugar(), "q/1", "") @@ -127,6 +138,7 @@ func TestFailRequest_TransitionsToError(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() err := failRequest(context.Background(), store, registry, zaptest.NewLogger(t).Sugar(), "q/1", "") @@ -151,6 +163,7 @@ func TestFailRequest_LogPublishErrorPropagates(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() err := failRequest(context.Background(), store, registry, zaptest.NewLogger(t).Sugar(), "q/1", "") @@ -164,6 +177,7 @@ func TestFailRequest_NotFoundIsNoOp(t *testing.T) { requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{}, storage.ErrNotFound) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() err := failRequest(context.Background(), store, consumer.TopicRegistry{}, zaptest.NewLogger(t).Sugar(), "q/1", "") @@ -177,6 +191,7 @@ func TestFailRequest_GenericGetErrorIsNonRetryable(t *testing.T) { requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{}, fmt.Errorf("boom")) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() err := failRequest(context.Background(), store, consumer.TopicRegistry{}, zaptest.NewLogger(t).Sugar(), "q/1", "") @@ -214,6 +229,7 @@ func TestFailBatch_TransitionsAndFansOut(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -243,6 +259,7 @@ func TestFailBatch_FailedFansOutForRepair(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -260,6 +277,7 @@ func TestFailBatch_DifferentTerminalOutcomeSkipsFanOut(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() err := failBatch(context.Background(), store, consumer.TopicRegistry{}, zaptest.NewLogger(t).Sugar(), "q/batch/1", "") @@ -297,6 +315,7 @@ func TestFailBatch_CancellingTransitionsToFailed(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -311,6 +330,7 @@ func TestFailBatch_NotFoundIsNoOp(t *testing.T) { batchStore.EXPECT().Get(gomock.Any(), "q/batch/1").Return(entity.Batch{}, storage.ErrNotFound) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() err := failBatch(context.Background(), store, consumer.TopicRegistry{}, zaptest.NewLogger(t).Sugar(), "q/batch/1", "") diff --git a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go index d93f8f4c..d0aeb10e 100644 --- a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go @@ -32,6 +32,7 @@ import ( func TestDLQMergeConflictSignalController_InterfaceAndAccessors(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQMergeConflictSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeConflictCheckSignal), "orchestrator-mergeconflictsignal-dlq") @@ -55,6 +56,7 @@ func TestDLQMergeConflictSignalController_Process_ReconcilesRequest(t *testing.T }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() c := NewDLQMergeConflictSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, registry, TopicKey(runwaymq.TopicKeyMergeConflictCheckSignal), "orchestrator-mergeconflictsignal-dlq") @@ -70,6 +72,7 @@ func TestDLQMergeConflictSignalController_Process_MalformedPayloadFails(t *testi ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQMergeConflictSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeConflictCheckSignal), "orchestrator-mergeconflictsignal-dlq") delivery := newMockDelivery(ctrl, []byte("garbage")) diff --git a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go index 214b4bbf..a33b034d 100644 --- a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go @@ -32,6 +32,7 @@ import ( func TestDLQMergeSignalController_InterfaceAndAccessors(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQMergeSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeSignal), "orchestrator-mergesignal-dlq") @@ -65,6 +66,7 @@ func TestDLQMergeSignalController_Process_ReconcilesBatch(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() @@ -81,6 +83,7 @@ func TestDLQMergeSignalController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQMergeSignalController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, TopicKey(runwaymq.TopicKeyMergeSignal), "orchestrator-mergesignal-dlq") delivery := newMockDelivery(ctrl, []byte("garbage")) diff --git a/submitqueue/orchestrator/controller/dlq/request_test.go b/submitqueue/orchestrator/controller/dlq/request_test.go index 5da1c10b..2d7b264b 100644 --- a/submitqueue/orchestrator/controller/dlq/request_test.go +++ b/submitqueue/orchestrator/controller/dlq/request_test.go @@ -33,6 +33,7 @@ import ( func TestDLQRequestController_InterfaceAndAccessors(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") @@ -56,6 +57,7 @@ func TestDLQRequestController_Process_LandRequestPayload(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, registry, DecodeLandRequestID, TopicKey(topickey.TopicKeyStart), "orchestrator-start-dlq") @@ -82,6 +84,7 @@ func TestDLQRequestController_Process_CancelRequestPayload(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, registry, DecodeCancelRequestID, TopicKey(topickey.TopicKeyCancel), "orchestrator-cancel-dlq") @@ -109,6 +112,7 @@ func TestDLQRequestController_Process_RequestIDPayload(t *testing.T) { }) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, registry, DecodeRequestID, TopicKey(topickey.TopicKeyBatch), "orchestrator-batch-dlq") @@ -129,6 +133,7 @@ func TestDLQRequestController_Process_DifferentTerminalOutcomeSkips(t *testing.T }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") @@ -144,6 +149,7 @@ func TestDLQRequestController_Process_MalformedPayloadFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() // no store calls expected c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") @@ -157,6 +163,7 @@ func TestDLQRequestController_Process_EmptyIDFails(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() // no store calls expected c := NewDLQRequestController(zaptest.NewLogger(t).Sugar(), testScope(), store, consumer.TopicRegistry{}, DecodeRequestID, TopicKey(topickey.TopicKeyValidate), "orchestrator-validate-dlq") diff --git a/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel b/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel index 78018f23..b1fb73d9 100644 --- a/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel +++ b/submitqueue/orchestrator/controller/mergesignal/BUILD.bazel @@ -11,6 +11,7 @@ go_library( "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", + "//submitqueue/core/batch:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/storage:go_default_library", diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go index 7738b361..51c5c015 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal.go @@ -31,6 +31,7 @@ import ( entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" + corebatch "github.com/uber/submitqueue/submitqueue/core/batch" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" "github.com/uber/submitqueue/submitqueue/extension/storage" @@ -115,10 +116,15 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er } // Idempotency: a previous delivery already transitioned this batch to a - // terminal state. Re-fan-out in case that attempt missed the downstream - // publishes, then ack. + // terminal state. Repair the membership record (a prior attempt may have + // CAS'd without completing the record move), re-fan-out in case that + // attempt missed the downstream publishes, then ack. if batch.State.IsTerminal() { metrics.NamedCounter(c.metricsScope, opName, "skipped_terminal", 1) + if err := corebatch.EnsureRecord(ctx, c.store, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "state_update_errors", 1) + return err + } return c.fanout(ctx, batch.ID, batch.Queue) } @@ -138,13 +144,11 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er ) } - newVersion := batch.Version + 1 - batch.State = newState - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + batch, err = corebatch.Transition(ctx, c.store, batch, newState) + if err != nil { metrics.NamedCounter(c.metricsScope, opName, "state_update_errors", 1) - return fmt.Errorf("failed to transition batch %s to %s: %w", batch.ID, newState, err) + return err } - batch.Version = newVersion return c.fanout(ctx, batch.ID, batch.Queue) } diff --git a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go index def73e6e..4e989ec5 100644 --- a/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/mergesignal/mergesignal_test.go @@ -34,6 +34,15 @@ import ( "go.uber.org/zap/zaptest" ) +// newQueueBatchStateStore returns a QueueBatchStateStore mock that accepts any +// membership-record write; these tests never list record buckets. +func newQueueBatchStateStore(ctrl *gomock.Controller) *storagemock.MockQueueBatchStateStore { + s := storagemock.NewMockQueueBatchStateStore(ctrl) + s.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().Delete(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return s +} + func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch { batch.State = state return batch @@ -91,6 +100,7 @@ func newController(t *testing.T, store *storagemock.MockStorage, registry consum func TestNewController(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() var got []string c := newController(t, store, recordingRegistry(t, ctrl, &got)) @@ -116,6 +126,7 @@ func TestProcess_MergedAdvancesBatch(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateSucceeded), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() var got []string @@ -150,6 +161,7 @@ func TestProcess_NotMergedMarksBatchFailed(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateFailed), int32(3), int32(4)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() var got []string @@ -171,6 +183,7 @@ func TestProcess_CancellingShortCircuit(t *testing.T) { entity.Batch{ID: testBatchID, Queue: testQueue, State: entity.BatchStateCancelling, Version: 4}, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() // No Update and no fan-out: gomock fails if either runs. @@ -193,6 +206,7 @@ func TestProcess_TerminalReFansOut(t *testing.T) { entity.Batch{ID: testBatchID, Queue: testQueue, State: entity.BatchStateSucceeded, Version: 5}, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() var got []string @@ -208,6 +222,7 @@ func TestProcess_DeserializeErrorRejects(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() var got []string c := newController(t, store, recordingRegistry(t, ctrl, &got)) @@ -222,6 +237,7 @@ func TestProcess_StorageErrorRejects(t *testing.T) { batchStore.EXPECT().Get(gomock.Any(), testBatchID).Return(entity.Batch{}, assert.AnError) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() var got []string diff --git a/submitqueue/orchestrator/controller/speculate/BUILD.bazel b/submitqueue/orchestrator/controller/speculate/BUILD.bazel index 830c8b73..df3d7f5a 100644 --- a/submitqueue/orchestrator/controller/speculate/BUILD.bazel +++ b/submitqueue/orchestrator/controller/speculate/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", "//platform/metrics:go_default_library", + "//submitqueue/core/batch:go_default_library", "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/storage:go_default_library", diff --git a/submitqueue/orchestrator/controller/speculate/speculate.go b/submitqueue/orchestrator/controller/speculate/speculate.go index c9f85d41..43ce01b9 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate.go +++ b/submitqueue/orchestrator/controller/speculate/speculate.go @@ -23,6 +23,7 @@ import ( entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" "github.com/uber/submitqueue/platform/metrics" + corebatch "github.com/uber/submitqueue/submitqueue/core/batch" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" "github.com/uber/submitqueue/submitqueue/extension/storage" @@ -118,6 +119,12 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // leave them stuck waiting on a Cancelled dep. if batch.State.IsTerminal() { metrics.NamedCounter(c.metricsScope, opName, "self_heal_terminal", 1) + // Repair the membership record for the same crash window: a prior + // attempt may have CAS'd to terminal without completing the record move. + if err := corebatch.EnsureRecord(ctx, c.store, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return err + } if batch.State == entity.BatchStateCancelled { if err := c.respeculateDependents(ctx, batch); err != nil { return err @@ -129,6 +136,12 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Merging is owned by the merge controller, which has its own self-heal. if batch.State == entity.BatchStateMerging { metrics.NamedCounter(c.metricsScope, opName, "noop_merging", 1) + // A redelivery can land here after a tryFinalize attempt crashed + // between its CAS and the record move; repair before acking. + if err := corebatch.EnsureRecord(ctx, c.store, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return err + } return nil } @@ -158,11 +171,9 @@ func (c *Controller) startSpeculation(ctx context.Context, batch entity.Batch) e // Optimistic CAS: if the version has already advanced (concurrent speculate), // the next event will see the new state and behave correctly. - newVersion := batch.Version + 1 - batch.State = entity.BatchStateSpeculating - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + if _, err := corebatch.Transition(ctx, c.store, batch, entity.BatchStateSpeculating); err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update batch %s state to speculating: %w", batch.ID, err) + return err } metrics.NamedCounter(c.metricsScope, opName, "started_speculation", 1) @@ -220,11 +231,9 @@ func (c *Controller) tryFinalize(ctx context.Context, batch entity.Batch) error return fmt.Errorf("failed to publish to merge: %w", err) } - newVersion := batch.Version + 1 - batch.State = entity.BatchStateMerging - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + if _, err := corebatch.Transition(ctx, c.store, batch, entity.BatchStateMerging); err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update batch %s state to merging: %w", batch.ID, err) + return err } return nil @@ -243,13 +252,11 @@ func (c *Controller) failOnDependency(ctx context.Context, batch entity.Batch, d "dependency_state", string(dep.State), ) - newVersion := batch.Version + 1 - batch.State = entity.BatchStateFailed - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + batch, err := corebatch.Transition(ctx, c.store, batch, entity.BatchStateFailed) + if err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update batch %s state to failed: %w", batch.ID, err) + return err } - batch.Version = newVersion if err := c.publish(ctx, topickey.TopicKeyConclude, batch.ID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) @@ -307,13 +314,11 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error return err } - newVersion := batch.Version + 1 - batch.State = entity.BatchStateCancelled - if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil { + batch, err := corebatch.Transition(ctx, c.store, batch, entity.BatchStateCancelled) + if err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update batch %s state to cancelled: %w", batch.ID, err) + return err } - batch.Version = newVersion if err := c.respeculateDependents(ctx, batch); err != nil { return err diff --git a/submitqueue/orchestrator/controller/speculate/speculate_test.go b/submitqueue/orchestrator/controller/speculate/speculate_test.go index 8a0f7b2e..c5761a71 100644 --- a/submitqueue/orchestrator/controller/speculate/speculate_test.go +++ b/submitqueue/orchestrator/controller/speculate/speculate_test.go @@ -35,6 +35,15 @@ import ( "go.uber.org/zap/zaptest" ) +// newQueueBatchStateStore returns a QueueBatchStateStore mock that accepts any +// membership-record write; these tests never list record buckets. +func newQueueBatchStateStore(ctrl *gomock.Controller) *storagemock.MockQueueBatchStateStore { + s := storagemock.NewMockQueueBatchStateStore(ctrl) + s.EXPECT().Put(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + s.EXPECT().Delete(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return s +} + func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch { batch.State = state return batch @@ -100,6 +109,7 @@ func runProcess(t *testing.T, ctrl *gomock.Controller, controller *Controller, b func TestNewController(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() controller := newTestController(t, ctrl, store, nil) require.NotNil(t, controller) @@ -128,6 +138,7 @@ func TestController_Process_StartSpeculation(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateSpeculating), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -146,6 +157,7 @@ func TestController_Process_FinalizeNoDeps(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateMerging), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -166,6 +178,7 @@ func TestController_Process_FinalizeAllDepsSucceeded(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateMerging), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -184,6 +197,7 @@ func TestController_Process_WaitingOnDep(t *testing.T) { // No Update expected — gomock will fail if it is called. store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -204,6 +218,7 @@ func TestController_Process_FailedDepFailsBatch(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateFailed), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -226,6 +241,7 @@ func TestController_Process_CancelledDepSkipped(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateMerging), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -242,6 +258,7 @@ func TestController_Process_MergingNoOp(t *testing.T) { // No Update expected. store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -266,6 +283,7 @@ func TestController_Process_TerminalSelfHeals(t *testing.T) { // No Update expected. store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() // Require exactly one publish to the conclude topic for self-healing. @@ -310,6 +328,7 @@ func TestController_Process_CancelledTerminalSelfHealsDependents(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() // BuildStore must NOT be touched on the terminal self-heal path. @@ -381,6 +400,7 @@ func TestController_Process_CancellingTerminalFlow(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -444,6 +464,7 @@ func TestController_Process_CancellingBuildAlreadyTerminal(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -473,6 +494,7 @@ func TestController_Process_CancellingNoBuildYet(t *testing.T) { }, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -500,6 +522,7 @@ func TestController_Process_CancellingNoDependents(t *testing.T) { depStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.BatchDependent{BatchID: batch.ID, Dependents: []string{}, Version: 1}, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() store.EXPECT().GetBatchDependentStore().Return(depStore).AnyTimes() @@ -541,6 +564,7 @@ func TestController_Process_CancellingTerminalCASVersionMismatch(t *testing.T) { buildStore.EXPECT().Get(gomock.Any(), batch.ID).Return(entity.Build{}, storage.ErrNotFound) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() store.EXPECT().GetBuildStore().Return(buildStore).AnyTimes() // BatchDependentStore must NOT be touched — terminal CAS failed before fan-out. @@ -576,6 +600,7 @@ func TestController_Process_UnrecognizedState(t *testing.T) { batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -591,6 +616,7 @@ func TestController_Process_StorageFailure(t *testing.T) { batchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.Batch{}, fmt.Errorf("db connection lost")) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, nil) @@ -609,6 +635,7 @@ func TestController_Process_PublishFailure(t *testing.T) { // No Update expected — publish fails before we get there. store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newTestController(t, ctrl, store, fmt.Errorf("publish failed")) @@ -619,6 +646,7 @@ func TestController_Process_PublishFailure(t *testing.T) { func TestController_Process_BadPayload(t *testing.T) { ctrl := gomock.NewController(t) store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetQueueBatchStateStore().Return(newQueueBatchStateStore(ctrl)).AnyTimes() controller := newTestController(t, ctrl, store, nil) msg := entityqueue.NewMessage("anything", []byte("not-json"), "test-queue", nil)