diff --git a/backend/workflow_manager/execution/tests/test_update_status_locking.py b/backend/workflow_manager/execution/tests/test_update_status_locking.py new file mode 100644 index 0000000000..d3f33d7f51 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py @@ -0,0 +1,533 @@ +"""Row locking on the workflow-execution status write. + +For PG-transport executions, ``update_status`` used to take ``SELECT ... FOR UPDATE`` +on the execution row twice (once in the view, again inside ``update_execution()``) and ran the file-row COUNTs +while holding it. ``FOR UPDATE`` conflicts with the ``FOR KEY SHARE`` every child-row +INSERT takes for its foreign-key check, so file-execution inserts for the same +execution queued behind each status write. + +These tests pin the fix: ONE ``FOR NO KEY UPDATE`` per call, one COUNT query (taken +after the lock, so it sees what any writer it waited behind committed), child inserts +not blocked — and status writers still serialized against each other, +so the terminal-one-way guard and the attempt counter stay correct under concurrency. +""" + +import threading +import time +import uuid +from unittest.mock import MagicMock, patch + +import pytest +from django.db import connection, connections, transaction +from django.test import TestCase, TransactionTestCase +from django.test.utils import CaptureQueriesContext + +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +# Generous: a hang must fail the test rather than the CI job. +WAIT_SECONDS = 10 + + +def _update_status(execution_id, data): + """Call the endpoint the way a worker does, on a fresh instance (as get_object + would load it). + """ + view = WorkflowExecutionInternalViewSet() + req = MagicMock() + req.data = data + fresh = WorkflowExecution.objects.get(pk=execution_id) + with patch.object(view, "get_object", return_value=fresh): + return view.update_status(req, id=str(execution_id)) + + +def _lock_statements(queries): + sqls = [q["sql"].upper() for q in queries] + no_key = [s for s in sqls if "FOR NO KEY UPDATE" in s] + # A plain FOR UPDATE, i.e. one that is not part of FOR NO KEY UPDATE. + full = [s for s in sqls if "FOR UPDATE" in s and "FOR NO KEY UPDATE" not in s] + return sqls, no_key, full + + +class SingleLockTests(TestCase): + def setUp(self): + self.wf = Workflow.objects.create(workflow_name="wf-lock") + + def _exec(self, status=ExecutionStatus.EXECUTING, pg=True, files=()): + ex = WorkflowExecution.objects.create( + workflow=self.wf, + status=status, + queue_message_id=123 if pg else None, + task_id=None if pg else uuid.uuid4(), + ) + for fstatus in files: + WorkflowFileExecution.objects.create( + workflow_execution=ex, file_name="f", status=fstatus.value + ) + return ex + + def _assert_one_no_key_lock(self, queries): + sqls, no_key, full = _lock_statements(queries) + assert len(no_key) == 1, sqls + assert full == [], full + return sqls + + def test_pg_update_status_takes_exactly_one_no_key_lock(self): + ex = self._exec(files=[ExecutionStatus.COMPLETED]) + with CaptureQueriesContext(connection) as ctx: + out = _update_status( + ex.id, + { + "status": ExecutionStatus.COMPLETED.value, + "total_files": 1, + "successful_files": 1, + "failed_files": 0, + }, + ) + assert out.data["status"] == "updated" + self._assert_one_no_key_lock(ctx.captured_queries) + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_celery_update_status_takes_exactly_one_no_key_lock(self): + ex = self._exec(pg=False, files=[ExecutionStatus.COMPLETED]) + with CaptureQueriesContext(connection) as ctx: + _update_status( + ex.id, + {"status": ExecutionStatus.COMPLETED.value, "total_files": 1}, + ) + self._assert_one_no_key_lock(ctx.captured_queries) + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_file_row_count_runs_once_and_after_the_lock(self): + ex = self._exec(files=[ExecutionStatus.COMPLETED, ExecutionStatus.ERROR]) + with CaptureQueriesContext(connection) as ctx: + _update_status( + ex.id, + { + "status": ExecutionStatus.ERROR.value, + "total_files": 1, + "failed_files": 0, + }, + ) + sqls = self._assert_one_no_key_lock(ctx.captured_queries) + lock_at = next(i for i, s in enumerate(sqls) if "FOR NO KEY UPDATE" in s) + counts = [ + i + for i, s in enumerate(sqls) + if "COUNT(" in s and "WORKFLOW_FILE_EXECUTION" in s + ] + assert len(counts) == 1, sqls + assert counts[0] > lock_at + # The floor still applies: the rows say 2 total / 1 errored. + ex.refresh_from_db() + assert ex.total_files == 2 + assert ex.failed_files == 1 + + def test_no_count_when_no_floored_counter_is_written(self): + ex = self._exec(files=[ExecutionStatus.COMPLETED]) + with CaptureQueriesContext(connection) as ctx: + _update_status(ex.id, {"status": ExecutionStatus.EXECUTING.value}) + assert not any("COUNT(" in q["sql"].upper() for q in ctx.captured_queries) + + def test_refused_pg_write_still_takes_one_lock(self): + ex = self._exec(status=ExecutionStatus.COMPLETED) + with CaptureQueriesContext(connection) as ctx: + out = _update_status(ex.id, {"status": ExecutionStatus.EXECUTING.value}) + assert out.data.get("reason") == "already_final" + self._assert_one_no_key_lock(ctx.captured_queries) + + def test_count_error_fails_the_request_and_commits_nothing(self): + """Any error in the count, not only a database one, fails the request + rather than committing counters the floor never checked. + """ + ex = self._exec(files=[ExecutionStatus.ERROR]) + + def broken_count(execution): + raise RuntimeError("count failed") + + with patch.object( + WorkflowExecutionInternalViewSet, + "_count_file_rows", + staticmethod(broken_count), + ): + out = _update_status( + ex.id, + { + "status": ExecutionStatus.ERROR.value, + "total_files": 0, + "failed_files": 0, + }, + ) + assert out.status_code == 500, out.data + ex.refresh_from_db() + assert ex.status == ExecutionStatus.EXECUTING.value + assert ex.total_files == 0 + assert not ex.failed_files + + def test_failed_count_query_fails_the_request_and_commits_nothing(self): + """A statement error in the count aborts the transaction, so the + request fails (500, retried by the worker) instead of committing counters + the floor never checked. Nothing is half-written. + """ + ex = self._exec(files=[ExecutionStatus.ERROR]) + calls = [] + + def broken_aggregate(*args, **kwargs): + calls.append(1) + with connection.cursor() as cur: + cur.execute("SELECT 1/0") + + with patch( + "django.db.models.query.QuerySet.aggregate", + side_effect=broken_aggregate, + autospec=True, + ): + out = _update_status( + ex.id, + { + "status": ExecutionStatus.ERROR.value, + "total_files": 0, + "failed_files": 0, + }, + ) + assert out.status_code == 500, out.data + assert len(calls) == 1 # tried once, not retried under the lock + ex.refresh_from_db() + assert ex.status == ExecutionStatus.EXECUTING.value + assert ex.total_files == 0 # the floor would have made it 1 + assert not ex.failed_files + + def test_reaper_recovery_takes_one_no_key_lock(self): + from datetime import timedelta + + from django.utils import timezone + + ex = self._exec(files=[ExecutionStatus.COMPLETED, ExecutionStatus.ERROR]) + stale = timezone.now() - timedelta(seconds=9999) + WorkflowExecution.objects.filter(pk=ex.pk).update(modified_at=stale) + WorkflowFileExecution.objects.filter(workflow_execution=ex).update( + modified_at=stale + ) + view = WorkflowExecutionInternalViewSet() + req = MagicMock() + req.data = {"stuck_seconds": 60, "limit": 10} + with CaptureQueriesContext(connection) as ctx: + out = view.recover_stuck_pg_executions(req).data + assert out["recovered"] == 1 + sqls, no_key, full = _lock_statements(ctx.captured_queries) + assert len(no_key) == 1 and full == [], sqls + ex.refresh_from_db() + assert ex.status == ExecutionStatus.ERROR.value + assert ex.total_files == 2 and ex.failed_files == 1 + + def test_update_execution_locks_once_by_itself(self): + ex = self._exec() + with CaptureQueriesContext(connection) as ctx: + ex.update_execution(status=ExecutionStatus.COMPLETED) + self._assert_one_no_key_lock(ctx.captured_queries) + + def test_row_locked_skips_the_second_lock(self): + ex = self._exec() + with transaction.atomic(): + locked = WorkflowExecution.objects.select_for_update(no_key=True).get( + pk=ex.pk + ) + with CaptureQueriesContext(connection) as ctx: + locked.update_execution( + status=ExecutionStatus.COMPLETED, + row_locked=True, + ) + _, no_key, full = _lock_statements(ctx.captured_queries) + assert no_key == [] and full == [] + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_row_locked_outside_a_transaction_is_rejected(self): + ex = self._exec() + # TestCase wraps each test in a transaction; step outside it the way a + # real autocommit caller would be. + with patch.object(connections["default"], "in_atomic_block", False): + with pytest.raises(RuntimeError, match="row_locked"): + ex.update_execution(status=ExecutionStatus.COMPLETED, row_locked=True) + ex.refresh_from_db() + assert ex.status == ExecutionStatus.EXECUTING.value + + +class ConcurrentStatusWriteTests(TransactionTestCase): + """Real concurrent transactions on separate connections.""" + + def setUp(self): + self.wf = Workflow.objects.create(workflow_name="wf-lock-concurrent") + + def _exec(self, status=ExecutionStatus.EXECUTING, pg=True): + return WorkflowExecution.objects.create( + workflow=self.wf, + status=status, + queue_message_id=123 if pg else None, + task_id=None if pg else uuid.uuid4(), + ) + + @staticmethod + def _run(target, *args): + """Run *target* on its own thread (own DB connection); collect its result + or exception. + """ + box = {} + + def runner(): + try: + box["result"] = target(*args) + except BaseException as e: # noqa: BLE001 — surfaced by the test + box["error"] = e + finally: + connections.close_all() + + t = threading.Thread(target=runner, daemon=True) + t.start() + return t, box + + @staticmethod + def _wait_until_a_backend_waits_on_a_lock(thread, box): + """Block until some other connection to this database is waiting on a + lock, so "it waited" is observed rather than inferred from a timeout. + """ + deadline = time.monotonic() + WAIT_SECONDS + while time.monotonic() < deadline: + assert thread.is_alive(), f"finished without waiting for the lock: {box}" + with connection.cursor() as cur: + cur.execute( + "SELECT count(*) FROM pg_stat_activity " + "WHERE datname = current_database() " + "AND pid <> pg_backend_pid() AND wait_event_type = 'Lock'" + ) + if cur.fetchone()[0]: + return + time.sleep(0.05) + raise AssertionError(f"no backend ever waited on the row lock: {box}") + + def _hold_lock_inside_update_status(self, execution_id, data): + """Start an update_status that pauses while it holds the row lock. + + Returns (thread, box, locked_event, release_event). The pause is inside + ``_update_file_aggregates`` — after the status write, before commit. + """ + locked, release = threading.Event(), threading.Event() + original = WorkflowExecutionInternalViewSet._update_file_aggregates + state = {"thread": None} + + def paused(execution, validated_data): + if threading.get_ident() == state["thread"]: + locked.set() + assert release.wait(WAIT_SECONDS), "holder was never released" + return original(execution, validated_data) + + patcher = patch.object( + WorkflowExecutionInternalViewSet, + "_update_file_aggregates", + staticmethod(paused), + ) + patcher.start() + self.addCleanup(patcher.stop) + + def target(): + state["thread"] = threading.get_ident() + return _update_status(execution_id, data) + + t, box = self._run(target) + assert locked.wait(WAIT_SECONDS), f"holder never reached the lock: {box}" + return t, box, release + + def test_child_insert_is_not_blocked_by_a_status_write(self): + """The fix itself: a file-execution INSERT (FK check → FOR KEY SHARE on the + execution row) must not wait for an in-flight status write. + """ + for pg in (True, False): + ex = self._exec(pg=pg) + t, box, release = self._hold_lock_inside_update_status( + ex.id, {"status": ExecutionStatus.COMPLETED.value, "total_files": 1} + ) + try: + + def insert(execution_id=ex.id): + with transaction.atomic(): + with connection.cursor() as cur: + # Fail fast instead of hanging if the insert queues. + cur.execute("SET LOCAL lock_timeout = '2s'") + WorkflowFileExecution.objects.create( + workflow_execution_id=execution_id, + file_name="child", + status=ExecutionStatus.PENDING.value, + ) + return True + + it, ibox = self._run(insert) + it.join(WAIT_SECONDS) + assert ( + ibox.get("result") is True + ), f"child insert blocked by the status lock (pg={pg}): {ibox}" + finally: + release.set() + t.join(WAIT_SECONDS) + assert box.get("result") is not None, box + assert box["result"].data["status"] == "updated" + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + assert ( + WorkflowFileExecution.objects.filter(workflow_execution=ex).count() == 1 + ) + + def test_count_sees_file_errors_committed_while_waiting_for_the_lock(self): + """A callback that waited behind another writer must floor its counters at + what that writer committed. Here the holder (reaper-style) cascades the + files to ERROR; the waiting success callback corrects the status to + COMPLETED (ERROR is correctable) but must not write failed_files=0. + """ + ex = self._exec() + WorkflowFileExecution.objects.create( + workflow_execution=ex, + file_name="f", + status=ExecutionStatus.EXECUTING.value, + ) + t, box, release = self._hold_lock_inside_update_status( + ex.id, + { + "status": ExecutionStatus.ERROR.value, + "cascade_terminal_files": True, + "error_message": "reaper", + }, + ) + st, sbox = self._run( + _update_status, + ex.id, + { + "status": ExecutionStatus.COMPLETED.value, + "total_files": 1, + "successful_files": 1, + "failed_files": 0, + }, + ) + try: + self._wait_until_a_backend_waits_on_a_lock(st, sbox) + finally: + release.set() + t.join(WAIT_SECONDS) + st.join(WAIT_SECONDS) + assert box["result"].data["status"] == "updated", box + assert sbox["result"].data["status"] == "updated", sbox + ex.refresh_from_db() + assert ex.failed_files == 1, "count missed rows committed during the wait" + + def test_uncommitted_file_insert_holds_no_lock_on_the_execution_row(self): + """Why counting after the lock can't "wait for" in-flight file inserts, under + either lock mode: Django's foreign keys are DEFERRABLE INITIALLY DEFERRED, + so an open INSERT takes its FOR KEY SHARE only at commit. Even main's plain + FOR UPDATE is granted at once, and the open row is invisible to the count. + """ + ex = self._exec() + inserted, release = threading.Event(), threading.Event() + + def insert_batch(execution_id=ex.id): + with transaction.atomic(): + WorkflowFileExecution.objects.create( + workflow_execution_id=execution_id, + file_name="new", + status=ExecutionStatus.PENDING.value, + ) + inserted.set() + assert release.wait(WAIT_SECONDS), "inserter was never released" + return True + + it, ibox = self._run(insert_batch) + assert inserted.wait(WAIT_SECONDS), ibox + + def lock_and_count(execution_id=ex.id): + with transaction.atomic(): + with connection.cursor() as cur: + cur.execute("SET LOCAL lock_timeout = '2s'") + WorkflowExecution.objects.select_for_update().get(pk=execution_id) + return WorkflowFileExecution.objects.filter( + workflow_execution_id=execution_id + ).count() + + try: + lt, lbox = self._run(lock_and_count) + lt.join(WAIT_SECONDS) + assert lbox.get("result") == 0, lbox + finally: + release.set() + it.join(WAIT_SECONDS) + assert ibox.get("result") is True, ibox + + def test_status_writers_still_exclude_each_other(self): + """A second writer must wait for the first to commit, then see its result: + here the straggler's EXECUTING is refused over the committed COMPLETED. + """ + ex = self._exec() + t, box, release = self._hold_lock_inside_update_status( + ex.id, {"status": ExecutionStatus.COMPLETED.value} + ) + st, sbox = self._run( + _update_status, ex.id, {"status": ExecutionStatus.EXECUTING.value} + ) + try: + self._wait_until_a_backend_waits_on_a_lock(st, sbox) + finally: + release.set() + t.join(WAIT_SECONDS) + st.join(WAIT_SECONDS) + assert box["result"].data["status"] == "updated", box + assert sbox["result"].data.get("reason") == "already_final", sbox + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_racing_writers_never_revert_completed(self): + """Many concurrent writers, one COMPLETED among stale EXECUTINGs: whatever + the interleaving, COMPLETED is the final status. + """ + for _ in range(3): + ex = self._exec() + start = threading.Barrier(8) + + def write(status, start=start, execution_id=ex.id): + start.wait(WAIT_SECONDS) + return _update_status(execution_id, {"status": status.value}) + + runs = [self._run(write, ExecutionStatus.COMPLETED)] + [ + self._run(write, ExecutionStatus.EXECUTING) for _ in range(7) + ] + for t, _ in runs: + t.join(WAIT_SECONDS) + assert not t.is_alive(), "a writer hung" + for _, b in runs: + assert "error" not in b, b + ex.refresh_from_db() + assert ex.status == ExecutionStatus.COMPLETED.value + + def test_concurrent_attempt_increments_are_not_lost(self): + """PG path increments from the locked value, so N racing increments land + as N — only true if the single lock still serializes the writers. + """ + ex = self._exec() + n = 6 + start = threading.Barrier(n) + + def bump(): + start.wait(WAIT_SECONDS) + # attempts=99 > any snapshot → increment_attempt is True every time. + return _update_status( + ex.id, {"status": ExecutionStatus.EXECUTING.value, "attempts": 99} + ) + + runs = [self._run(bump) for _ in range(n)] + for t, _ in runs: + t.join(WAIT_SECONDS) + assert not t.is_alive(), "a writer hung" + for _, b in runs: + assert "error" not in b and b["result"].data["status"] == "updated", b + ex.refresh_from_db() + assert ex.attempts == n diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 8b568818e9..bbb7259d69 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -548,6 +548,7 @@ def update_status(self, request, id=None): # with failed_files=None, which is_failure_run() reads as a success and # silently bypasses notify_on_failures subscribers. with transaction.atomic(): + row_locked = False # Concurrency guard for PG-transport executions, identified by the # durable ``queue_message_id`` marker — NOT a re-evaluated Flipt flag. # Keying on the row's own transport keeps the invariant correct even @@ -572,10 +573,18 @@ def update_status(self, request, id=None): # succeeded — blocking that would freeze a genuinely successful run at # a premature ERROR. Celery executions (queue_message_id NULL) keep the # legacy path. + # + # ONE lock per call: the row read here is handed to update_execution() + # with row_locked=True, so it does not lock the same row a second time. The + # Celery path takes no lock here and lets update_execution() take its + # own — still one. FOR NO KEY UPDATE (no_key=True) keeps writers + # mutually exclusive without blocking the FOR KEY SHARE that every + # child-row INSERT takes on this row for its foreign-key check. if execution.queue_message_id is not None: - execution = WorkflowExecution.objects.select_for_update().get( - pk=execution.pk - ) + execution = WorkflowExecution.objects.select_for_update( + no_key=True + ).get(pk=execution.pk) + row_locked = True protected = { ExecutionStatus.COMPLETED.value, ExecutionStatus.STOPPED.value, @@ -602,9 +611,18 @@ def update_status(self, request, id=None): status=status_enum, error=error_message, increment_attempt=increment_attempt, - ) - - # File aggregates are not handled by update_execution; persist separately. + row_locked=row_locked, + ) + + # File aggregates are not handled by update_execution; persist + # separately. The floor's file-row count runs here, AFTER the lock, + # on purpose: any status writer we waited behind (e.g. the reaper + # cascading files to ERROR) has committed by now, and a count taken + # before the wait would miss its rows and let failed_files be written + # too low. (A file row whose INSERT has not committed is not counted + # either way: Django's foreign keys are DEFERRABLE INITIALLY DEFERRED, + # so an open insert holds no lock on this row until it commits, and + # FOR UPDATE never waited for it.) self._update_file_aggregates(execution, validated_data) # Reaper recovery: cascade the terminal status to any non-terminal @@ -791,7 +809,9 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str: try: with transaction.atomic(): - execution = WorkflowExecution.objects.select_for_update().get(pk=exec_id) + execution = WorkflowExecution.objects.select_for_update(no_key=True).get( + pk=exec_id + ) # Re-check under lock: skip if it reached terminal or was just # touched (no longer stuck) — never race a live execution. if execution.status in terminal or execution.modified_at >= cutoff: @@ -826,7 +846,9 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str: computed = ExecutionStatus.STOPPED else: computed = ExecutionStatus.COMPLETED - execution.update_execution(status=computed) + # Already locked above; hand the locked read over so the row is + # not locked a second time. + execution.update_execution(status=computed, row_locked=True) self._restamp_execution_time_from_files(execution) self._update_file_aggregates( execution, @@ -952,6 +974,26 @@ def _update_file_aggregates(execution, validated_data) -> None: ) execution.save(update_fields=update_fields) + @staticmethod + def _count_file_rows(execution) -> dict: + """Count this execution's file rows, total and ERROR, in one query. + + Returns ``{"total": int, "errored": int}``. Errors are left to propagate: + the request then fails and the transaction rolls back, so the worker retries + instead of committing counters the floor never checked (an under-counted + failed_files would read as a success and skip failure notifications). + """ + from django.db.models import Count, Q + + from workflow_manager.workflow_v2.enums import ExecutionStatus + + return WorkflowFileExecution.objects.filter( + workflow_execution_id=execution.pk + ).aggregate( + total=Count("id"), + errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)), + ) + @staticmethod def _floor_counters_at_file_rows(execution, update_fields: "list[str]") -> None: """Raise the counters being written to match the file rows on disk. @@ -964,23 +1006,11 @@ def _floor_counters_at_file_rows(execution, update_fields: "list[str]") -> None: if not (wants_failed or wants_total): return - from workflow_manager.workflow_v2.enums import ExecutionStatus - - try: - rows = execution.file_executions - error_rows = rows.filter(status=ExecutionStatus.ERROR.value).count() - row_total = rows.count() - except Exception: - logger.exception( - "Could not read file rows to floor the counters for execution %s", - execution.id, - ) - return - + counts = WorkflowExecutionInternalViewSet._count_file_rows(execution) if wants_failed: - execution.failed_files = max(execution.failed_files or 0, error_rows) + execution.failed_files = max(execution.failed_files or 0, counts["errored"]) if wants_total: - execution.total_files = max(execution.total_files or 0, row_total) + execution.total_files = max(execution.total_files or 0, counts["total"]) class FileBatchCreateAPIView(APIView): diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 89d7e3b709..536fda11f2 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -401,6 +401,8 @@ def update_execution( status: ExecutionStatus | None = None, error: str | None = None, increment_attempt: bool = False, + *, + row_locked: bool = False, ) -> None: """Update the execution status and related fields. @@ -408,23 +410,54 @@ def update_execution( status (Optional[ExecutionStatus], optional): New execution status. Defaults to None. error (Optional[str], optional): Error message if any. Defaults to None. increment_attempt (bool, optional): Whether to increment attempt counter. Defaults to False. + row_locked (bool, optional): True only when THIS instance was itself + read with ``select_for_update()`` (either mode) in the current, still-open + transaction. Its ``queue_message_id`` / ``status`` / ``attempts`` + are then the persisted values under the lock, so the second + ``SELECT ... FOR NO KEY UPDATE`` round-trip is skipped. Passing it + for an instance read without the lock reintroduces the + stale-routing race described below. """ - from django.db import transaction + from django.db import connection, transaction + + if row_locked and not connection.in_atomic_block: + # Outside a transaction the caller's lock was already released when its + # statement autocommitted, so the instance proves nothing. + raise RuntimeError( + "update_execution(row_locked=True) requires the caller's open " + "transaction that holds the row lock" + ) should_release_rate_limit = False with transaction.atomic(): # Lock the row and read the PERSISTED marker + status together, so routing # AND the write are atomic. Never trust the (possibly stale) in-memory - # self.queue_message_id: a PG row snapshotted before dispatch recorded its + # self.queue_message_id — unless the caller passes row_locked=True, i.e. + # this instance was itself read under the lock (see the docstring): a PG row snapshotted before dispatch recorded its # handle, mis-routed to an unlocked legacy full save(), would write NULL # back into queue_message_id and permanently disable every guard. Under # this lock a concurrent _record_dispatch_handle serializes after us. - locked = ( - WorkflowExecution.objects.select_for_update() - .filter(pk=self.pk) - .values("queue_message_id", "status", "attempts") - .first() - ) + # + # no_key=True takes FOR NO KEY UPDATE rather than FOR UPDATE. Status + # writers still exclude each other (the mode conflicts with itself), but + # it no longer conflicts with the FOR KEY SHARE that every child-row + # INSERT takes on this row for its foreign-key check — so file-execution + # inserts and log writes for this execution stop queueing behind a + # status update. Nothing here changes the primary key, which is the only + # thing FOR KEY SHARE protects. + if row_locked: + locked = { + "queue_message_id": self.queue_message_id, + "status": self.status, + "attempts": self.attempts, + } + else: + locked = ( + WorkflowExecution.objects.select_for_update(no_key=True) + .filter(pk=self.pk) + .values("queue_message_id", "status", "attempts") + .first() + ) if locked is None: return if locked["queue_message_id"] is None: