Skip to content

feat(batch): make creation resumable #455

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ubettigole wants to merge 1 commit into from
+711 −103
Open

feat(batch): make creation resumable #455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
280 changes: 202 additions & 78 deletions submitqueue/orchestrator/controller/batch/batch.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"context"
"errors"
"fmt"
"slices"

"github.com/uber-go/tally"
entityqueue "github.com/uber/submitqueue/platform/base/messagequeue"
Expand Down Expand Up @@ -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",
Expand All @@ -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.
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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,
})
Expand All @@ -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)
Expand Down
Loading
Loading