fix: cancel in progress has incorrect ordering with new in-memory index

.github/workflows/sdk-python.yml

@@ -64,6 +64,7 @@ jobs:
       matrix:
         python-version: ${{ github.event_name == 'pull_request' && fromJSON('["3.14"]') || fromJSON('["3.10", "3.11", "3.12", "3.13", "3.14"]') }}
         optimistic-scheduling: ["true", "false"]
+        concurrency-in-memory-index: ["true", "false"]
     timeout-minutes: 20
     steps:
       - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
@@ -105,6 +106,7 @@ jobs:
           export SERVER_DEFAULT_ENGINE_VERSION=V1
           export SERVER_MSGQUEUE_RABBITMQ_URL="amqp://user:password@localhost:5672/"
           export SERVER_OPTIMISTIC_SCHEDULING_ENABLED=${{ matrix.optimistic-scheduling }}
+          export SERVER_CONCURRENCY_IN_MEMORY_INDEX_ENABLED=${{ matrix.concurrency-in-memory-index }}
           export SERVER_OBSERVABILITY_ENABLED=true
 
           go run ./cmd/hatchet-admin quickstart
@@ -161,14 +163,14 @@ jobs:
         if: always()
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
         with:
-          name: ${{ env.HATCHET_CLIENT_NAMESPACE }}-opt-${{ matrix.optimistic-scheduling }}-engine-logs
+          name: ${{ env.HATCHET_CLIENT_NAMESPACE }}-opt-${{ matrix.optimistic-scheduling }}-cimi-${{ matrix.concurrency-in-memory-index }}-engine-logs
           path: engine.log
 
       - name: Upload API logs
         if: always()
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
         with:
-          name: ${{ env.HATCHET_CLIENT_NAMESPACE }}-opt-${{ matrix.optimistic-scheduling }}-api-logs
+          name: ${{ env.HATCHET_CLIENT_NAMESPACE }}-opt-${{ matrix.optimistic-scheduling }}-cimi-${{ matrix.concurrency-in-memory-index }}-api-logs
           path: api.log
 
   publish:

internal/services/scheduler/v1/scheduler.go

