From ab3dff748a71663228f6c40eb50058ccdd241a5e Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 09:40:34 +0000
Subject: [PATCH 1/6] [FIX] Take one non-blocking row lock per execution status
 write

update_status locked the workflow_execution row twice with FOR UPDATE (once
in the view, again inside update_execution) and counted the file rows 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 (staging waits up to 8 min).

- Lock once: the view hands its locked read to update_execution(locked_row=),
  and the reaper's recovery does the same.
- Lock with FOR NO KEY UPDATE: status writers still exclude each other, child
  inserts no longer wait.
- Count the file rows for the counter floor in one query before the lock.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 .../tests/test_update_status_locking.py       | 405 ++++++++++++++++++
 backend/workflow_manager/internal_views.py    | 115 ++++-
 .../workflow_v2/models/execution.py           |  41 +-
 3 files changed, 533 insertions(+), 28 deletions(-)
 create mode 100644 backend/workflow_manager/execution/tests/test_update_status_locking.py

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..1fa25b64b4
--- /dev/null
+++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py
@@ -0,0 +1,405 @@
+"""Row locking on the workflow-execution status write.
+
+``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, the COUNTs before it,
+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 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: only a broken lock makes the threads wait at all, and 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_counts_run_before_the_lock_and_only_once(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_failure_writes_counters_as_given(self):
+        ex = self._exec(files=[ExecutionStatus.ERROR])
+        with patch.object(
+            WorkflowExecutionInternalViewSet,
+            "_count_file_rows",
+            staticmethod(lambda execution: None),
+        ):
+            out = _update_status(
+                ex.id,
+                {
+                    "status": ExecutionStatus.ERROR.value,
+                    "total_files": 0,
+                    "failed_files": 0,
+                },
+            )
+        assert out.data["status"] == "updated"
+        ex.refresh_from_db()
+        assert ex.total_files == 0 and ex.failed_files == 0
+
+    def test_reaper_recovery_takes_one_lock_and_no_extra_count(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
+        lock_at = sqls.index(no_key[0])
+        # Only the reaper's own aggregate counts the files under the lock.
+        counts_after_lock = [
+            s for s in sqls[lock_at:] if "COUNT(" in s and "WORKFLOW_FILE_EXECUTION" in s
+        ]
+        assert len(counts_after_lock) == 1, counts_after_lock
+        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_locked_row_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,
+                    locked_row={
+                        "queue_message_id": locked.queue_message_id,
+                        "status": locked.status,
+                        "attempts": locked.attempts,
+                    },
+                )
+        _, 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_locked_row_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="locked_row"):
+                ex.update_execution(
+                    status=ExecutionStatus.COMPLETED,
+                    locked_row={
+                        "queue_message_id": 123,
+                        "status": ExecutionStatus.EXECUTING.value,
+                        "attempts": 0,
+                    },
+                )
+        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
+
+    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, file_row_counts=None):
+            if threading.get_ident() == state["thread"]:
+                locked.set()
+                assert release.wait(WAIT_SECONDS), "holder was never released"
+            return original(execution, validated_data, file_row_counts=file_row_counts)
+
+        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_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:
+            st.join(1.0)
+            assert st.is_alive(), f"second writer did not wait for the row lock: {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)
+            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)
+        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 6d6c969022..882bfa3b20 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -543,11 +543,19 @@ def update_status(self, request, id=None):
 
             status_enum = ExecutionStatus(validated_data["status"])
             logger.info(f"Updating status for execution {id} to {status_enum}")
+            # Count the file rows BEFORE taking the row lock, so the COUNTs no
+            # longer run while every other status writer for this execution waits.
+            # The lock never froze the file rows anyway — they are a separate table —
+            # so reading them first sees the same committed state the old in-lock
+            # read did, just without holding anyone up.
+            file_row_counts = self._count_file_rows_for_floor(execution, validated_data)
+
             # Status and file-count aggregates must commit together. Otherwise a
             # failure between the two writes can leave the row at a terminal status
             # with failed_files=None, which is_failure_run() reads as a success and
             # silently bypasses notify_on_failures subscribers.
             with transaction.atomic():
+                locked_row = None
                 # 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 +580,22 @@ 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()
+                # as locked_row, 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)
+                    locked_row = {
+                        "queue_message_id": execution.queue_message_id,
+                        "status": execution.status,
+                        "attempts": execution.attempts,
+                    }
                     protected = {
                         ExecutionStatus.COMPLETED.value,
                         ExecutionStatus.STOPPED.value,
@@ -602,10 +622,13 @@ def update_status(self, request, id=None):
                     status=status_enum,
                     error=error_message,
                     increment_attempt=increment_attempt,
+                    locked_row=locked_row,
                 )
 
                 # File aggregates are not handled by update_execution; persist separately.
-                self._update_file_aggregates(execution, validated_data)
+                self._update_file_aggregates(
+                    execution, validated_data, file_row_counts=file_row_counts
+                )
 
                 # Reaper recovery: cascade the terminal status to any non-terminal
                 # file executions so the execution and its files agree (atomic with
@@ -791,7 +814,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,8 +851,19 @@ 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,
+                    locked_row={
+                        "queue_message_id": execution.queue_message_id,
+                        "status": execution.status,
+                        "attempts": execution.attempts,
+                    },
+                )
                 self._restamp_execution_time_from_files(execution)
+                # The aggregate above already counted the file rows; reuse it rather
+                # than counting them again under the lock.
                 self._update_file_aggregates(
                     execution,
                     {
@@ -835,6 +871,7 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str:
                         "successful_files": counts["completed"],
                         "failed_files": counts["errored"],
                     },
+                    file_row_counts={"total": total, "errored": counts["errored"]},
                 )
                 logger.warning(
                     "Reaper safety-net: finalized stranded PG execution %s to %s "
@@ -933,11 +970,15 @@ def _restamp_execution_time_from_files(execution) -> None:
         execution.save(update_fields=["execution_time"])
 
     @staticmethod
-    def _update_file_aggregates(execution, validated_data) -> None:
+    def _update_file_aggregates(
+        execution, validated_data, file_row_counts: "dict | None" = None
+    ) -> None:
         """Persist total / failed / successful file counts when present.
 
         Floored at what the file rows say, because the worker's aggregate only
         counts files it was dispatched and misses any rejected before that.
