From e055992dc2ec274c1f8d440f7b0d9936f0b13d17 Mon Sep 17 00:00:00 2001 From: abettigole Date: Tue, 28 Jul 2026 22:45:42 +0000 Subject: [PATCH] fix(cancel): resolve durable batch ownership Use request-to-batch assignments with a legacy all-state fallback and retry while a claimed request's batch is not yet visible. Jira Issues: CODEM-304 --- submitqueue/entity/batch.go | 15 + submitqueue/entity/batch_test.go | 13 + .../orchestrator/controller/cancel/cancel.go | 103 ++++--- .../controller/cancel/cancel_test.go | 286 +++++++++++++++++- 4 files changed, 375 insertions(+), 42 deletions(-) diff --git a/submitqueue/entity/batch.go b/submitqueue/entity/batch.go index 39e21e5d..b360e329 100644 --- a/submitqueue/entity/batch.go +++ b/submitqueue/entity/batch.go @@ -62,6 +62,21 @@ func (s BatchState) IsTerminal() bool { } } +var nonCancellableBatchStates = map[BatchState]bool{ + BatchStateUnknown: true, + BatchStateCreating: true, + BatchStateMerging: true, + BatchStateSucceeded: true, + BatchStateFailed: true, + BatchStateCancelled: true, +} + +// IsCancellable returns true if cancellation should transition or republish a batch in this state. +// New non-terminal states are cancellable by default unless explicitly excluded above. +func (s BatchState) IsCancellable() bool { + return !nonCancellableBatchStates[s] +} + // IsBatchStateHalted returns true if the batch is either terminal or in the process of being cancelled. // Forward-progress controllers (build, buildsignal, speculate, merge) use this to short-circuit // work for batches that the user has asked to cancel — even though Cancelling is non-terminal, no diff --git a/submitqueue/entity/batch_test.go b/submitqueue/entity/batch_test.go index ebdfcd04..0afc38f2 100644 --- a/submitqueue/entity/batch_test.go +++ b/submitqueue/entity/batch_test.go @@ -45,6 +45,19 @@ func TestBatchState_IsTerminal(t *testing.T) { } } +func TestIsCancellable(t *testing.T) { + assert.True(t, BatchStateCreated.IsCancellable()) + assert.True(t, BatchStateSpeculating.IsCancellable()) + assert.True(t, BatchStateCancelling.IsCancellable()) + assert.True(t, BatchState("future").IsCancellable()) + assert.False(t, BatchStateUnknown.IsCancellable()) + assert.False(t, BatchStateCreating.IsCancellable()) + assert.False(t, BatchStateMerging.IsCancellable()) + assert.False(t, BatchStateSucceeded.IsCancellable()) + assert.False(t, BatchStateFailed.IsCancellable()) + assert.False(t, BatchStateCancelled.IsCancellable()) +} + func TestActiveBatchStates_ExcludesCreating(t *testing.T) { assert.NotContains(t, ActiveBatchStates(), BatchStateCreating) } diff --git a/submitqueue/orchestrator/controller/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 5f78689f..764f4009 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -23,13 +23,11 @@ // RequestStatusCancelled log entry. This path is fully owned by the cancel // controller. // -// - The request is already part of an active batch — the controller performs -// a single intent CAS on the batch (advancing it to BatchStateCancelling) -// and hands off to the speculate controller by publishing the batch ID to -// TopicKeySpeculate. The speculate controller then owns: cancelling any -// in-flight Build entity for the batch, fanning out to dependents, the -// terminal CAS to BatchStateCancelled, and publishing to conclude. Cancel -// does no terminal write and no downstream fan-out on the batch path. +// - The request is associated with one or more batch attempts — the controller +// records cancellation intent on every cancellable attempt and hands each one +// to speculate. Creating attempts are ignored because their reverse indexes +// may be incomplete, while Merging and terminal attempts retain their existing +// outcome for conclude to reconcile. // // The split exists so that the terminal write and the work that must precede // it (cancelling builds, respeculating dependents) live in the same controller @@ -40,11 +38,9 @@ // its terminal state. // // The controller is idempotent: re-delivery of the same CancelRequest after -// the terminal request transition is a no-op; re-delivery after the -// Cancelling write skips the mark-cancelling step and proceeds straight to -// the batch lookup. On the batch path, re-delivery against an already -// Cancelling batch re-publishes to TopicKeySpeculate (a cheap no-op nudge -// the speculate controller absorbs). +// the terminal request transition is a no-op. Re-delivery after a Cancelling +// write skips that CAS, finds every batch ID containing the request again, and +// republishes matching attempts already in BatchStateCancelling. // // Concurrent producers surface as the intrinsically retryable // storage.ErrVersionMismatch; the controller returns the wrapped error as-is @@ -56,7 +52,9 @@ package cancel import ( "context" + "errors" "fmt" + "sort" "github.com/uber-go/tally" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" @@ -144,16 +142,46 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er return err } - // Look for an active batch that already contains this request. - batch, found, err := c.findActiveBatch(ctx, request) + // Find every batch associated with this request. Retries may create multiple batch IDs, and each persisted attempt must be handled. + batches, err := c.findBatches(ctx, request) if err != nil { return err } - if !found { + // Return the first failure so the existing linear error-classification chain remains intact. + // Matches are sorted by ID below, making the selected error deterministic while every failure is still logged and counted. + var firstErr error + foundApplicableBatch := false + for _, batch := range batches { + switch { + case batch.State.IsCancellable(): + foundApplicableBatch = true + if err := c.cancelBatch(ctx, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_cancel_errors", 1) + c.logger.Errorw("failed to cancel batch", + "batch_id", batch.ID, + "batch_state", string(batch.State), + "error", err, + ) + if firstErr == nil { + firstErr = err + } + } + case batch.State == entity.BatchStateMerging: + // Merge owns the outcome once it has started. Conclude will reconcile the request with that outcome. + foundApplicableBatch = true + metrics.NamedCounter(c.metricsScope, opName, "batch_merging", 1) + case batch.State.IsTerminal(): + // The terminal batch outcome wins; conclude may not have reconciled the request yet. + foundApplicableBatch = true + metrics.NamedCounter(c.metricsScope, opName, "batch_already_terminal", 1) + } + } + + if !foundApplicableBatch { return c.cancelRequest(ctx, request, cancelReq.Reason) } - return c.cancelBatch(ctx, batch) + return firstErr } // markCancelling transitions the request to RequestStateCancelling (intent) if it @@ -182,31 +210,36 @@ func (c *Controller) markCancelling(ctx context.Context, request entity.Request) return request, nil } -// findActiveBatch scans all active batches in the request's queue for one whose -// Contains list includes the request. Returns (batch, true, nil) on a hit, -// (zero, false, nil) when the request is not yet batched, and any storage -// error otherwise. -// -// BatchStateCancelling is included in the active-state list so an idempotent -// redelivery of the cancel message (the prior pass wrote the intent but the -// speculate hand-off publish failed) still resolves the batch and re-attempts -// the publish. -func (c *Controller) findActiveBatch(ctx context.Context, request entity.Request) (entity.Batch, bool, error) { - // TODO: Scans all the batches in flight - make it more efficient? - active, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, request.Queue, entity.ActiveBatchStates()) +// findBatches resolves every batch attempt associated with the request. +// Associations whose batch was never persisted are stale retry artifacts and are ignored. +func (c *Controller) findBatches(ctx context.Context, request entity.Request) ([]entity.Batch, error) { + associations, err := c.store.GetRequestBatchStore().GetByRequestID(ctx, request.ID) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) - return entity.Batch{}, false, fmt.Errorf("failed to get active batches for queue=%s: %w", request.Queue, err) + metrics.NamedCounter(c.metricsScope, opName, "request_batch_store_errors", 1) + return nil, fmt.Errorf("failed to get batch associations for request %s: %w", request.ID, err) } - for _, b := range active { - for _, rid := range b.Contains { - if rid == request.ID { - return b, true, nil + var batches []entity.Batch + for _, association := range associations { + batch, err := c.store.GetBatchStore().Get(ctx, association.BatchID) + if err != nil { + if errors.Is(err, storage.ErrNotFound) { + // The association may precede batch persistence or may outlive a failed attempt. + // If the batch is later persisted and published, speculate re-checks the contained request state before starting work. + metrics.NamedCounter(c.metricsScope, opName, "stale_batch_associations", 1) + continue } + metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) + return nil, fmt.Errorf("failed to get associated batch %s for request %s: %w", association.BatchID, request.ID, err) } + batches = append(batches, batch) } - return entity.Batch{}, false, nil + + // The batches are independent, but deterministic order stabilizes logs, tests, and first-error selection. + sort.Slice(batches, func(i, j int) bool { + return batches[i].ID < batches[j].ID + }) + return batches, nil } // cancelRequest performs the terminal CAS (Cancelling → Cancelled) for a request diff --git a/submitqueue/orchestrator/controller/cancel/cancel_test.go b/submitqueue/orchestrator/controller/cancel/cancel_test.go index 162b78d9..f803a985 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel_test.go +++ b/submitqueue/orchestrator/controller/cancel/cancel_test.go @@ -71,6 +71,28 @@ func newDelivery(t *testing.T, ctrl *gomock.Controller, payload []byte, partitio return d } +func expectBatchLookup( + ctrl *gomock.Controller, + store *storagemock.MockStorage, + batchStore *storagemock.MockBatchStore, + requestID string, + batches ...entity.Batch, +) { + associations := make([]entity.RequestBatch, 0, len(batches)) + for _, batch := range batches { + associations = append(associations, entity.RequestBatch{ + RequestID: requestID, + BatchID: batch.ID, + Version: 1, + }) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + } + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().GetByRequestID(gomock.Any(), requestID).Return(associations, nil) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() +} + func TestNewController(t *testing.T) { ctrl := gomock.NewController(t) registry, pub := newRegistry(t, ctrl) @@ -144,11 +166,11 @@ func TestProcess_CancelsUnbatchedRequest(t *testing.T) { ) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return(nil, nil) store := storagemock.NewMockStorage(ctrl) 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", "user changed mind"), "q/1")) @@ -175,11 +197,11 @@ func TestProcess_AlreadyCancelling_SkipsMarkCancelling(t *testing.T) { reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateCancelled).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return(nil, nil) store := storagemock.NewMockStorage(ctrl) 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")) @@ -228,11 +250,11 @@ func TestProcess_UnbatchedVersionMismatch_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().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")) @@ -276,13 +298,13 @@ func TestProcess_BatchPath_HandsOffToSpeculate(t *testing.T) { reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return([]entity.Batch{batch}, nil) // Single batch CAS: intent only. No terminal CAS. batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(3), int32(4), entity.BatchStateCancelling).Return(nil) store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, req.ID, batch) // BatchDependentStore and BuildStore must NOT be touched — speculate owns those now. controller := newController(t, store, registry) @@ -292,10 +314,181 @@ func TestProcess_BatchPath_HandsOffToSpeculate(t *testing.T) { assert.Equal(t, []pubRec{{topic: "speculate", msgID: "q/batch/1"}}, records) } +func TestProcess_CancelsEveryApplicableBatch(t *testing.T) { + ctrl := gomock.NewController(t) + registry, publisher := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + batch1 := entity.Batch{ID: "q/batch/1", Queue: "q", Contains: []string{request.ID}, State: entity.BatchStateCreated, Version: 1} + batch2 := entity.Batch{ID: "q/batch/2", Queue: "q", Contains: []string{request.ID}, State: entity.BatchStateSpeculating, Version: 3} + terminalBatch := entity.Batch{ID: "q/batch/3", Queue: "q", Contains: []string{request.ID}, State: entity.BatchStateSucceeded, Version: 4} + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + + var operations []string + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().UpdateState(gomock.Any(), batch1.ID, int32(1), int32(2), entity.BatchStateCancelling).DoAndReturn( + func(context.Context, string, int32, int32, entity.BatchState) error { + operations = append(operations, "update:"+batch1.ID) + return nil + }, + ) + batchStore.EXPECT().UpdateState(gomock.Any(), batch2.ID, int32(3), int32(4), entity.BatchStateCancelling).DoAndReturn( + func(context.Context, string, int32, int32, entity.BatchState) error { + operations = append(operations, "update:"+batch2.ID) + return nil + }, + ) + publisher.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + operations = append(operations, "publish:"+msg.ID) + return nil + }, + ).Times(2) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, request.ID, terminalBatch, batch2, batch1) + + controller := newController(t, store, registry) + assert.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID))) + assert.Equal(t, []string{ + "update:q/batch/1", + "publish:q/batch/1", + "update:q/batch/2", + "publish:q/batch/2", + }, operations) +} + +func TestProcess_BatchFailureDoesNotPreventLaterCancellation(t *testing.T) { + ctrl := gomock.NewController(t) + registry, publisher := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + batch1 := entity.Batch{ID: "q/batch/1", Queue: "q", Contains: []string{request.ID}, State: entity.BatchStateCreated, Version: 1} + batch2 := entity.Batch{ID: "q/batch/2", Queue: "q", Contains: []string{request.ID}, State: entity.BatchStateCreated, Version: 2} + storeErr := fmt.Errorf("storage failed") + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().UpdateState(gomock.Any(), batch1.ID, int32(1), int32(2), entity.BatchStateCancelling).Return(storeErr) + batchStore.EXPECT().UpdateState(gomock.Any(), batch2.ID, int32(2), int32(3), entity.BatchStateCancelling).Return(nil) + publisher.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + assert.Equal(t, batch2.ID, msg.ID) + return nil + }, + ) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, request.ID, batch2, batch1) + + controller := newController(t, store, registry) + err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID)) + assert.ErrorContains(t, err, "failed to mark batch q/batch/1 as cancelling") +} + +func TestProcess_NonCancellableBatchSuppressesRequestCancellation(t *testing.T) { + tests := []struct { + name string + state entity.BatchState + }{ + {name: "merging", state: entity.BatchStateMerging}, + {name: "succeeded", state: entity.BatchStateSucceeded}, + {name: "failed", state: entity.BatchStateFailed}, + {name: "cancelled", state: entity.BatchStateCancelled}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + batch := entity.Batch{ID: "q/batch/1", Queue: "q", Contains: []string{request.ID}, State: tt.state, Version: 4} + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, request.ID, batch) + + controller := newController(t, store, registry) + assert.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID))) + }) + } +} + +func TestProcess_BatchedWithoutMatchCancelsRequest(t *testing.T) { + ctrl := gomock.NewController(t) + registry, publisher := newRegistry(t, ctrl) + publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + requestStore := storagemock.NewMockRequestStore(ctrl) + gomock.InOrder( + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil), + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil), + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(3), int32(4), entity.RequestStateCancelled).Return(nil), + ) + + batchStore := storagemock.NewMockBatchStore(ctrl) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, request.ID) + + controller := newController(t, store, registry) + assert.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, "stop"), request.ID))) +} + +func TestProcess_CreatingBatchDoesNotSuppressRequestCancellation(t *testing.T) { + ctrl := gomock.NewController(t) + registry, publisher := newRegistry(t, ctrl) + publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + batch := entity.Batch{ + ID: "q/batch/1", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateCreating, + Version: 1, + } + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(3), int32(4), entity.RequestStateCancelled).Return(nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, request.ID, batch) + + controller := newController(t, store, registry) + assert.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, "stop"), request.ID))) +} + // TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate covers the // idempotent redelivery path on the batch flow: a prior pass wrote the // intent CAS but the speculate publish failed. The cancel controller must -// observe the active (Cancelling) batch via findActiveBatch, skip the +// observe the Cancelling batch via findBatches, skip the // intent CAS, and re-publish the batch ID to TopicKeySpeculate so the // speculate controller can pick the work up. func TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate(t *testing.T) { @@ -325,12 +518,12 @@ func TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate(t *testing.T) { // No request UpdateState — already in Cancelling. batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return([]entity.Batch{batch}, nil) // No batch UpdateState — already in Cancelling. store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, req.ID, batch) controller := newController(t, store, registry) err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) @@ -355,13 +548,13 @@ func TestProcess_BatchIntentVersionMismatch_Retryable(t *testing.T) { reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) - batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return([]entity.Batch{batch}, nil) batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCancelling). Return(storage.ErrVersionMismatch) store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + expectBatchLookup(ctrl, store, batchStore, req.ID, batch) controller := newController(t, store, registry) err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) @@ -393,3 +586,82 @@ func TestProcess_RequestStoreError(t *testing.T) { err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) require.Error(t, err) } + +func TestFindBatches(t *testing.T) { + request := entity.Request{ID: "q/1", Queue: "q"} + batch1 := entity.Batch{ID: "q/batch/1", Contains: []string{request.ID}} + batch2 := entity.Batch{ID: "q/batch/2", Contains: []string{"q/other", request.ID}} + storeErr := fmt.Errorf("storage failed") + + tests := []struct { + name string + mockFunc func(*storagemock.MockRequestBatchStore, *storagemock.MockBatchStore) + want []entity.Batch + errMsg string + }{ + { + name: "association lookup failure", + mockFunc: func(store *storagemock.MockRequestBatchStore, _ *storagemock.MockBatchStore) { + store.EXPECT().GetByRequestID(gomock.Any(), request.ID).Return(nil, storeErr) + }, + errMsg: "failed to get batch associations", + }, + { + name: "stale association is ignored", + mockFunc: func(store *storagemock.MockRequestBatchStore, batchStore *storagemock.MockBatchStore) { + store.EXPECT().GetByRequestID(gomock.Any(), request.ID).Return([]entity.RequestBatch{{ + RequestID: request.ID, + BatchID: "q/batch/missing", + Version: 1, + }}, nil) + batchStore.EXPECT().Get(gomock.Any(), "q/batch/missing").Return(entity.Batch{}, storage.ErrNotFound) + }, + }, + { + name: "batch lookup failure", + mockFunc: func(store *storagemock.MockRequestBatchStore, batchStore *storagemock.MockBatchStore) { + store.EXPECT().GetByRequestID(gomock.Any(), request.ID).Return([]entity.RequestBatch{{ + RequestID: request.ID, + BatchID: batch1.ID, + Version: 1, + }}, nil) + batchStore.EXPECT().Get(gomock.Any(), batch1.ID).Return(entity.Batch{}, storeErr) + }, + errMsg: "failed to get associated batch", + }, + { + name: "batches are sorted", + mockFunc: func(store *storagemock.MockRequestBatchStore, batchStore *storagemock.MockBatchStore) { + store.EXPECT().GetByRequestID(gomock.Any(), request.ID).Return([]entity.RequestBatch{ + {RequestID: request.ID, BatchID: batch2.ID, Version: 1}, + {RequestID: request.ID, BatchID: batch1.ID, Version: 1}, + }, nil) + batchStore.EXPECT().Get(gomock.Any(), batch2.ID).Return(batch2, nil) + batchStore.EXPECT().Get(gomock.Any(), batch1.ID).Return(batch1, nil) + }, + want: []entity.Batch{batch1, batch2}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + batchStore := storagemock.NewMockBatchStore(ctrl) + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + tt.mockFunc(requestBatchStore, batchStore) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := &Controller{metricsScope: tally.NoopScope, store: store} + got, err := controller.findBatches(context.Background(), request) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + return + } + assert.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +}