From 61a7de79505196a3a58b205621a44ff7120360c8 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 26 Jul 2026 21:31:58 -0400 Subject: [PATCH] fix(session queue): count only rows actually transitioned by a bulk cancel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two concurrent bulk cancellations could select the same in-progress row; the first canceled it and the second, seeing the already-terminal row returned by _set_queue_item_status(), counted it again — both responses reporting that they canceled the same item. The terminal guard and the UPDATE already ran atomically (a single transaction under SqliteDatabase's process-wide lock); what was missing was exposing WHO performed the transition. _set_queue_item_status() now delegates to _transition_queue_item_status(), which additionally returns whether THIS call moved the row, and _cancel_in_progress_matching() counts only actual transitions — so exactly one racing caller observes itself as the canceller. The docstring records that the atomicity comes from the database lock, not the SQL, so a future move to per-thread connections doesn't silently reintroduce the double count. Includes a two-thread barrier test (in-memory SQLite) asserting each item is counted exactly once across concurrent bulk cancels; the pre-fix counter reports 4 for 2 items under the same interleaving. Follow-up to PR #9263 (JPPhoto review, 2026-07-25). Co-Authored-By: Claude Fable 5 --- .../session_queue/session_queue_sqlite.py | 50 +++++++++++++--- .../test_session_queue_multigpu_cancel.py | 60 +++++++++++++++++-- 2 files changed, 99 insertions(+), 11 deletions(-) diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 00e82b34064..6b63141a5af 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -472,6 +472,38 @@ def _set_queue_item_status( device: Optional[str] = None, queue_item: Optional[SessionQueueItem] = None, ) -> SessionQueueItem: + return self._transition_queue_item_status( + item_id=item_id, + status=status, + error_type=error_type, + error_message=error_message, + error_traceback=error_traceback, + device=device, + queue_item=queue_item, + )[0] + + def _transition_queue_item_status( + self, + item_id: int, + status: QUEUE_ITEM_STATUS, + error_type: Optional[str] = None, + error_message: Optional[str] = None, + error_traceback: Optional[str] = None, + device: Optional[str] = None, + queue_item: Optional[SessionQueueItem] = None, + ) -> tuple[SessionQueueItem, bool]: + """Move a queue item to `status` unless it is already finished (completed, failed or + canceled), returning the item and whether THIS call performed the transition. + + The terminal guard and the UPDATE run in one transaction, so two callers racing to cancel + the same row cannot both observe themselves as the one that canceled it — exactly one sees + `transitioned=True` (the bulk-cancel counters rely on this). Note the atomicity comes from + SqliteDatabase.transaction() holding its process-wide lock for the whole block, not from + the SQL: the guard is a Python check between the SELECT and the UPDATE, valid only while + every writer goes through that lock. When no transition happens, the item is returned + unchanged and no status-changed event is emitted; a vanished row raises + SessionQueueItemNotFoundError. + """ if queue_item is not None and queue_item.item_id != item_id: raise ValueError(f"Queue item {queue_item.item_id} does not match requested item {item_id}") @@ -519,7 +551,8 @@ def _set_queue_item_status( updated_status_row = cast(sqlite3.Row, cursor.fetchone()) if updated_status_row is None: - return self.get_queue_item(item_id) + # Already finished (return it unchanged) or deleted (get_queue_item raises). + return self.get_queue_item(item_id), False if queue_item is None: queue_item = self.get_queue_item(item_id) @@ -543,7 +576,7 @@ def _set_queue_item_status( ) self.__invoker.services.events.emit_queue_item_status_changed(queue_item, batch_status, queue_status) - return queue_item + return queue_item, True def _get_workflow_call_child_ids(self, item_id: int) -> list[int]: with self._db.transaction() as cursor: @@ -828,12 +861,15 @@ def _cancel_in_progress_matching(self, match_filter: str, params: list[Any]) -> canceled: list[int] = [] for item_id in item_ids: - # _set_queue_item_status no-ops (and returns the existing item) if the item finished - # between the SELECT and now, so count only the ones we actually moved to 'canceled'. - # It raises if the row vanished entirely (a concurrent clear/delete); such an item - # needs no cancellation, so skip it rather than failing the whole bulk operation. + # Count only the items THIS call actually moved to 'canceled'. An item that finished + # between the SELECT and now — including one canceled by a concurrent bulk request + # that selected the same row — is a no-op transition and must not be counted again. + # The transition raises if the row vanished entirely (a concurrent clear/delete); + # such an item needs no cancellation, so skip it rather than failing the whole bulk + # operation. try: - if self._set_queue_item_status(item_id, "canceled").status == "canceled": + _, transitioned = self._transition_queue_item_status(item_id, "canceled") + if transitioned: canceled.append(item_id) except SessionQueueItemNotFoundError: continue diff --git a/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py b/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py index fea29457839..c013d386c60 100644 --- a/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py +++ b/tests/app/services/session_queue/test_session_queue_multigpu_cancel.py @@ -206,15 +206,15 @@ def test_bulk_cancel_tolerates_rows_deleted_mid_cancel( _insert(session_queue, batch_id=batch_id) _dequeue_two_on_separate_devices(session_queue) - original_set_status = session_queue._set_queue_item_status + original_transition = session_queue._transition_queue_item_status - def delete_then_set_status(item_id: int, status: str, **kwargs): + def delete_then_transition(item_id: int, status: str, **kwargs): # Simulate the concurrent deletion landing just before the per-item cancel. with session_queue._db.transaction() as cursor: cursor.execute("DELETE FROM session_queue WHERE item_id = ?", (item_id,)) - return original_set_status(item_id, status, **kwargs) + return original_transition(item_id, status, **kwargs) - monkeypatch.setattr(session_queue, "_set_queue_item_status", delete_then_set_status) + monkeypatch.setattr(session_queue, "_transition_queue_item_status", delete_then_transition) result = session_queue.cancel_by_batch_ids("default", [batch_id]) @@ -223,6 +223,58 @@ def delete_then_set_status(item_id: int, status: str, **kwargs): assert result.canceled == 0 +def test_concurrent_bulk_cancels_count_each_item_once( + session_queue: SqliteSessionQueue, mock_invoker: Invoker, monkeypatch: pytest.MonkeyPatch +): + """Two bulk cancellations that select the same in-progress rows before either updates them must + not both report those rows as newly canceled. The terminal guard is part of the cancel UPDATE + itself, so for each row exactly one request observes itself as the canceller and the combined + count equals the number of items (JPPhoto follow-up on PR #9263).""" + import threading + + batch_id = str(uuid.uuid4()) + _insert(session_queue, batch_id=batch_id) + _insert(session_queue, batch_id=batch_id) + id_a, id_b = _dequeue_two_on_separate_devices(session_queue) + + # Rendezvous both requests after their item-id SELECTs but before either per-item update, so + # both have selected the same still-in-progress rows. + barrier = threading.Barrier(2) + arrived = threading.local() + original_transition = session_queue._transition_queue_item_status + + def rendezvous_then_transition(item_id: int, status: str, **kwargs): + if not getattr(arrived, "done", False): + arrived.done = True + barrier.wait(timeout=10) + return original_transition(item_id, status, **kwargs) + + monkeypatch.setattr(session_queue, "_transition_queue_item_status", rendezvous_then_transition) + + results = [] + errors = [] + + def run() -> None: + try: + results.append(session_queue.cancel_by_batch_ids("default", [batch_id])) + except Exception as e: # noqa: BLE001 - surface any thread failure in the assertion below + errors.append(e) + + threads = [threading.Thread(target=run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert not errors + assert len(results) == 2 + # Each of the two items is counted by exactly one request (previously both requests could + # report the same item, summing to 4). + assert sum(r.canceled for r in results) == 2 + assert session_queue.get_queue_item(id_a).status == "canceled" + assert session_queue.get_queue_item(id_b).status == "canceled" + + def test_cancel_by_batch_ids_respects_user_scope(session_queue: SqliteSessionQueue, mock_invoker: Invoker): """A user-scoped cancel must not cancel another user's in-progress item in the same batch.""" batch_id = str(uuid.uuid4())