+        *file_row_counts* is the ``{"total", "errored"}`` the caller already read
+        (see :meth:`_count_file_rows`); without it the rows are counted here.
         """
         update_fields: list[str] = []
         for field in ("total_files", "successful_files", "failed_files"):
@@ -948,39 +989,71 @@ def _update_file_aggregates(execution, validated_data) -> None:
             return
 
         WorkflowExecutionInternalViewSet._floor_counters_at_file_rows(
-            execution, update_fields
+            execution, update_fields, file_row_counts
         )
         execution.save(update_fields=update_fields)
 
     @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.
+    def _count_file_rows_for_floor(execution, validated_data) -> "dict | None":
+        """The file-row counts :meth:`_floor_counters_at_file_rows` will need, or
+        ``None`` when the request writes no counter the floor applies to.
+        """
+        if (
+            validated_data.get("total_files") is None
+            and validated_data.get("failed_files") is None
+        ):
+            return None
+        return WorkflowExecutionInternalViewSet._count_file_rows(execution)
 
-        Only ever corrects upward, and only to a count already true, so a run
-        with nothing rejected is left alone.
+    @staticmethod
+    def _count_file_rows(execution) -> "dict | None":
+        """Count this execution's file rows, total and ERROR, in one query.
+
+        Returns ``{"total": int, "errored": int}``, or ``None`` if the count failed
+        (logged; the counters are then written as given, exactly as before).
         """
-        wants_failed = "failed_files" in update_fields
-        wants_total = "total_files" in update_fields
-        if not (wants_failed or wants_total):
-            return
+        from django.db.models import Count, Q
 
         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()
+            return WorkflowFileExecution.objects.filter(
+                workflow_execution_id=execution.pk
+            ).aggregate(
+                total=Count("id"),
+                errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
+            )
         except Exception:
             logger.exception(
                 "Could not read file rows to floor the counters for execution %s",
                 execution.id,
             )
+            return None
+
+    @staticmethod
+    def _floor_counters_at_file_rows(
+        execution, update_fields: "list[str]", file_row_counts: "dict | None" = None
+    ) -> None:
+        """Raise the counters being written to match the file rows on disk.
+
+        Only ever corrects upward, and only to a count already true, so a run
+        with nothing rejected is left alone.
+        """
+        wants_failed = "failed_files" in update_fields
+        wants_total = "total_files" in update_fields
+        if not (wants_failed or wants_total):
+            return
+
+        counts = file_row_counts
+        if counts is None:
+            counts = WorkflowExecutionInternalViewSet._count_file_rows(execution)
+        if counts is None:
             return
 
         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 1f39804783..b7fd984693 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,
+        *,
+        locked_row: dict | None = None,
     ) -> None:
         """Update the execution status and related fields.
 
