From 3c288ddb46cccb5f2aed7aa65b2ac960bfcc83fe Mon Sep 17 00:00:00 2001 From: abettigole Date: Tue, 28 Jul 2026 22:49:48 +0000 Subject: [PATCH] feat(batch): make creation resumable Persist batches in Creating, initialize assignments and reverse indexes idempotently, and resume the same batch across queue redelivery. Jira Issues: CODEM-304 --- .../orchestrator/controller/batch/batch.go | 280 ++++++--- .../controller/batch/batch_test.go | 534 +++++++++++++++++- 2 files changed, 711 insertions(+), 103 deletions(-) diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 88c37f4c..2955eaed 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -18,6 +18,7 @@ import ( "context" "errors" "fmt" + "slices" "github.com/uber-go/tally" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" @@ -103,9 +104,23 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er "partition_key", msg.PartitionKey, ) - // Short-circuit if the request has been halted — either it already reached a - // terminal state, or the cancel controller has recorded a cancellation intent - // (RequestStateCancelling). A halted request must never spawn a new batch. + // Cancellation can mark the request Cancelling after its Creating batch is persisted but before initialization finishes. + // A redelivery must finish the batch's structural initialization before cancellation can safely make it terminal. + if request.State == entity.RequestStateCancelling { + batch, found, err := corerequest.FindBatch(ctx, c.store.GetBatchStore(), request, entity.ActiveBatchStates()) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_lookup_errors", 1) + return err + } + if found && batch.State == entity.BatchStateCreating { + metrics.NamedCounter(c.metricsScope, opName, "completed_creating_for_cancel", 1) + _, err := c.initializeBatch(ctx, batch) + return err + } + } + + // Short-circuit if the request has been halted. + // A halted request must never spawn a new batch or publish normal forward progress. if entity.IsRequestStateHalted(request.State) { metrics.NamedCounter(c.metricsScope, opName, "skipped_halted", 1) c.logger.Infow("skipping batch for halted request", @@ -115,6 +130,18 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er return nil } + // A prior delivery may have claimed the request and persisted some or all of its batch before failing. + // Resume that batch instead of minting a new ID so every durable step below is safe under queue redelivery. + if request.State == entity.RequestStateBatched { + batch, found, err := c.findAssignedBatch(ctx, request) + if err != nil { + return err + } + if found { + return c.resumeBatch(ctx, request, batch) + } + } + // TODO: if capacity is full, wait here for other requests to accumulate to batch them together, or include a request into an existing batch if it's not too late. // Generate a globally unique batch ID. @@ -128,22 +155,19 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er ID: fmt.Sprintf("%s/batch/%d", request.Queue, seq), Queue: request.Queue, Contains: []string{request.ID}, - State: entity.BatchStateCreated, + State: entity.BatchStateCreating, Version: 1, } - // 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. + // 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()) 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) } - // Dedupe by batch ID since a single (analyzed, in-flight) pair may be - // reported with multiple Conflict entries when different conflict types - // apply; the dependency graph only tracks the relation. + // Dedupe by batch ID because a single pair may have multiple conflict types while the dependency graph tracks only the relation. analyzer, err := c.analyzers.For(conflict.Config{QueueName: batch.Queue}) if err != nil { metrics.NamedCounter(c.metricsScope, opName, "conflict_analyzer_errors", 1) @@ -167,36 +191,6 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er batch.Dependencies = conflictingIDs - // Update reverse index for each conflicting batch (BatchDependent = - // "batches that depend on me"). One UpdateDependents call per conflict. - for _, depID := range conflictingIDs { - existing, err := c.store.GetBatchDependentStore().Get(ctx, depID) - if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to get batch dependent for batchID=%s: %w", depID, err) - } - - dependents := append(existing.Dependents, batch.ID) - - newVersion := existing.Version + 1 - if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, depID, existing.Version, newVersion, dependents); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", depID, batch.ID, err) - } - } - - // Create new reverse index entry for the new batch. It would be empty for now, but will be updated as new batches are created that conflict with this batch. - bd := entity.BatchDependent{ - BatchID: batch.ID, - Dependents: []string{}, - Version: 1, - } - - if err := c.store.GetBatchDependentStore().Create(ctx, bd); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to create batch dependent index for new batchID=%s: %w", batch.ID, err) - } - // Claim the request for this batch with a CAS-write that transitions the // request to RequestStateBatched. This CAS is the serialization point // between the batch controller and the cancel controller — without it, the @@ -209,47 +203,30 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // T1 batch.Get(R) → R{State: Validated, Version: 1} // T2 cancel.Get(R) → R{State: Validated, Version: 1} // T3 cancel.markCancelling CAS 1→2 → R{State: Cancelling, Version: 2} - // T4 cancel.findActiveBatch(R) → none (batch has not been Created yet) + // T4 cancel searches active batches → none (batch has not been Created yet) // T5 cancel.cancelRequest CAS 2→3 → R{State: Cancelled, Version: 3} // T6 batch.IsRequestStateHalted(R) → false (stale in-memory copy from T1) // T7 batch.BatchStore.Create(B{[R]}) → orphan batch containing a cancelled R // - // After T7 the orphan batch flows through speculate → merge → conclude; - // conclude does NOT gate on the source request state when writing the terminal - // state, so it would CAS the request from Cancelled back to Landed, silently - // undoing the user's cancel. + // After T7 the orphan batch can still flow through speculate → merge even though its only request was cancelled. + // Conclude preserves a different terminal request outcome, but the batch still performs invalid work and can participate in the dependency graph without owning a live request. // // The CAS below collapses that window. Whichever of batch.UpdateState(..., // RequestStateBatched) and cancel.markCancelling(... RequestStateCancelling) // reaches storage first wins; the loser sees storage.ErrVersionMismatch: // - If cancel won: this CAS fails. We ack the message (cancel will drive R - // to its terminal state on its own; no batch is needed). The reverse-index - // entry above becomes a dangling BatchDependent — tolerated per the - // "downstream should handle stale entries" contract on this store. + // to its terminal state on its own; no batch or reverse-index residue was + // written). // - If batch won: cancel.markCancelling will fail with ErrVersionMismatch // on its next attempt, re-fetch R, observe RequestStateBatched, and take // the batch-cancellation branch (which terminates the whole batch). // - // Note on re-delivery: a retry of a batch message that already CAS'd R to - // Batched but failed before/after BatchStore.Create lands in this code with - // R already in RequestStateBatched. The top-level IsRequestStateHalted check - // does NOT include Batched (Batched is forward-progress, not halted), so we - // reach here and re-CAS Batched → Batched (a version-only bump). The bump - // keeps the same serialization invariant on every attempt — if cancel sneaks - // in between our Get and this CAS, our version is stale and we abandon, just - // like the first-delivery case. The cost is an extra batch (the previous - // attempt may have already created one) which is tolerated per the comment - // on BatchStore.Create below. + // On redelivery, a persisted batch is resumed with the same ID. + // If the prior attempt failed after the request CAS but before BatchStore.Create, no batch exists to resume. + // The next attempt performs a version-only Batched → Batched CAS before creating a fresh batch so cancellation still races against a current request version. // - // Residual window: a thin race remains between this CAS and BatchStore.Create. - // During that window cancel.findActiveBatch can still observe R in Batched - // with no batch yet persisted, and take the request-only cancel path — which - // then leaves R in Cancelled and the batch we are about to create orphaned. - // Fully closing this requires cancel-side wait/retry when its pre-CAS - // observation was RequestStateBatched; deferred to a follow-up since the - // window is narrow (one storage round-trip) and the user-visible outcome - // (request cancelled) is still correct — the orphan batch just gets - // reconciled by conclude as if it had no requests to act on. + // During the CAS → Create window, the cancel controller observes Batched but cannot yet resolve an assignment. + // It retries rather than taking the request-only cancellation path, closing the remaining orphan-batch race. newRequestVersion := request.Version + 1 if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newRequestVersion, entity.RequestStateBatched); err != nil { // ErrVersionMismatch == cancel (or another writer) advanced R first. Ack @@ -270,9 +247,8 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er request.Version = newRequestVersion request.State = entity.RequestStateBatched - // Persist batch to storage. - // This is the final operation that concludes the batch creation process. If it fails, BatchDependents will be pointing to a batch id that does not exist. - // We do not reuse batch ids, a retry of this operation will create a new batch with a new ID. The downstream logic that operates on BatchDependent should be able to handle stale entries. + // Persist the batch before creating any references to it. + // Creating batches are visible to ownership and cancellation lookups but excluded from dependency analysis and normal processing. if err := c.store.GetBatchStore().Create(ctx, batch); err != nil { metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) return fmt.Errorf("failed to create batch in batch store: %w", err) @@ -285,12 +261,162 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er "dependency_count", len(batch.Dependencies), ) - // Record the "batched" status in the request log. This status corresponds to - // the RequestStateBatched transition CAS'd above, so it carries the request - // version for reconciliation. The message ID is scoped to (requestID, status), - // so a redelivery that creates a fresh batch re-emits "batched" with a - // different batch_id but is deduped to the first entry — acceptable, the - // request is batched either way. + return c.initializeAndPublish(ctx, request, batch) +} + +// findAssignedBatch resolves the durable request-to-batch assignment. +// The all-state fallback recovers batches persisted before their assignment and batches created by an older writer. +func (c *Controller) findAssignedBatch(ctx context.Context, request entity.Request) (entity.Batch, bool, error) { + // This is the normal resume path after initializeBatch has persisted the request-to-batch assignment. + assignment, err := c.store.GetRequestBatchStore().Get(ctx, request.ID) + if err == nil { + 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 + } + if !errors.Is(err, storage.ErrNotFound) { + 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) + } + + // A missing assignment means the previous attempt may have persisted the batch and failed before recording the mapping. + // Scan every state so terminal batches and batches written before assignment persistence was introduced can also be recovered. + batch, found, err := corerequest.FindBatch(ctx, c.store.GetBatchStore(), request, entity.AllBatchStates()) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_lookup_errors", 1) + } + if err != nil || !found { + return batch, found, err + } + if err := c.ensureRequestAssignments(ctx, batch); err != nil { + return entity.Batch{}, false, err + } + return batch, true, nil +} + +// resumeBatch continues the durable batch-creation handoff after redelivery. +func (c *Controller) resumeBatch(ctx context.Context, request entity.Request, batch entity.Batch) error { + switch batch.State { + case entity.BatchStateCreating: + metrics.NamedCounter(c.metricsScope, opName, "resumed_creating", 1) + return c.initializeAndPublish(ctx, request, batch) + case entity.BatchStateCreated: + metrics.NamedCounter(c.metricsScope, opName, "resumed_created", 1) + return c.publishBatch(ctx, request, batch) + default: + // Speculating or Merging means the handoff already succeeded. + // Cancelling means the cancellation pipeline owns the batch, and terminal states leave request reconciliation to conclude. + metrics.NamedCounter(c.metricsScope, opName, "resume_already_handed_off", 1) + return nil + } +} + +// initializeAndPublish creates the batch's reverse-index structure, marks it ready, and hands it to speculation. +// Every step is idempotent so redelivery can resume the same Creating batch. +func (c *Controller) initializeAndPublish(ctx context.Context, request entity.Request, batch entity.Batch) error { + batch, err := c.initializeBatch(ctx, batch) + if err != nil { + return err + } + return c.publishBatch(ctx, request, batch) +} + +// initializeBatch completes the durable structure of a Creating batch without publishing normal forward progress. +// Cancellation redelivery uses this path to make a partially initialized batch safe to terminate. +func (c *Controller) initializeBatch(ctx context.Context, batch entity.Batch) (entity.Batch, error) { + if err := c.ensureRequestAssignments(ctx, batch); err != nil { + return entity.Batch{}, err + } + + bd := entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + } + if err := c.store.GetBatchDependentStore().Create(ctx, bd); err != nil && !errors.Is(err, storage.ErrAlreadyExists) { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return entity.Batch{}, fmt.Errorf("failed to create batch dependent index for new batchID=%s: %w", batch.ID, err) + } + + for _, dependencyID := range batch.Dependencies { + if err := c.ensureDependentSubscription(ctx, dependencyID, batch.ID); err != nil { + return entity.Batch{}, err + } + } + + newVersion := batch.Version + 1 + if err := c.store.GetBatchStore().UpdateState(ctx, batch.ID, batch.Version, newVersion, entity.BatchStateCreated); 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 + batch.State = entity.BatchStateCreated + + return batch, nil +} + +// ensureRequestAssignments idempotently records the owning batch for every contained request. +// The assignment remains available after the batch becomes terminal so delayed redeliveries can find the original batch. +func (c *Controller) ensureRequestAssignments(ctx context.Context, batch entity.Batch) error { + store := c.store.GetRequestBatchStore() + for _, requestID := range batch.Contains { + assignment := entity.RequestBatch{ + RequestID: requestID, + BatchID: batch.ID, + Version: 1, + } + if err := store.Create(ctx, assignment); err != nil { + if !errors.Is(err, storage.ErrAlreadyExists) { + metrics.NamedCounter(c.metricsScope, opName, "request_batch_store_errors", 1) + return fmt.Errorf("failed to assign request %s to batch %s: %w", requestID, batch.ID, err) + } + + existingAssignment, getErr := store.Get(ctx, requestID) + if getErr != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_batch_store_errors", 1) + return fmt.Errorf("failed to verify batch assignment for request %s: %w", requestID, getErr) + } + if existingAssignment.BatchID != batch.ID { + // Another batch owns this request while the current batch also lists it. + // Fail closed and leave the current batch Creating so it is never published or used as a dependency. + metrics.NamedCounter(c.metricsScope, opName, "request_batch_conflicts", 1) + return fmt.Errorf("request %s is already assigned to batch %s instead of %s", requestID, existingAssignment.BatchID, batch.ID) + } + } + } + return nil +} + +// ensureDependentSubscription idempotently adds dependentID to the reverse index for dependencyID. +func (c *Controller) ensureDependentSubscription(ctx context.Context, dependencyID, dependentID string) error { + existing, err := c.store.GetBatchDependentStore().Get(ctx, dependencyID) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return fmt.Errorf("failed to get batch dependent for batchID=%s: %w", dependencyID, err) + } + + // The dependent subscription has already been created (e.g. on a prior attempt) - no-op. + if slices.Contains(existing.Dependents, dependentID) { + return nil + } + + dependents := append(append([]string{}, existing.Dependents...), dependentID) + newVersion := existing.Version + 1 + if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", dependencyID, dependentID, err) + } + return nil +} + +// publishBatch records the batched request status and performs the reliable handoff to speculation. +// Both publishes use deterministic IDs, so repeating them after redelivery is safe. +func (c *Controller) publishBatch(ctx context.Context, request entity.Request, batch entity.Batch) error { + // Record the "batched" status corresponding to the RequestStateBatched transition. + // The message ID is scoped to (requestID, status), so redelivery is deduplicated. logEntry := entity.NewRequestLog(request.ID, entity.RequestStatusBatched, request.Version, "", map[string]string{ "batch_id": batch.ID, }) @@ -300,8 +426,6 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er } // Publish to speculate topic for further processing. - // If it fails and the controller retries, a new batch will be created with the new batch ID but the same request ID. - // The downstream logic should be able to handle stale entries by looking at the state of the batch. if err := c.publish(ctx, topickey.TopicKeySpeculate, batch.ID, batch.Queue); err != nil { metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) return fmt.Errorf("failed to publish batch ID to speculate topic: %w", err) diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index fdfef8e0..f881851f 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -72,6 +72,13 @@ func testRequest() entity.Request { } } +func newRequestBatchStore(ctrl *gomock.Controller) *storagemock.MockRequestBatchStore { + store := storagemock.NewMockRequestBatchStore(ctrl) + store.EXPECT().Get(gomock.Any(), gomock.Any()).Return(entity.RequestBatch{}, storage.ErrNotFound).AnyTimes() + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + return store +} + // newTestController creates a controller with test dependencies. // If mockStorage is nil, a default MockStorage with an empty batch store is created. // If analyzer is nil, the "all" conflict analyzer is used (every active batch becomes a dependency). @@ -86,6 +93,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M 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().UpdateState(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), entity.BatchStateCreated).Return(nil).AnyTimes() mockReqStore := storagemock.NewMockRequestStore(ctrl) req := testRequest() @@ -98,6 +106,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M mockStorage = storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() } @@ -168,6 +177,7 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { 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().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) @@ -179,6 +189,7 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() // Capture messages published to the log topic. @@ -298,7 +309,11 @@ 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().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, batch entity.Batch) error { + assert.Equal(t, entity.BatchStateCreating, batch.State) + return nil + }) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) // batch/1 has no existing dependents. @@ -324,6 +339,7 @@ func TestController_Process_WithDependencies(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -351,6 +367,7 @@ 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().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) // Only batch/2 is selected by the analyzer, so only it gets a reverse-index update. @@ -368,6 +385,7 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() // Analyzer returns duplicate Conflict entries for the same batch (different @@ -425,20 +443,10 @@ func TestController_InterfaceImplementation(t *testing.T) { var _ consumer.Controller = controller } -// A request that is halted (terminal OR Cancelling) must be short-circuited -// before the batch controller queries the batch store, allocates a batch ID, -// CAS-claims the request, or publishes. We verify by configuring a batch -// store and counter with NO EXPECTs (gomock fails on any call), a request -// store that only expects the initial Get (no UpdateState), and a publisher -// that returns a sentinel error if invoked. -// -// Cancelling is non-terminal but must halt forward progress: the cancel -// controller has already recorded the cancellation intent on the request and -// owns the terminal write. Any new batch spawned here would be an orphan -// containing a request that is about to become Cancelled. +// A terminal request must be short-circuited before the batch controller queries the batch store, allocates an ID, claims the request, or publishes. +// The mocks have no expectations for those operations, so any unexpected forward progress fails the test. func TestController_Process_HaltedShortCircuit(t *testing.T) { for _, state := range []entity.RequestState{ - entity.RequestStateCancelling, entity.RequestStateCancelled, entity.RequestStateLanded, entity.RequestStateError, @@ -475,6 +483,66 @@ func TestController_Process_HaltedShortCircuit(t *testing.T) { } } +func TestController_Process_CancellingRequestCompletesCreatingBatch(t *testing.T) { + ctrl := gomock.NewController(t) + + // Simulate a cancellation that became visible after the batch row was persisted but before initialization completed. + request := testRequest() + request.State = entity.RequestStateCancelling + request.Version = 3 + batch := entity.Batch{ + ID: "test-queue/batch/1", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateCreating, + Version: 1, + } + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.ActiveBatchStates()).Return([]entity.Batch{batch}, nil) + batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(nil) + + // The batch's reverse-index row was already created by the interrupted attempt, so initialization must tolerate the duplicate. + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + batchDependentStore.EXPECT().Create(gomock.Any(), entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + }).Return(storage.ErrAlreadyExists) + + requestStore := storagemock.NewMockRequestStore(ctrl) + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + + publisher := queuemock.NewMockPublisher(ctrl) + queue := queuemock.NewMockQueue(ctrl) + queue.EXPECT().Publisher().Return(publisher).AnyTimes() + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: queue}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: queue}, + }) + require.NoError(t, err) + + analyzerFactory := conflictmock.NewMockFactory(ctrl) + controller := NewController( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, registry, countermock.NewMockCounter(ctrl), + store, analyzerFactory, topickey.TopicKeyBatch, "orchestrator-batch", + ) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(2).AnyTimes() + + // The delivery finishes structural initialization without publishing normal forward progress. + require.NoError(t, controller.Process(context.Background(), delivery)) +} + // Race-lost path: the cancel controller's markCancelling CAS landed first, // so the batch controller's request-claim CAS (Validated → Batched) fails // with storage.ErrVersionMismatch. The controller must ack the message (the @@ -495,9 +563,6 @@ func TestController_Process_CASLostToCancel(t *testing.T) { // Create must NOT be called — gomock fails if it is. mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - // The reverse-index Create still runs because it precedes the CAS; this is - // tolerated per the "downstream handles stale entries" contract. - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) @@ -508,6 +573,7 @@ func TestController_Process_CASLostToCancel(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() // Publisher with no EXPECTs — must not be called. @@ -548,7 +614,6 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { // Create must NOT be called — gomock fails if it is. mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) casErr := fmt.Errorf("db connection lost") mockReqStore := storagemock.NewMockRequestStore(ctrl) @@ -560,6 +625,7 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -575,13 +641,8 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { assert.True(t, errors.Is(err, casErr)) } -// Recovery path: a re-delivered batch message whose prior attempt CAS'd the -// request to RequestStateBatched but failed before BatchStore.Create. The -// halted check at the top of Process does NOT include Batched (Batched is -// forward-progress, not halted), so we reach the CAS again and re-bump the -// version on the request (Batched → Batched, version+1). The batch is then -// re-created with a new batch ID, which is tolerated per the existing -// duplicate-handling comment on BatchStore.Create. +// A prior attempt claimed the request but failed before persisting its batch. +// Redelivery performs a version-only Batched → Batched CAS before creating a replacement batch. func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { ctrl := gomock.NewController(t) @@ -590,8 +651,9 @@ 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().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil).AnyTimes() mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) @@ -605,6 +667,7 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(newRequestBatchStore(ctrl)).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) @@ -616,3 +679,424 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) } + +// A redelivery resumes every incomplete initialization step on the original Creating batch. +func TestController_Process_ResumesCreatingBatch(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + request.State = entity.RequestStateBatched + request.Version = 2 + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + Dependencies: []string{"test-queue/batch/1", "test-queue/batch/2"}, + State: entity.BatchStateCreating, + Version: 1, + } + + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + mockBatchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(nil) + + mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + mockBatchDependentStore.EXPECT().Create(gomock.Any(), entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + }).Return(storage.ErrAlreadyExists) + mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/1", + Dependents: []string{batch.ID}, + Version: 2, + }, nil) + mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/2", + Version: 4, + }, nil) + mockBatchDependentStore.EXPECT().UpdateDependents( + gomock.Any(), "test-queue/batch/2", int32(4), int32(5), []string{batch.ID}, + ).Return(nil) + + mockReqStore := storagemock.NewMockRequestStore(ctrl) + mockReqStore.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).AnyTimes() + requestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ + RequestID: request.ID, + BatchID: batch.ID, + Version: 1, + }).Return(storage.ErrAlreadyExists) + + mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() + + counter := countermock.NewMockCounter(ctrl) + controller := newTestController(t, ctrl, counter, mockStorage, nil, nil) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(2).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) +} + +// A redelivery after initialization but before a successful publish reuses and republishes the Created batch. +func TestController_Process_RepublishesCreatedBatch(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + request.State = entity.RequestStateBatched + request.Version = 2 + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateCreated, + Version: 2, + } + + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + mockReqStore := storagemock.NewMockRequestStore(ctrl) + mockReqStore.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) + + mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() + + counter := countermock.NewMockCounter(ctrl) + controller := newTestController(t, ctrl, counter, mockStorage, nil, nil) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(2).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) +} + +// A delayed batch delivery must not restart work after the assigned batch has become terminal. +func TestController_Process_TerminalAssignedBatchNoOp(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + request.State = entity.RequestStateBatched + request.Version = 2 + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateSucceeded, + Version: 4, + } + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + 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) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + + publisher := queuemock.NewMockPublisher(ctrl) + queue := queuemock.NewMockQueue(ctrl) + queue.EXPECT().Publisher().Return(publisher).AnyTimes() + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: queue}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: queue}, + }) + require.NoError(t, err) + + controller := NewController( + zaptest.NewLogger(t).Sugar(), tally.NoopScope, registry, countermock.NewMockCounter(ctrl), + store, conflictmock.NewMockFactory(ctrl), topickey.TopicKeyBatch, "orchestrator-batch", + ) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(3).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) +} + +// A terminal batch written before assignment persistence is recovered and backfilled rather than duplicated. +func TestController_Process_TerminalUnassignedBatchBackfillsAndNoOps(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + request.State = entity.RequestStateBatched + request.Version = 2 + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateSucceeded, + Version: 4, + } + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.AllBatchStates()).Return([]entity.Batch{batch}, nil) + + 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) + requestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ + RequestID: request.ID, + BatchID: batch.ID, + Version: 1, + }).Return(nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore).AnyTimes() + store.EXPECT().GetRequestStore().Return(requestStore).AnyTimes() + + controller := newTestController(t, ctrl, countermock.NewMockCounter(ctrl), store, nil, nil) + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(3).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) +} + +func TestFindAssignedBatchErrors(t *testing.T) { + request := testRequest() + request.State = entity.RequestStateBatched + storeErr := errors.New("storage failed") + + tests := []struct { + name string + mockFunc func(*gomock.Controller, *storagemock.MockStorage) + errMsg string + }{ + { + name: "assignment get fails", + mockFunc: func(ctrl *gomock.Controller, store *storagemock.MockStorage) { + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{}, storeErr) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + }, + errMsg: "failed to get batch assignment", + }, + { + name: "assigned batch get fails", + mockFunc: func(ctrl *gomock.Controller, store *storagemock.MockStorage) { + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(entity.RequestBatch{ + RequestID: request.ID, + BatchID: "test-queue/batch/1", + Version: 1, + }, nil) + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.Batch{}, storeErr) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + store.EXPECT().GetBatchStore().Return(batchStore) + }, + errMsg: "failed to get assigned batch", + }, + { + name: "fallback scan fails", + mockFunc: func(ctrl *gomock.Controller, store *storagemock.MockStorage) { + 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, storeErr) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + store.EXPECT().GetBatchStore().Return(batchStore) + }, + errMsg: "failed to get batches", + }, + { + name: "fallback scan is ambiguous", + mockFunc: func(ctrl *gomock.Controller, store *storagemock.MockStorage) { + 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: "test-queue/batch/1", Contains: []string{request.ID}}, + {ID: "test-queue/batch/2", Contains: []string{request.ID}}, + }, nil) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + store.EXPECT().GetBatchStore().Return(batchStore) + }, + errMsg: "contained by multiple batches", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockStorage(ctrl) + tt.mockFunc(ctrl, store) + + controller := &Controller{metricsScope: tally.NoopScope, store: store} + _, _, err := controller.findAssignedBatch(context.Background(), request) + assert.ErrorContains(t, err, tt.errMsg) + }) + } +} + +func TestEnsureRequestAssignmentsErrors(t *testing.T) { + batch := entity.Batch{ID: "test-queue/batch/1", Contains: []string{"test-queue/1"}} + storeErr := errors.New("storage failed") + + tests := []struct { + name string + mockFunc func(*storagemock.MockRequestBatchStore) + errMsg string + }{ + { + name: "create fails", + mockFunc: func(store *storagemock.MockRequestBatchStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(storeErr) + }, + errMsg: "failed to assign request", + }, + { + name: "duplicate verification fails", + mockFunc: func(store *storagemock.MockRequestBatchStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(storage.ErrAlreadyExists) + store.EXPECT().Get(gomock.Any(), "test-queue/1").Return(entity.RequestBatch{}, storeErr) + }, + errMsg: "failed to verify batch assignment", + }, + { + name: "duplicate conflicts", + mockFunc: func(store *storagemock.MockRequestBatchStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(storage.ErrAlreadyExists) + store.EXPECT().Get(gomock.Any(), "test-queue/1").Return(entity.RequestBatch{ + RequestID: "test-queue/1", + BatchID: "test-queue/batch/other", + Version: 1, + }, nil) + }, + errMsg: "already assigned", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + tt.mockFunc(requestBatchStore) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + + controller := &Controller{metricsScope: tally.NoopScope, store: store} + assert.ErrorContains(t, controller.ensureRequestAssignments(context.Background(), batch), tt.errMsg) + }) + } +} + +func TestInitializeBatchErrors(t *testing.T) { + batch := entity.Batch{ + ID: "test-queue/batch/1", + Contains: []string{"test-queue/1"}, + Dependencies: []string{"test-queue/batch/0"}, + State: entity.BatchStateCreating, + Version: 1, + } + storeErr := errors.New("storage failed") + + tests := []struct { + name string + mockFunc func(*storagemock.MockBatchStore, *storagemock.MockBatchDependentStore) + errMsg string + }{ + { + name: "dependent create fails", + mockFunc: func(_ *storagemock.MockBatchStore, store *storagemock.MockBatchDependentStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(storeErr) + }, + errMsg: "failed to create batch dependent index", + }, + { + name: "dependency get fails", + mockFunc: func(_ *storagemock.MockBatchStore, store *storagemock.MockBatchDependentStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + store.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{}, storeErr) + }, + errMsg: "failed to get batch dependent", + }, + { + name: "dependency update fails", + mockFunc: func(_ *storagemock.MockBatchStore, store *storagemock.MockBatchDependentStore) { + store.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + store.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Version: 2, + }, nil) + store.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{batch.ID}).Return(storeErr) + }, + errMsg: "failed to update batch dependent index", + }, + { + name: "created transition fails", + mockFunc: func(store *storagemock.MockBatchStore, dependentStore *storagemock.MockBatchDependentStore) { + dependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + dependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/0").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Dependents: []string{batch.ID}, + Version: 2, + }, nil) + store.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr) + }, + errMsg: "failed to mark batch", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + requestBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + batchStore := storagemock.NewMockBatchStore(ctrl) + batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + tt.mockFunc(batchStore, batchDependentStore) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestBatchStore().Return(requestBatchStore) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + store.EXPECT().GetBatchDependentStore().Return(batchDependentStore).AnyTimes() + + controller := &Controller{metricsScope: tally.NoopScope, store: store} + _, err := controller.initializeBatch(context.Background(), batch) + assert.ErrorContains(t, err, tt.errMsg) + }) + } +}