diff --git a/submitqueue/core/request/BUILD.bazel b/submitqueue/core/request/BUILD.bazel index 63c149bc..a76ff26b 100644 --- a/submitqueue/core/request/BUILD.bazel +++ b/submitqueue/core/request/BUILD.bazel @@ -3,6 +3,7 @@ load("@rules_go//go:def.bzl", "go_library", "go_test") go_library( name = "go_default_library", srcs = [ + "batch.go", "log.go", "materializer.go", "request.go", @@ -21,6 +22,7 @@ go_library( go_test( name = "go_default_test", srcs = [ + "batch_test.go", "log_test.go", "materializer_test.go", "request_test.go", diff --git a/submitqueue/core/request/batch.go b/submitqueue/core/request/batch.go new file mode 100644 index 00000000..8fcbfb21 --- /dev/null +++ b/submitqueue/core/request/batch.go @@ -0,0 +1,46 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request + +import ( + "context" + "fmt" + "slices" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +// FindBatch returns the unique batch in the requested states that contains the request. +func FindBatch(ctx context.Context, store storage.BatchStore, request entity.Request, states []entity.BatchState) (entity.Batch, bool, error) { + batches, err := store.GetByQueueAndStates(ctx, request.Queue, states) + if err != nil { + return entity.Batch{}, false, fmt.Errorf("failed to get batches for queue=%s states=%v: %w", request.Queue, states, err) + } + + var match entity.Batch + found := false + for _, batch := range batches { + if !slices.Contains(batch.Contains, request.ID) { + continue + } + if found { + return entity.Batch{}, false, fmt.Errorf("request %s is contained by multiple batches in queue %s", request.ID, request.Queue) + } + match = batch + found = true + } + return match, found, nil +} diff --git a/submitqueue/core/request/batch_test.go b/submitqueue/core/request/batch_test.go new file mode 100644 index 00000000..f06327d5 --- /dev/null +++ b/submitqueue/core/request/batch_test.go @@ -0,0 +1,88 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package request + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/uber/submitqueue/submitqueue/entity" + storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" + "go.uber.org/mock/gomock" +) + +func TestFindBatch(t *testing.T) { + request := entity.Request{ID: "q/1", Queue: "q"} + states := entity.AllBatchStates() + batch := entity.Batch{ID: "q/batch/1", Queue: request.Queue, Contains: []string{request.ID}} + storeErr := errors.New("storage failed") + + tests := map[string]struct { + mockFunc func(*storagemock.MockBatchStore) + want entity.Batch + found bool + errMsg string + }{ + "store failure": { + mockFunc: func(store *storagemock.MockBatchStore) { + store.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, states).Return(nil, storeErr) + }, + errMsg: "storage failed", + }, + "no matching batch": { + mockFunc: func(store *storagemock.MockBatchStore) { + store.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, states).Return( + []entity.Batch{{ID: "q/batch/other", Contains: []string{"q/other"}}}, nil, + ) + }, + }, + "multiple matching batches": { + mockFunc: func(store *storagemock.MockBatchStore) { + store.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, states).Return( + []entity.Batch{batch, {ID: "q/batch/2", Queue: request.Queue, Contains: []string{request.ID}}}, nil, + ) + }, + errMsg: "contained by multiple batches", + }, + "success": { + mockFunc: func(store *storagemock.MockBatchStore) { + store.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, states).Return( + []entity.Batch{{ID: "q/batch/other", Contains: []string{"q/other"}}, batch}, nil, + ) + }, + want: batch, + found: true, + }, + } + + for name, tt := range tests { + t.Run(name, func(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + tt.mockFunc(store) + + got, found, err := FindBatch(context.Background(), store, request, states) + if tt.errMsg != "" { + assert.ErrorContains(t, err, tt.errMsg) + return + } + assert.NoError(t, err) + assert.Equal(t, tt.want, got) + assert.Equal(t, tt.found, found) + }) + } +} diff --git a/submitqueue/orchestrator/controller/cancel/BUILD.bazel b/submitqueue/orchestrator/controller/cancel/BUILD.bazel index 66f99b3f..fafc5bfd 100644 --- a/submitqueue/orchestrator/controller/cancel/BUILD.bazel +++ b/submitqueue/orchestrator/controller/cancel/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", + "//platform/errs:go_default_library", "//platform/metrics:go_default_library", "//submitqueue/core/request:go_default_library", "//submitqueue/core/topickey:go_default_library", @@ -25,6 +26,7 @@ go_test( deps = [ "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", + "//platform/errs:go_default_library", "//platform/extension/messagequeue/mock: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 5f78689f..789e33e9 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -31,6 +31,9 @@ // terminal CAS to BatchStateCancelled, and publishing to conclude. Cancel // does no terminal write and no downstream fan-out on the batch path. // +// A batch found in BatchStateCreating is retried without changing its state. +// Creating batches may not have their own reverse-index row yet, so allowing initialization to reach Created preserves the terminal fan-out invariant. +// // The split exists so that the terminal write and the work that must precede // it (cancelling builds, respeculating dependents) live in the same controller // — speculate is the single writer of every non-Cancelling batch state and is @@ -56,11 +59,13 @@ package cancel import ( "context" + "errors" "fmt" "github.com/uber-go/tally" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" + "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/metrics" corerequest "github.com/uber/submitqueue/submitqueue/core/request" "github.com/uber/submitqueue/submitqueue/core/topickey" @@ -135,27 +140,99 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er return nil } - // Step 1: record the cancellation intent on the request itself by transitioning - // to RequestStateCancelling. This is non-terminal; forward-progress controllers - // (validate, batch) treat it as halted, but conclude may still write a different - // terminal state if a concurrent merge or failure wins the race. + initialRequestState := request.State + var batch entity.Batch + var foundBatch bool + if initialRequestState == entity.RequestStateBatched { + // Batched requests resolve their assignment before recording cancellation intent. + batch, foundBatch, err = c.findAssignedBatch(ctx, request) + if err != nil { + return err + } + if !foundBatch { + // The batch controller claimed the request but has not yet persisted a batch. + // Retry instead of cancelling the request immediately before its batch becomes visible. + metrics.NamedCounter(c.metricsScope, opName, "batch_assignment_pending", 1) + return errs.NewRetryableError(fmt.Errorf("batch assignment for request %s is not visible yet", request.ID)) + } + if batch.State.IsTerminal() { + // Batch state becomes terminal before conclude reconciles the contained requests, so the request may still be Batched here. + // Let conclude apply the batch outcome instead of overwriting it with cancellation. + metrics.NamedCounter(c.metricsScope, opName, "batch_already_terminal", 1) + return nil + } + } + + // Record cancellation intent by transitioning the request to RequestStateCancelling. + // This state halts new forward progress, but a concurrent merge or failure may still win with a different terminal outcome. request, err = c.markCancelling(ctx, request) if err != nil { return err } - // Look for an active batch that already contains this request. - batch, found, err := c.findActiveBatch(ctx, request) - if err != nil { - return err + if !foundBatch { + switch initialRequestState { + case entity.RequestStateCancelling: + // A previous delivery may have recorded intent and failed before cancelling its assigned batch. + batch, foundBatch, err = c.findAssignedBatch(ctx, request) + default: + // The request was not Batched when cancellation began, but a batch worker may have persisted an active batch from a stale read. + batch, foundBatch, err = corerequest.FindBatch(ctx, c.store.GetBatchStore(), request, entity.ActiveBatchStates()) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_lookup_errors", 1) + } + } + if err != nil { + return err + } } - if !found { + if !foundBatch { return c.cancelRequest(ctx, request, cancelReq.Reason) } + if batch.State.IsTerminal() { + // A terminal batch can be visible before conclude reconciles the request from Cancelling to its terminal outcome. + metrics.NamedCounter(c.metricsScope, opName, "batch_already_terminal", 1) + return nil + } + if batch.State == entity.BatchStateCreating { + // A Creating batch may not have its BatchDependent row yet, so it cannot safely enter terminal cancellation. + // Retry until initialization reaches Created; RequestStateCancelling prevents new forward progress in the meantime. + metrics.NamedCounter(c.metricsScope, opName, "batch_creating", 1) + return errs.NewRetryableError(fmt.Errorf("batch %s is still creating", batch.ID)) + } + // The gateway recorded the cancellation reason with RequestStatusCancelling. + // The batch path only hands the batch ID to speculate; conclude emits the terminal request logs. return c.cancelBatch(ctx, batch) } +// findAssignedBatch resolves the immutable request-to-batch assignment. +// The all-state fallback supports batches persisted before their assignment and batches created before assignment persistence was introduced. +func (c *Controller) findAssignedBatch(ctx context.Context, request entity.Request) (entity.Batch, bool, error) { + // This is the normal lookup after the batch controller has recorded the durable request-to-batch assignment. + assignment, err := c.store.GetRequestBatchStore().Get(ctx, request.ID) + if errors.Is(err, storage.ErrNotFound) { + // The batch may have been persisted before its assignment, or by an older writer that never created assignments. + // Search every state so cancellation can recover active and terminal ownership. + batch, found, findErr := corerequest.FindBatch(ctx, c.store.GetBatchStore(), request, entity.AllBatchStates()) + if findErr != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_lookup_errors", 1) + } + return batch, found, findErr + } + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_batch_store_errors", 1) + return entity.Batch{}, false, fmt.Errorf("failed to get batch assignment for request %s: %w", request.ID, err) + } + + batch, err := c.store.GetBatchStore().Get(ctx, assignment.BatchID) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) + return entity.Batch{}, false, fmt.Errorf("failed to get assigned batch %s for request %s: %w", assignment.BatchID, request.ID, err) + } + return batch, true, nil +} + // markCancelling transitions the request to RequestStateCancelling (intent) if it // isn't already in that state. Returns the updated request (with the post-CAS // version and state) on success, or the original request on the idempotent path @@ -182,33 +259,6 @@ 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()) - 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) - } - - for _, b := range active { - for _, rid := range b.Contains { - if rid == request.ID { - return b, true, nil - } - } - } - return entity.Batch{}, false, nil -} - // cancelRequest performs the terminal CAS (Cancelling → Cancelled) for a request // that is not part of any active batch, and emits the RequestStatusCancelled log // entry. storage.ErrVersionMismatch here means a concurrent writer (typically diff --git a/submitqueue/orchestrator/controller/cancel/cancel_test.go b/submitqueue/orchestrator/controller/cancel/cancel_test.go index 162b78d9..05aeff32 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel_test.go +++ b/submitqueue/orchestrator/controller/cancel/cancel_test.go @@ -24,6 +24,7 @@ import ( "github.com/uber-go/tally" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" + "github.com/uber/submitqueue/platform/errs" queuemock "github.com/uber/submitqueue/platform/extension/messagequeue/mock" "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" @@ -177,8 +178,12 @@ func TestProcess_AlreadyCancelling_SkipsMarkCancelling(t *testing.T) { batchStore := storagemock.NewMockBatchStore(ctrl) batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return(nil, nil) + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.RequestBatch{}, storage.ErrNotFound) + store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newController(t, store, registry) @@ -292,12 +297,205 @@ func TestProcess_BatchPath_HandsOffToSpeculate(t *testing.T) { assert.Equal(t, []pubRec{{topic: "speculate", msgID: "q/batch/1"}}, records) } +// Cancellation records request intent but waits for a Creating batch to finish structural initialization. +func TestProcess_CreatingBatchRetries(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateStarted, Version: 2} + batch := entity.Batch{ + ID: "q/batch/1", + Queue: "q", + 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) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.ActiveBatchStates()).Return([]entity.Batch{batch}, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID)) + require.Error(t, err) + assert.True(t, errs.IsRetryable(err)) +} + +// A Batched request with no visible batch is in the claim-to-create window and must be retried rather than cancelled directly. +func TestProcess_BatchedWithoutAssignmentRetries(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{}, storage.ErrNotFound) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.AllBatchStates()).Return(nil, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID)) + require.Error(t, err) + assert.True(t, errs.IsRetryable(err)) +} + +// An active batch written before assignment persistence is still discovered and cancelled. +func TestProcess_BatchedWithoutAssignmentFindsLegacyActiveBatch(t *testing.T) { + ctrl := gomock.NewController(t) + registry, pub := newRegistry(t, ctrl) + pub.EXPECT().Publish(gomock.Any(), "speculate", 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.BatchStateSpeculating, + Version: 3, + } + + 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) + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{}, storage.ErrNotFound) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.AllBatchStates()).Return([]entity.Batch{batch}, nil) + batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(3), int32(4), entity.BatchStateCancelling).Return(nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + require.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID))) +} + +// Recovery fails closed when historical data contains more than one possible owning batch. +func TestProcess_BatchedWithoutAssignmentRejectsAmbiguousLegacyBatches(t *testing.T) { + ctrl := gomock.NewController(t) + registry, _ := newRegistry(t, ctrl) + + request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{}, storage.ErrNotFound) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.AllBatchStates()).Return([]entity.Batch{ + {ID: "q/batch/1", Contains: []string{request.ID}}, + {ID: "q/batch/2", Contains: []string{request.ID}}, + }, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + assert.Error(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID))) +} + +// The durable assignment path also waits while its batch is still Creating. +func TestProcess_BatchedAssignedCreatingRetries(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: 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) + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{ + RequestID: request.ID, + BatchID: batch.ID, + Version: 1, + }, nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), request.ID)) + require.Error(t, err) + assert.True(t, errs.IsRetryable(err)) +} + +// A terminal assigned batch wins over a late cancellation request. +func TestProcess_BatchedAssignedTerminalBatchNoOp(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: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateSucceeded, + Version: 4, + } + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{ + RequestID: request.ID, + BatchID: batch.ID, + Version: 1, + }, nil) + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + controller := newController(t, store, registry) + require.NoError(t, controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, request.ID, ""), 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 -// intent CAS, and re-publish the batch ID to TopicKeySpeculate so the -// speculate controller can pick the work up. +// observe the active Cancelling batch, skip the intent CAS, and re-publish the batch ID to TopicKeySpeculate. +// This gives the speculate controller another chance to pick the work up. func TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate(t *testing.T) { ctrl := gomock.NewController(t) registry, pub := newRegistry(t, ctrl) @@ -328,8 +526,12 @@ func TestProcess_BatchAlreadyCancelling_RepublishesToSpeculate(t *testing.T) { batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "q", gomock.Any()).Return([]entity.Batch{batch}, nil) // No batch UpdateState — already in Cancelling. + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), req.ID).Return(entity.RequestBatch{}, storage.ErrNotFound) + store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() controller := newController(t, store, registry)