@@ -408,8 +410,23 @@ 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.
+            locked_row (Optional[dict], optional): For a caller that ALREADY holds
+                this row's lock in the current transaction: the persisted
+                ``queue_message_id`` / ``status`` / ``attempts`` it read under that
+                lock. Skips the second ``SELECT ... FOR NO KEY UPDATE`` round-trip.
+                Must come from a locked read in the same, still-open transaction —
+                a value read without the lock reintroduces the stale-routing race
+                described below.
         """
-        from django.db import transaction
+        from django.db import connection, transaction
+
+        if locked_row is not None and not connection.in_atomic_block:
+            # Outside a transaction the caller's lock was already released when its
+            # statement autocommitted, so the snapshot proves nothing.
+            raise RuntimeError(
+                "update_execution(locked_row=...) requires the caller's open "
+                "transaction that holds the row lock"
+            )
 
         should_release_rate_limit = False
         with transaction.atomic():
@@ -419,12 +436,22 @@ def update_execution(
             # 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.
+            locked = locked_row
+            if locked is None:
+                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:

From 7683bfe85a10ea18f5f4367d420ee07f82ef9a93 Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 09:49:47 +0000
Subject: [PATCH 2/6] [FIX] Count file rows after the execution lock, in a
 savepoint

Counting before the lock let a status write that waited behind another
writer (e.g. the reaper cascading files to ERROR) floor its counters at a
stale count and save failed_files=0. Count after the lock again (still one
aggregate query), and run it in a savepoint so a failed count no longer
aborts the status transaction.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 .../tests/test_update_status_locking.py       | 82 ++++++++++++++++++-
 backend/workflow_manager/internal_views.py    | 51 +++++-------
 2 files changed, 97 insertions(+), 36 deletions(-)

diff --git a/backend/workflow_manager/execution/tests/test_update_status_locking.py b/backend/workflow_manager/execution/tests/test_update_status_locking.py
index 1fa25b64b4..249c780e42 100644
--- a/backend/workflow_manager/execution/tests/test_update_status_locking.py
+++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py
@@ -6,8 +6,9 @@
 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, the COUNTs before it,
-child inserts not blocked — and status writers still serialized against each other,
+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.
 """
 
@@ -102,7 +103,7 @@ def test_celery_update_status_takes_exactly_one_no_key_lock(self):
         ex.refresh_from_db()
         assert ex.status == ExecutionStatus.COMPLETED.value
 
-    def test_file_row_counts_run_before_the_lock_and_only_once(self):
+    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(
@@ -121,7 +122,7 @@ def test_file_row_counts_run_before_the_lock_and_only_once(self):
             if "COUNT(" in s and "WORKFLOW_FILE_EXECUTION" in s
         ]
         assert len(counts) == 1, sqls
-        assert counts[0] < lock_at
+        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
@@ -159,6 +160,37 @@ def test_count_failure_writes_counters_as_given(self):
         ex.refresh_from_db()
         assert ex.total_files == 0 and ex.failed_files == 0
 
+    def test_failed_count_query_does_not_roll_back_the_status_write(self):
+        """A statement error in the count aborts a Postgres transaction; the count's
+        own savepoint must contain it so the status and counters still commit.
+        """
+        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.data["status"] == "updated", out.data
+        assert len(calls) == 1  # tried once, not retried
+        ex.refresh_from_db()
+        assert ex.status == ExecutionStatus.ERROR.value
+        assert ex.total_files == 0 and ex.failed_files == 0
+
     def test_reaper_recovery_takes_one_lock_and_no_extra_count(self):
         from datetime import timedelta
 
@@ -336,6 +368,48 @@ def insert(execution_id=ex.id):
                 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:
+            st.join(1.0)
+            assert st.is_alive(), f"callback did not wait for the row lock: {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_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.
diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index 882bfa3b20..dd1144a1a7 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -543,13 +543,6 @@ def update_status(self, request, id=None):
 
             status_enum = ExecutionStatus(validated_data["status"])
             logger.info(f"Updating status for execution {id} to {status_enum}")
-            # Count the file rows BEFORE taking the row lock, so the COUNTs no
-            # longer run while every other status writer for this execution waits.
-            # The lock never froze the file rows anyway — they are a separate table —
-            # so reading them first sees the same committed state the old in-lock
-            # read did, just without holding anyone up.
-            file_row_counts = self._count_file_rows_for_floor(execution, validated_data)
-
             # Status and file-count aggregates must commit together. Otherwise a
             # failure between the two writes can leave the row at a terminal status
             # with failed_files=None, which is_failure_run() reads as a success and
@@ -625,10 +618,12 @@ def update_status(self, request, id=None):
                     locked_row=locked_row,
                 )
 
-                # File aggregates are not handled by update_execution; persist separately.
-                self._update_file_aggregates(
-                    execution, validated_data, file_row_counts=file_row_counts
-                )
+                # File aggregates are not handled by update_execution; persist
+                # separately. The floor's file-row count runs here, AFTER the lock,
+                # on purpose: a 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.
+                self._update_file_aggregates(execution, validated_data)
 
                 # Reaper recovery: cascade the terminal status to any non-terminal
                 # file executions so the execution and its files agree (atomic with
