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

Author: ubettigole

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)