@@ -723,6 +723,12 @@ func (s *Scheduler) notifyAfterConcurrency(ctx context.Context, tenantId uuid.UU
 
 	s.pool.NotifyQueues(ctx, tenantId, queues)
 
+	// The in-memory concurrency flush releases the task runtime in-transaction, so the
+	// MsgIDTaskCancelled handler below (handleTaskCancelled -> CancelTasks) can no longer resolve the
+	// worker for a task that was running when it got cancelled. Signal those workers' dispatchers
+	// directly using the worker id captured during the flush, mirroring sendTaskCancellationsToDispatcher.
+	s.signalWorkersToCancelInProgress(ctx, tenantId, res.Cancelled)
+
 	// handle cancellations
 	for _, cancelled := range res.Cancelled {
 		eventType := sqlcv1.V1EventTypeOlapCANCELLED
@@ -766,6 +772,74 @@ func (s *Scheduler) notifyAfterConcurrency(ctx context.Context, tenantId uuid.UU
 	}
 }
 
+// signalWorkersToCancelInProgress tells the relevant dispatchers to cancel tasks that were running on
+// a worker when a concurrency strategy cancelled them. The in-memory concurrency flush releases the
+// task runtime in its transaction, so the MsgIDTaskCancelled handler (handleTaskCancelled) can't map
+// these tasks back to a worker; we do it here from the worker id captured during the flush. Tasks that
+// were only queued (no runtime) carry a nil worker id and are skipped - their cancellation still flows
+// through the MsgIDTaskCancelled handler. This mirrors sendTaskCancellationsToDispatcher in the tasks
+// controller.
+func (s *Scheduler) signalWorkersToCancelInProgress(ctx context.Context, tenantId uuid.UUID, cancelled []repov1.TaskWithCancelledReason) {
+	signals := make([]tasktypes.SignalTaskCancelledPayload, 0, len(cancelled))
+	workerIds := make([]uuid.UUID, 0, len(cancelled))
+
+	for _, c := range cancelled {
+		if c.WorkerId == uuid.Nil {
+			continue
+		}
+
+		signals = append(signals, tasktypes.SignalTaskCancelledPayload{
+			TaskId:     c.Id,
+			InsertedAt: c.InsertedAt,
+			RetryCount: c.RetryCount,
+			WorkerId:   c.WorkerId,
+		})
+		workerIds = append(workerIds, c.WorkerId)
+	}
+
+	if len(signals) == 0 {
+		return
+	}
+
+	workerIdToDispatcherId, _, err := s.repov1.Workers().GetDispatcherIdsForWorkers(ctx, tenantId, workerIds)
+
+	if err != nil {
+		s.l.Error().Ctx(ctx).Err(err).Msg("could not list dispatcher ids for cancelled in-progress tasks")
+		return
+	}
+
+	dispatcherIdsToPayloads := make(map[uuid.UUID][]tasktypes.SignalTaskCancelledPayload)
+
+	for _, sig := range signals {
+		dispatcherId, ok := workerIdToDispatcherId[sig.WorkerId]
+
+		if !ok {
+			continue
+		}
+
+		dispatcherIdsToPayloads[dispatcherId] = append(dispatcherIdsToPayloads[dispatcherId], sig)
+	}
+
+	for dispatcherId, payloads := range dispatcherIdsToPayloads {
+		msg, err := msgqueue.NewTenantMessage(
+			tenantId,
+			msgqueue.MsgIDTaskCancelled,
+			false,
+			true,
+			payloads...,
+		)
+
+		if err != nil {
+			s.l.Error().Ctx(ctx).Err(err).Msg("could not create task cancellation signal for dispatcher")
+			continue
+		}
+
+		if err := s.mq.SendMessage(ctx, msgqueue.QueueTypeFromDispatcherID(dispatcherId), msg); err != nil {
+			s.l.Error().Ctx(ctx).Err(err).Msg("could not send task cancellation signal to dispatcher")
+		}
+	}
+}
+
 func taskBulkAssignedTask(tenantId uuid.UUID, workerIdsToTaskIds map[uuid.UUID][]int64) (*msgqueue.Message, error) {
 	return msgqueue.NewTenantMessage(
 		tenantId,

pkg/repository/scheduler_concurrency.go

@@ -36,6 +36,12 @@ type TaskWithCancelledReason struct {
 	TaskExternalId uuid.UUID
 
 	WorkflowRunId uuid.UUID
+
+	// WorkerId is the worker running the task at the moment it was cancelled, or uuid.Nil if the task
+	// was not running (queued). It's captured here because the concurrency flush releases the task
+	// runtime in-transaction, so the downstream MsgIDTaskCancelled handler can no longer resolve the
+	// worker; the scheduler uses it to signal the dispatcher directly for in-progress cancellations.
+	WorkerId uuid.UUID
 }
 
 // CancelledSlotInput identifies a concurrency slot to cancel along with the reason it's being
@@ -992,6 +998,10 @@ func (c *ConcurrencyRepositoryImpl) UpdateConcurrencySlotsTx(
 			CancelledReason: reason,
 			TaskExternalId:  released.ExternalID,
 			WorkflowRunId:   released.WorkflowRunID,
+			// releaseTasks deleted v1_task_runtime here, so the downstream MsgIDTaskCancelled handler
+			// can no longer resolve the worker for an in-progress cancellation. Capture the worker id
+			// now so the scheduler can signal the dispatcher directly (see notifyAfterConcurrency).
+			WorkerId: released.WorkerID,
 		})
 	}
 

pkg/scheduling/v1/concurrency/cancel_in_progress_test.go

@@ -89,6 +89,82 @@ func TestCancelInProgress_RanksByPriority(t *testing.T) {
 	}
 }
 
+// Among equal-priority slots, CANCEL_IN_PROGRESS keeps the NEWEST maxRuns and cancels the oldest
+// (matching the legacy runCancelInProgress ORDER BY sort_id DESC). This is the case that actually
+// exercises the recency tiebreak - if it were inverted we would keep the oldest instead.
+func TestCancelInProgress_KeepsNewestAmongEqualPriority(t *testing.T) {
+	now := time.Now().UTC()
+	future := now.Add(time.Hour)
+
+	repo := &mockConcurrencyRepo{}
+	c := newCancelInProgressStrategy(repo, 2)
+
+	// same priority, same inserted_at; recency is distinguished by taskId (higher = newer).
+	msgs := []walMessage{
+		walInsert("a", 10, 5, now, future), // oldest
+		walInsert("a", 11, 5, now, future),
+		walInsert("a", 12, 5, now, future), // newest
+	}
+
+	if _, err := c.processWALMessages(context.Background(), nil, msgs); err != nil {
+		t.Fatalf("processWALMessages: %v", err)
+	}
+
+	// newest two (12, 11) run; oldest (10) cancelled
+	filled := filledIDs(repo.lastFilled)
+	if len(filled) != 2 || !containsID(filled, 12) || !containsID(filled, 11) {
+		t.Fatalf("filled = %v, want {11,12} (newest maxRuns)", filled)
+	}
+	cancelled := cancelledByReason(repo.lastCancelled, repository.CancelledReasonConcurrencyLimit)
+	if len(cancelled) != 1 || !containsID(cancelled, 10) {
+		t.Fatalf("cancelled = %v, want [10] (oldest)", cancelled)
+	}
+}
+
+// At capacity with equal priorities, a newer arrival preempts the OLDEST running task (in-progress
+// cancellation) - the defining CANCEL_IN_PROGRESS behaviour. Distinct from the priority-driven
+// preemption above; this pins the recency direction for running slots.
+func TestCancelInProgress_NewerArrivalPreemptsOldestRunning(t *testing.T) {
+	now := time.Now().UTC()
+	future := now.Add(time.Hour)
+
+	repo := &mockConcurrencyRepo{
+		indexRows: []*sqlcv1.ListConcurrencySlotsForIndexingRow{
+			indexRow("a", 1, 5, 0, now, future, true), // running, oldest
+			indexRow("a", 2, 5, 0, now, future, true), // running
+		},
+	}
+	c := newCancelInProgressStrategy(repo, 2)
+
+	if err := c.buildIndex(context.Background()); err != nil {
+		t.Fatalf("buildIndex: %v", err)
+	}
+
+	// a newer equal-priority slot arrives at a full group
+	msgs := []walMessage{walInsert("a", 3, 5, now, future)}
+
+	if _, err := c.processWALMessages(context.Background(), nil, msgs); err != nil {
+		t.Fatalf("processWALMessages: %v", err)
+	}
+
+	// newcomer (3) promoted; oldest runner (1) preempted
+	if got := filledIDs(repo.lastFilled); len(got) != 1 || got[0] != 3 {
+		t.Fatalf("filled = %v, want [3]", got)
+	}
+	cancelled := cancelledByReason(repo.lastCancelled, repository.CancelledReasonConcurrencyLimit)
+	if len(cancelled) != 1 || !containsID(cancelled, 1) {
+		t.Fatalf("cancelled = %v, want [1] (oldest running preempted)", cancelled)
+	}
+
+	sq := c.getOrCreateSubQueue("a")
+	if _, ok := sq.running.get(1); ok {
+		t.Fatalf("oldest task 1 should have been preempted")
+	}
+	if _, ok := sq.running.get(3); !ok {
+		t.Fatalf("newcomer 3 should be running")
+	}
+}
+
 // Queued slots past their scheduling timeout are cancelled as SCHEDULING_TIMED_OUT and excluded from
 // the run/cancel ranking, even if they would otherwise outrank the survivors.
 func TestCancelInProgress_TimedOutExcludedFromRanking(t *testing.T) {