@@ -977,8 +972,9 @@ def _update_file_aggregates(
 
         Floored at what the file rows say, because the worker's aggregate only
         counts files it was dispatched and misses any rejected before that.
-        *file_row_counts* is the ``{"total", "errored"}`` the caller already read
-        (see :meth:`_count_file_rows`); without it the rows are counted here.
+        *file_row_counts* is the ``{"total", "errored"}`` a caller already read
+        under the row lock (the reaper's own aggregate); without it the rows are
+        counted here, which must also be under the lock.
         """
         update_fields: list[str] = []
         for field in ("total_files", "successful_files", "failed_files"):
@@ -993,36 +989,27 @@ def _update_file_aggregates(
         )
         execution.save(update_fields=update_fields)
 
-    @staticmethod
-    def _count_file_rows_for_floor(execution, validated_data) -> "dict | None":
-        """The file-row counts :meth:`_floor_counters_at_file_rows` will need, or
-        ``None`` when the request writes no counter the floor applies to.
-        """
-        if (
-            validated_data.get("total_files") is None
-            and validated_data.get("failed_files") is None
-        ):
-            return None
-        return WorkflowExecutionInternalViewSet._count_file_rows(execution)
-
     @staticmethod
     def _count_file_rows(execution) -> "dict | None":
         """Count this execution's file rows, total and ERROR, in one query.
 
         Returns ``{"total": int, "errored": int}``, or ``None`` if the count failed
-        (logged; the counters are then written as given, exactly as before).
+        (logged; the counters are then written as given). The query runs in its own
+        savepoint so a failed statement does not abort the caller's transaction and
+        take the status write down with it.
         """
         from django.db.models import Count, Q
 
         from workflow_manager.workflow_v2.enums import ExecutionStatus
 
         try:
-            return WorkflowFileExecution.objects.filter(
-                workflow_execution_id=execution.pk
-            ).aggregate(
-                total=Count("id"),
-                errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
-            )
+            with transaction.atomic():
+                return WorkflowFileExecution.objects.filter(
+                    workflow_execution_id=execution.pk
+                ).aggregate(
+                    total=Count("id"),
+                    errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
+                )
         except Exception:
             logger.exception(
                 "Could not read file rows to floor the counters for execution %s",

From 96a1b858eb884af218cbd3e81d0bd390229fd0d6 Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 11:05:48 +0000
Subject: [PATCH 3/6] UN-4513 [FIX] Share the status-lock snapshot and drop the
 reaper-only count plumbing

- WorkflowExecution.STATUS_LOCK_FIELDS + status_lock_snapshot() replace the
  hand-built locked_row dicts, so callers and update_execution's own locked
  read cannot drift.
- The reaper no longer passes file_row_counts; the floor recounts as on main.
- Test documents that an uncommitted file insert holds no lock on the
  execution row (deferred FK), under either lock mode.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 .../tests/test_update_status_locking.py       | 59 ++++++++++++++-----
 backend/workflow_manager/internal_views.py    | 43 ++++----------
 .../workflow_v2/models/execution.py           | 15 ++++-
 3 files changed, 72 insertions(+), 45 deletions(-)

diff --git a/backend/workflow_manager/execution/tests/test_update_status_locking.py b/backend/workflow_manager/execution/tests/test_update_status_locking.py
index 249c780e42..9bcc2b4990 100644
--- a/backend/workflow_manager/execution/tests/test_update_status_locking.py
+++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py
@@ -191,7 +191,7 @@ def broken_aggregate(*args, **kwargs):
         assert ex.status == ExecutionStatus.ERROR.value
         assert ex.total_files == 0 and ex.failed_files == 0
 
-    def test_reaper_recovery_takes_one_lock_and_no_extra_count(self):
+    def test_reaper_recovery_takes_one_no_key_lock(self):
         from datetime import timedelta
 
         from django.utils import timezone
@@ -210,12 +210,6 @@ def test_reaper_recovery_takes_one_lock_and_no_extra_count(self):
         assert out["recovered"] == 1
         sqls, no_key, full = _lock_statements(ctx.captured_queries)
         assert len(no_key) == 1 and full == [], sqls
-        lock_at = sqls.index(no_key[0])
-        # Only the reaper's own aggregate counts the files under the lock.
-        counts_after_lock = [
-            s for s in sqls[lock_at:] if "COUNT(" in s and "WORKFLOW_FILE_EXECUTION" in s
-        ]
-        assert len(counts_after_lock) == 1, counts_after_lock
         ex.refresh_from_db()
         assert ex.status == ExecutionStatus.ERROR.value
         assert ex.total_files == 2 and ex.failed_files == 1
@@ -235,11 +229,7 @@ def test_locked_row_skips_the_second_lock(self):
             with CaptureQueriesContext(connection) as ctx:
                 locked.update_execution(
                     status=ExecutionStatus.COMPLETED,
-                    locked_row={
-                        "queue_message_id": locked.queue_message_id,
-                        "status": locked.status,
-                        "attempts": locked.attempts,
-                    },
+                    locked_row=locked.status_lock_snapshot(),
                 )
         _, no_key, full = _lock_statements(ctx.captured_queries)
         assert no_key == [] and full == []
@@ -307,11 +297,11 @@ def _hold_lock_inside_update_status(self, execution_id, data):
         original = WorkflowExecutionInternalViewSet._update_file_aggregates
         state = {"thread": None}
 
-        def paused(execution, validated_data, file_row_counts=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, file_row_counts=file_row_counts)
+            return original(execution, validated_data)
 
         patcher = patch.object(
             WorkflowExecutionInternalViewSet,
@@ -410,6 +400,47 @@ def test_count_sees_file_errors_committed_while_waiting_for_the_lock(self):
         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.
diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index dd1144a1a7..895e6cdba6 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -584,11 +584,7 @@ def update_status(self, request, id=None):
                     execution = WorkflowExecution.objects.select_for_update(
                         no_key=True
                     ).get(pk=execution.pk)
-                    locked_row = {
-                        "queue_message_id": execution.queue_message_id,
-                        "status": execution.status,
-                        "attempts": execution.attempts,
-                    }
+                    locked_row = execution.status_lock_snapshot()
                     protected = {
                         ExecutionStatus.COMPLETED.value,
                         ExecutionStatus.STOPPED.value,
@@ -620,9 +616,13 @@ def update_status(self, request, id=None):
 
                 # File aggregates are not handled by update_execution; persist
                 # separately. The floor's file-row count runs here, AFTER the lock,
-                # on purpose: a 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.
+                # 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
@@ -849,16 +849,9 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str:
                 # Already locked above; hand the locked read over so the row is
                 # not locked a second time.
                 execution.update_execution(
-                    status=computed,
-                    locked_row={
-                        "queue_message_id": execution.queue_message_id,
-                        "status": execution.status,
-                        "attempts": execution.attempts,
-                    },
+                    status=computed, locked_row=execution.status_lock_snapshot()
                 )
                 self._restamp_execution_time_from_files(execution)
-                # The aggregate above already counted the file rows; reuse it rather
-                # than counting them again under the lock.
                 self._update_file_aggregates(
                     execution,
                     {
@@ -866,7 +859,6 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str:
                         "successful_files": counts["completed"],
                         "failed_files": counts["errored"],
                     },
-                    file_row_counts={"total": total, "errored": counts["errored"]},
                 )
                 logger.warning(
                     "Reaper safety-net: finalized stranded PG execution %s to %s "
@@ -965,16 +957,11 @@ def _restamp_execution_time_from_files(execution) -> None:
         execution.save(update_fields=["execution_time"])
 
     @staticmethod
-    def _update_file_aggregates(
-        execution, validated_data, file_row_counts: "dict | None" = None
-    ) -> None:
+    def _update_file_aggregates(execution, validated_data) -> None:
         """Persist total / failed / successful file counts when present.
 
         Floored at what the file rows say, because the worker's aggregate only
         counts files it was dispatched and misses any rejected before that.
-        *file_row_counts* is the ``{"total", "errored"}`` a caller already read
-        under the row lock (the reaper's own aggregate); without it the rows are
-        counted here, which must also be under the lock.
         """
         update_fields: list[str] = []
         for field in ("total_files", "successful_files", "failed_files"):
@@ -985,7 +972,7 @@ def _update_file_aggregates(
             return
 
         WorkflowExecutionInternalViewSet._floor_counters_at_file_rows(
-            execution, update_fields, file_row_counts
+            execution, update_fields
         )
         execution.save(update_fields=update_fields)
 
@@ -1018,9 +1005,7 @@ def _count_file_rows(execution) -> "dict | None":
             return None
 
     @staticmethod
-    def _floor_counters_at_file_rows(
-        execution, update_fields: "list[str]", file_row_counts: "dict | None" = None
-    ) -> None:
+    def _floor_counters_at_file_rows(execution, update_fields: "list[str]") -> None:
         """Raise the counters being written to match the file rows on disk.
 
         Only ever corrects upward, and only to a count already true, so a run
@@ -1031,9 +1016,7 @@ def _floor_counters_at_file_rows(
         if not (wants_failed or wants_total):
             return
 
-        counts = file_row_counts
-        if counts is None:
-            counts = WorkflowExecutionInternalViewSet._count_file_rows(execution)
+        counts = WorkflowExecutionInternalViewSet._count_file_rows(execution)
         if counts is None:
             return
 
diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py
index b7fd984693..d014d54c17 100644
--- a/backend/workflow_manager/workflow_v2/models/execution.py
+++ b/backend/workflow_manager/workflow_v2/models/execution.py
@@ -396,6 +396,19 @@ def __str__(self) -> str:
             f"error message: {self.error_message})"
         )
 
+    # The persisted fields update_execution() routes and guards on, read under the
+    # row lock. One list, so a caller's locked_row and the method's own locked read
+    # cannot drift apart.
+    STATUS_LOCK_FIELDS = ("queue_message_id", "status", "attempts")
+
+    def status_lock_snapshot(self) -> dict:
+        """The ``locked_row`` for :meth:`update_execution`, taken from this instance.
+
+        Only valid when this instance was itself read with ``select_for_update()``
+        in the still-open transaction.
+        """
+        return {field: getattr(self, field) for field in self.STATUS_LOCK_FIELDS}
+
     def update_execution(
         self,
         status: ExecutionStatus | None = None,
@@ -449,7 +462,7 @@ def update_execution(
                 locked = (
                     WorkflowExecution.objects.select_for_update(no_key=True)
                     .filter(pk=self.pk)
-                    .values("queue_message_id", "status", "attempts")
+                    .values(*self.STATUS_LOCK_FIELDS)
                     .first()
                 )
             if locked is None:

From df224ddbf82cac7a7a1a0337a73d848ed1e8f9f5 Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 11:13:26 +0000
Subject: [PATCH 4/6] UN-4513 [FIX] Fail the status write when the floor count
 fails, as on main

- Drop the savepoint around the file-row count: a failed count aborts the
  transaction and the request is retried, instead of committing counters the
  floor never checked (failed_files too low reads as success).
- Replace locked_row dicts with a row_locked flag that reads the locked
  instance's own fields, so a mismatched snapshot cannot be passed.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 .../tests/test_update_status_locking.py       | 33 ++++++-------
 backend/workflow_manager/internal_views.py    | 32 ++++++-------
 .../workflow_v2/models/execution.py           | 46 ++++++++-----------
 3 files changed, 48 insertions(+), 63 deletions(-)

diff --git a/backend/workflow_manager/execution/tests/test_update_status_locking.py b/backend/workflow_manager/execution/tests/test_update_status_locking.py
index 9bcc2b4990..25ec099738 100644
--- a/backend/workflow_manager/execution/tests/test_update_status_locking.py
+++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py
@@ -160,9 +160,10 @@ def test_count_failure_writes_counters_as_given(self):
         ex.refresh_from_db()
         assert ex.total_files == 0 and ex.failed_files == 0
 
-    def test_failed_count_query_does_not_roll_back_the_status_write(self):
-        """A statement error in the count aborts a Postgres transaction; the count's
-        own savepoint must contain it so the status and counters still commit.
+    def test_failed_count_query_fails_the_request_and_commits_nothing(self):
+        """As on main: 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 = []
@@ -185,11 +186,12 @@ def broken_aggregate(*args, **kwargs):
                     "failed_files": 0,
                 },
             )
-        assert out.data["status"] == "updated", out.data
-        assert len(calls) == 1  # tried once, not retried
+        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.ERROR.value
-        assert ex.total_files == 0 and ex.failed_files == 0
+        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
@@ -220,7 +222,7 @@ def test_update_execution_locks_once_by_itself(self):
             ex.update_execution(status=ExecutionStatus.COMPLETED)
         self._assert_one_no_key_lock(ctx.captured_queries)
 
-    def test_locked_row_skips_the_second_lock(self):
+    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(
@@ -229,27 +231,20 @@ def test_locked_row_skips_the_second_lock(self):
             with CaptureQueriesContext(connection) as ctx:
                 locked.update_execution(
                     status=ExecutionStatus.COMPLETED,
-                    locked_row=locked.status_lock_snapshot(),
+                    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_locked_row_outside_a_transaction_is_rejected(self):
+    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="locked_row"):
-                ex.update_execution(
-                    status=ExecutionStatus.COMPLETED,
-                    locked_row={
-                        "queue_message_id": 123,
-                        "status": ExecutionStatus.EXECUTING.value,
-                        "attempts": 0,
-                    },
-                )
+            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
 
diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index 895e6cdba6..7a1b2ad01f 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -548,7 +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():
-                locked_row = None
+                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
@@ -575,7 +575,7 @@ def update_status(self, request, id=None):
                 # legacy path.
                 #
                 # ONE lock per call: the row read here is handed to update_execution()
-                # as locked_row, so it does not lock the same row a second time. The
+                # 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
@@ -584,7 +584,7 @@ def update_status(self, request, id=None):
                     execution = WorkflowExecution.objects.select_for_update(
                         no_key=True
                     ).get(pk=execution.pk)
-                    locked_row = execution.status_lock_snapshot()
+                    row_locked = True
                     protected = {
                         ExecutionStatus.COMPLETED.value,
                         ExecutionStatus.STOPPED.value,
@@ -611,7 +611,7 @@ def update_status(self, request, id=None):
                     status=status_enum,
                     error=error_message,
                     increment_attempt=increment_attempt,
-                    locked_row=locked_row,
+                    row_locked=row_locked,
                 )
 
                 # File aggregates are not handled by update_execution; persist
@@ -848,9 +848,7 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str:
                     computed = ExecutionStatus.COMPLETED
                 # Already locked above; hand the locked read over so the row is
                 # not locked a second time.
-                execution.update_execution(
-                    status=computed, locked_row=execution.status_lock_snapshot()
-                )
+                execution.update_execution(status=computed, row_locked=True)
                 self._restamp_execution_time_from_files(execution)
                 self._update_file_aggregates(
                     execution,
@@ -981,22 +979,22 @@ def _count_file_rows(execution) -> "dict | None":
         """Count this execution's file rows, total and ERROR, in one query.
 
         Returns ``{"total": int, "errored": int}``, or ``None`` if the count failed
-        (logged; the counters are then written as given). The query runs in its own
-        savepoint so a failed statement does not abort the caller's transaction and
-        take the status write down with it.
+        (logged). Deliberately NOT in its own savepoint: as on main, a failed count
+        aborts the caller's transaction, so the request fails and is retried rather
+        than 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
 
         try:
-            with transaction.atomic():
-                return WorkflowFileExecution.objects.filter(
-                    workflow_execution_id=execution.pk
-                ).aggregate(
-                    total=Count("id"),
-                    errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
-                )
+            return WorkflowFileExecution.objects.filter(
+                workflow_execution_id=execution.pk
+            ).aggregate(
+                total=Count("id"),
+                errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
+            )
         except Exception:
             logger.exception(
                 "Could not read file rows to floor the counters for execution %s",
diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py
index d014d54c17..b852a0e089 100644
--- a/backend/workflow_manager/workflow_v2/models/execution.py
+++ b/backend/workflow_manager/workflow_v2/models/execution.py
@@ -396,26 +396,13 @@ def __str__(self) -> str:
             f"error message: {self.error_message})"
         )
 
-    # The persisted fields update_execution() routes and guards on, read under the
-    # row lock. One list, so a caller's locked_row and the method's own locked read
-    # cannot drift apart.
-    STATUS_LOCK_FIELDS = ("queue_message_id", "status", "attempts")
-
-    def status_lock_snapshot(self) -> dict:
-        """The ``locked_row`` for :meth:`update_execution`, taken from this instance.
-
-        Only valid when this instance was itself read with ``select_for_update()``
-        in the still-open transaction.
-        """
-        return {field: getattr(self, field) for field in self.STATUS_LOCK_FIELDS}
-
     def update_execution(
         self,
         status: ExecutionStatus | None = None,
         error: str | None = None,
         increment_attempt: bool = False,
         *,
-        locked_row: dict | None = None,
+        row_locked: bool = False,
     ) -> None:
         """Update the execution status and related fields.
 
@@ -423,21 +410,21 @@ 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.
-            locked_row (Optional[dict], optional): For a caller that ALREADY holds
-                this row's lock in the current transaction: the persisted
-                ``queue_message_id`` / ``status`` / ``attempts`` it read under that
-                lock. Skips the second ``SELECT ... FOR NO KEY UPDATE`` round-trip.
-                Must come from a locked read in the same, still-open transaction —
-                a value read without the lock reintroduces the stale-routing race
-                described below.
+            row_locked (bool, optional): True only when THIS instance was itself
+                read with ``select_for_update()`` 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 connection, transaction
 
-        if locked_row is not None and not connection.in_atomic_block:
+        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 snapshot proves nothing.
+            # statement autocommitted, so the instance proves nothing.
             raise RuntimeError(
-                "update_execution(locked_row=...) requires the caller's open "
+                "update_execution(row_locked=True) requires the caller's open "
                 "transaction that holds the row lock"
             )
 
@@ -457,12 +444,17 @@ def update_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.
-            locked = locked_row
-            if locked is None:
+            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(*self.STATUS_LOCK_FIELDS)
+                    .values("queue_message_id", "status", "attempts")
                     .first()
                 )
             if locked is None:

From 6e3bd7762416bee68fdb365038fe3cd423d37558 Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 11:15:04 +0000
Subject: [PATCH 5/6] UN-4513 [FIX] Say exactly which count failures fail the
 request

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 backend/workflow_manager/internal_views.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index 7a1b2ad01f..04baac727e 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -979,10 +979,11 @@ def _count_file_rows(execution) -> "dict | None":
         """Count this execution's file rows, total and ERROR, in one query.
 
         Returns ``{"total": int, "errored": int}``, or ``None`` if the count failed
-        (logged). Deliberately NOT in its own savepoint: as on main, a failed count
-        aborts the caller's transaction, so the request fails and is retried rather
-        than committing counters the floor never checked (an under-counted
-        failed_files would read as a success and skip failure notifications).
+        (logged; the floor is then skipped, as on main). Deliberately NOT in its own
+        savepoint: as on main, a database error here aborts the caller's
+        transaction, so the request fails and is retried rather than 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 295ee8627e6404d99db91cf6f69922b0130233ae Mon Sep 17 00:00:00 2001
From: ali <117142933+muhammad-ali-e@users.noreply.github.com>
Date: Mon, 26 Oct 2026 11:42:03 +0000
Subject: [PATCH 6/6] UN-4513 [FIX] Address review: observe the lock wait, let
 count errors fail the request

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015YNhfqAovVwDvPZDYrPyE3
---
 .../tests/test_update_status_locking.py       | 54 ++++++++++++++-----
 backend/workflow_manager/internal_views.py    | 34 ++++--------
 .../workflow_v2/models/execution.py           |  5 +-
 3 files changed, 55 insertions(+), 38 deletions(-)

diff --git a/backend/workflow_manager/execution/tests/test_update_status_locking.py b/backend/workflow_manager/execution/tests/test_update_status_locking.py
index 25ec099738..d3f33d7f51 100644
--- a/backend/workflow_manager/execution/tests/test_update_status_locking.py
+++ b/backend/workflow_manager/execution/tests/test_update_status_locking.py
@@ -1,7 +1,7 @@
 """Row locking on the workflow-execution status write.
 
-``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
+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.
@@ -13,6 +13,7 @@
 """
 
 import threading
+import time
 import uuid
 from unittest.mock import MagicMock, patch
 
@@ -27,8 +28,7 @@
 from workflow_manager.workflow_v2.models.execution import WorkflowExecution
 from workflow_manager.workflow_v2.models.workflow import Workflow
 
-# Generous: only a broken lock makes the threads wait at all, and a hang must fail
-# the test rather than the CI job.
+# Generous: a hang must fail the test rather than the CI job.
 WAIT_SECONDS = 10
 
 
@@ -141,12 +141,19 @@ def test_refused_pg_write_still_takes_one_lock(self):
         assert out.data.get("reason") == "already_final"
         self._assert_one_no_key_lock(ctx.captured_queries)
 
-    def test_count_failure_writes_counters_as_given(self):
+    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(lambda execution: None),
+            staticmethod(broken_count),
         ):
             out = _update_status(
                 ex.id,
@@ -156,12 +163,14 @@ def test_count_failure_writes_counters_as_given(self):
                     "failed_files": 0,
                 },
             )
-        assert out.data["status"] == "updated"
+        assert out.status_code == 500, out.data
         ex.refresh_from_db()
-        assert ex.total_files == 0 and ex.failed_files == 0
+        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):
-        """As on main: a statement error in the count aborts the transaction, so the
+        """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.
         """
@@ -282,6 +291,25 @@ def runner():
         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.
 
@@ -384,8 +412,7 @@ def test_count_sees_file_errors_committed_while_waiting_for_the_lock(self):
             },
         )
         try:
-            st.join(1.0)
-            assert st.is_alive(), f"callback did not wait for the row lock: {sbox}"
+            self._wait_until_a_backend_waits_on_a_lock(st, sbox)
         finally:
             release.set()
             t.join(WAIT_SECONDS)
@@ -448,8 +475,7 @@ def test_status_writers_still_exclude_each_other(self):
             _update_status, ex.id, {"status": ExecutionStatus.EXECUTING.value}
         )
         try:
-            st.join(1.0)
-            assert st.is_alive(), f"second writer did not wait for the row lock: {sbox}"
+            self._wait_until_a_backend_waits_on_a_lock(st, sbox)
         finally:
             release.set()
             t.join(WAIT_SECONDS)
@@ -476,6 +502,7 @@ def write(status, start=start, execution_id=ex.id):
             ]
             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()
