fix(cancel): resolve durable batch ownership

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",

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
+}

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)
+		})
+	}
+}

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",

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