@@ -499,6 +526,7 @@ def bump():
         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()
diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index 04baac727e..988b797d73 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -975,33 +975,24 @@ def _update_file_aggregates(execution, validated_data) -> None:
         execution.save(update_fields=update_fields)
 
     @staticmethod
-    def _count_file_rows(execution) -> "dict | None":
+    def _count_file_rows(execution) -> dict:
         """Count this execution's file rows, total and ERROR, in one query.
 
-        Returns ``{"total": int, "errored": int}``, or ``None`` if the count failed
-        (logged; the floor is then skipped, as on main). Deliberately NOT in its own
-        savepoint: as on main, a database error here aborts the caller's
-        transaction, so the request fails and is retried rather than committing
-        counters the floor never checked (an under-counted failed_files would read
-        as a success and skip failure notifications).
+        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
 
-        try:
-            return WorkflowFileExecution.objects.filter(
-                workflow_execution_id=execution.pk
-            ).aggregate(
-                total=Count("id"),
-                errored=Count("id", filter=Q(status=ExecutionStatus.ERROR.value)),
-            )
-        except Exception:
-            logger.exception(
-                "Could not read file rows to floor the counters for execution %s",
-                execution.id,
-            )
-            return None
+        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:
@@ -1016,9 +1007,6 @@ def _floor_counters_at_file_rows(execution, update_fields: "list[str]") -> None:
             return
 
         counts = WorkflowExecutionInternalViewSet._count_file_rows(execution)
-        if counts is None:
-            return
-
         if wants_failed:
             execution.failed_files = max(execution.failed_files or 0, counts["errored"])
         if wants_total:
diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py
index b852a0e089..ede36ea652 100644
--- a/backend/workflow_manager/workflow_v2/models/execution.py
+++ b/backend/workflow_manager/workflow_v2/models/execution.py
@@ -411,7 +411,7 @@ def update_execution(
             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()`` in the current, still-open
+                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
@@ -432,7 +432,8 @@ def update_execution(
         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.