From cdf3a1372274873d4982a7caf734996356c823e0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 09:36:48 +0000 Subject: [PATCH 1/4] UN-4177 [FIX] Release the request's DB connection before the API sync wait The API deployment sync path sleep-polls for up to the request timeout with the request's Django connection still open. Behind pgbouncer's session pool that pins one Postgres server connection per waiting request, doing nothing, which is what lets the session pool grow past what the instance can serve under load. Close the connection before each sleep in execute_workflow_async and wait_for_execution. The next ORM call reconnects lazily and custom_db re-applies search_path on connect. Skipped inside an atomic block (ATOMIC_REQUESTS), where closing would break the rest of the request. This must ship before the session pool is capped, or the cap starves the backend of connections. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01RAwdv4B9xPDFEg9c18f5Az --- .../test_sync_wait_releases_db_connection.py | 195 ++++++++++++++++++ .../workflow_v2/workflow_helper.py | 23 ++- 2 files changed, 217 insertions(+), 1 deletion(-) create mode 100644 backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py diff --git a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py new file mode 100644 index 0000000000..e32e4e289a --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py @@ -0,0 +1,195 @@ +"""UN-4177 — the API sync wait must not pin a DB connection while it sleeps. + +Django holds one connection per request, and behind pgbouncer's session pool +that is one Postgres server connection for as long as the request lives. The +sync wait in ``execute_workflow_async`` (and ``wait_for_execution``) used to +sleep-poll for up to the API timeout with that connection still open, so every +waiting API request held a slot doing nothing. + +These tests pin that the connection is closed before each sleep, that nothing +changes on the async path (``timeout=-1``), and that an open atomic block is +never closed under the caller. + +DB-free: the module's ``connection``, the model, the context and ``time`` are +mocked. +""" + +from unittest.mock import MagicMock, call, patch + +import pytest + +from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.workflow_helper import WorkflowHelper + +_MOD = "workflow_manager.workflow_v2.workflow_helper" + + +@pytest.fixture +def db_connection(): + with patch(f"{_MOD}.connection") as conn: + conn.in_atomic_block = False + yield conn + + +@pytest.fixture +def async_deps(): + """Everything ``execute_workflow_async`` touches besides the connection.""" + exec_row = MagicMock(status=ExecutionStatus.EXECUTING.value) + with ( + patch(f"{_MOD}.UserContext") as user_ctx, + patch(f"{_MOD}.StateStore") as state_store, + patch(f"{_MOD}.time") as time_mod, + patch(f"{_MOD}.WorkflowExecution") as wf_exec, + patch(f"{_MOD}.WorkflowExecutionServiceHelper"), + patch(f"{_MOD}.ResultCacheUtils"), + patch.object(WorkflowHelper, "_dispatch_orchestrator_task", return_value="1"), + patch.object(WorkflowHelper, "_record_dispatch_handle"), + patch.object(WorkflowHelper, "set_result_acknowledge"), + patch.object(WorkflowHelper, "_get_execution_status") as get_status, + ): + user_ctx.get_organization_identifier.return_value = "org1" + state_store.get.return_value = None + wf_exec.objects.get.return_value = exec_row + yield time_mod, get_status + + +def _run_async(timeout: int) -> ExecutionResponse: + return WorkflowHelper.execute_workflow_async( + workflow_id="wf-1", + execution_id="exec-1", + hash_values_of_files={}, + timeout=timeout, + ) + + +class TestExecuteWorkflowAsyncWait: + def test_connection_closed_before_every_sleep(self, db_connection, async_deps): + time_mod, get_status = async_deps + get_status.side_effect = [ + ExecutionStatus.EXECUTING, + ExecutionStatus.EXECUTING, + ExecutionStatus.COMPLETED, + ] + order = MagicMock() + order.attach_mock(db_connection.close, "close") + order.attach_mock(time_mod.sleep, "sleep") + + response = _run_async(timeout=10) + + assert response.execution_status == ExecutionStatus.COMPLETED + # Close precedes each sleep — including the ones after a status read, + # since a cache miss there reopens the connection. + assert order.mock_calls == [call.close(), call.sleep(2)] * 3 + + def test_async_path_leaves_connection_alone(self, db_connection, async_deps): + time_mod, get_status = async_deps + + _run_async(timeout=-1) + + db_connection.close.assert_not_called() + time_mod.sleep.assert_not_called() + get_status.assert_not_called() + + def test_already_completed_skips_release(self, db_connection, async_deps): + _, get_status = async_deps + with patch(f"{_MOD}.WorkflowExecution") as wf_exec: + wf_exec.objects.get.return_value = MagicMock( + status=ExecutionStatus.COMPLETED.value + ) + _run_async(timeout=10) + + db_connection.close.assert_not_called() + get_status.assert_not_called() + + def test_atomic_block_is_never_closed(self, db_connection, async_deps): + _, get_status = async_deps + db_connection.in_atomic_block = True + get_status.return_value = ExecutionStatus.COMPLETED + + response = _run_async(timeout=10) + + db_connection.close.assert_not_called() + # The wait itself still runs normally. + assert get_status.called + assert response.execution_status == ExecutionStatus.COMPLETED + + +class TestWaitForExecution: + def test_connection_closed_before_every_sleep(self, db_connection): + result = ExecutionResponse("wf-1", "exec-1", ExecutionStatus.EXECUTING.value) + with ( + patch(f"{_MOD}.time") as time_mod, + patch.object( + WorkflowHelper, + "_get_execution_status", + side_effect=[ExecutionStatus.EXECUTING, ExecutionStatus.COMPLETED], + ), + ): + order = MagicMock() + order.attach_mock(db_connection.close, "close") + order.attach_mock(time_mod.sleep, "sleep") + + WorkflowHelper.wait_for_execution(result, timeout=10) + + assert order.mock_calls == [call.close(), call.sleep(2)] * 2 + assert result.execution_status == ExecutionStatus.COMPLETED + + def test_no_timeout_leaves_connection_alone(self, db_connection): + result = ExecutionResponse("wf-1", "exec-1", ExecutionStatus.EXECUTING.value) + + WorkflowHelper.wait_for_execution(result, timeout=None) + + db_connection.close.assert_not_called() + + +class TestReleaseDbConnection: + def test_closes_outside_atomic(self, db_connection): + WorkflowHelper._release_db_connection() + db_connection.close.assert_called_once_with() + + def test_skips_inside_atomic(self, db_connection): + db_connection.in_atomic_block = True + WorkflowHelper._release_db_connection() + db_connection.close.assert_not_called() + + +@pytest.mark.django_db(transaction=True) +class TestReleaseAgainstRealConnection: + """The mocks above prove *when* we close; this proves closing is safe. + + ``transaction=True`` so the test body is not itself inside an atomic + block — otherwise the release would (correctly) do nothing. + """ + + def test_orm_reconnects_with_search_path_after_release(self): + from django.conf import settings + from django.db import connection as real_connection + + from workflow_manager.workflow_v2.models.execution import WorkflowExecution + + WorkflowExecution.objects.count() + assert real_connection.connection is not None + + WorkflowHelper._release_db_connection() + assert real_connection.connection is None + + # The post-wait reads (row refetch, acknowledge) reconnect lazily... + assert WorkflowExecution.objects.count() == 0 + assert real_connection.connection is not None + # ...and custom_db re-applies the schema on that new connection. + with real_connection.cursor() as cursor: + cursor.execute("SHOW search_path") + assert cursor.fetchone()[0] == settings.DB_SCHEMA + + def test_release_inside_atomic_keeps_transaction_usable(self): + from django.db import connection as real_connection + from django.db import transaction + + from workflow_manager.workflow_v2.models.execution import WorkflowExecution + + with transaction.atomic(): + WorkflowExecution.objects.count() + WorkflowHelper._release_db_connection() + assert real_connection.connection is not None + assert WorkflowExecution.objects.count() == 0 diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index f50ba77713..27e4af9c62 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -11,7 +11,7 @@ from celery.result import AsyncResult from configuration.enums import ConfigKey from configuration.models import Configuration -from django.db import IntegrityError +from django.db import IntegrityError, connection from pg_queue.producer import DEFAULT_PRIORITY as PG_DEFAULT_PRIORITY from pg_queue.producer import enqueue_task as pg_enqueue_task from pipeline_v2.models import Pipeline @@ -487,6 +487,23 @@ def _get_execution_status( ) return execution_cache.status + @staticmethod + def _release_db_connection() -> None: + """Hand the request's DB connection back before a long sleep-poll. + + Django holds one connection for the whole request (no CONN_MAX_AGE), + and behind pgbouncer's session pool that pins a server connection for + as long as the sync wait runs — up to the API timeout, doing nothing. + Closing it here returns the slot; the next ORM call reconnects lazily + (and ``custom_db`` re-applies ``search_path`` on connect). + + Skipped inside an atomic block (e.g. ``ATOMIC_REQUESTS``): closing + there would poison the transaction for everything after the wait. + """ + if connection.in_atomic_block: + return + connection.close() + @staticmethod def _dispatch_orchestrator_task( *, @@ -689,6 +706,9 @@ def execute_workflow_async( execution_status = workflow_execution.status if timeout > -1: while not ExecutionStatus.is_completed(execution_status) and timeout > 0: + # Before every sleep, not just the first: a cache miss in + # _get_execution_status reads the row and reopens one. + cls._release_db_connection() time.sleep(2) timeout -= 2 @@ -990,6 +1010,7 @@ def wait_for_execution( execution_id = result.execution_id if timeout > 0: while not ExecutionStatus.is_completed(execution_status) and timeout > 0: + cls._release_db_connection() time.sleep(2) timeout -= 2 From 3b0c353b5f9d9f375ff3554c071575e1b9d97c18 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 11:02:44 +0000 Subject: [PATCH 2/4] UN-4177 [FIX] Release via Django's expiry check and never abort the wait Review round 1: - Release with close_if_unusable_or_obsolete over every initialized alias instead of a hard close of `default`, so persistent connections (CONN_MAX_AGE) keep working if they are ever enabled. - Also skip a connection with autocommit off, where closing would roll back pending writes; read `.autocommit` so a closed connection is not reopened just to be closed. - Swallow and log a failing release: a raise would reach the post-dispatch handler and return EXECUTING without waiting. - Test docstring now says which tests need Postgres; compare current_schema() instead of the raw search_path string. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01RAwdv4B9xPDFEg9c18f5Az --- .../test_sync_wait_releases_db_connection.py | 85 ++++++++++++++++--- .../workflow_v2/workflow_helper.py | 50 ++++++++--- 2 files changed, 111 insertions(+), 24 deletions(-) diff --git a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py index e32e4e289a..4ca130f2e7 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py +++ b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py @@ -10,8 +10,9 @@ changes on the async path (``timeout=-1``), and that an open atomic block is never closed under the caller. -DB-free: the module's ``connection``, the model, the context and ``time`` are -mocked. +The unit tests are DB-free: the module's ``connections``, the model, the +context and ``time`` are mocked. ``TestReleaseAgainstRealConnection`` needs a +real Postgres, so the conftest marks it ``integration``. """ from unittest.mock import MagicMock, call, patch @@ -25,11 +26,23 @@ _MOD = "workflow_manager.workflow_v2.workflow_helper" +def _open_conn(alias: str = "default") -> MagicMock: + return MagicMock(alias=alias, in_atomic_block=False, autocommit=True) + + +@pytest.fixture +def db_connections(): + with patch(f"{_MOD}.connections") as conns: + conns.all.return_value = [_open_conn()] + yield conns + + @pytest.fixture -def db_connection(): - with patch(f"{_MOD}.connection") as conn: - conn.in_atomic_block = False - yield conn +def db_connection(db_connections): + """The single open ``default`` connection; ``.close`` is the release call.""" + conn = db_connections.all.return_value[0] + conn.close = conn.close_if_unusable_or_obsolete + return conn @pytest.fixture @@ -82,6 +95,16 @@ def test_connection_closed_before_every_sleep(self, db_connection, async_deps): # since a cache miss there reopens the connection. assert order.mock_calls == [call.close(), call.sleep(2)] * 3 + def test_failing_release_does_not_skip_the_wait(self, db_connection, async_deps): + _, get_status = async_deps + db_connection.close_if_unusable_or_obsolete.side_effect = RuntimeError("boom") + get_status.side_effect = [ExecutionStatus.EXECUTING, ExecutionStatus.COMPLETED] + + response = _run_async(timeout=10) + + assert get_status.call_count == 2 + assert response.execution_status == ExecutionStatus.COMPLETED + def test_async_path_leaves_connection_alone(self, db_connection, async_deps): time_mod, get_status = async_deps @@ -144,14 +167,42 @@ def test_no_timeout_leaves_connection_alone(self, db_connection): class TestReleaseDbConnection: - def test_closes_outside_atomic(self, db_connection): + def test_releases_outside_a_transaction(self, db_connection): WorkflowHelper._release_db_connection() - db_connection.close.assert_called_once_with() + db_connection.close_if_unusable_or_obsolete.assert_called_once_with() def test_skips_inside_atomic(self, db_connection): db_connection.in_atomic_block = True WorkflowHelper._release_db_connection() - db_connection.close.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() + + def test_skips_when_autocommit_is_off(self, db_connection): + # Pending writes outside an atomic block: closing would roll them back. + db_connection.autocommit = False + WorkflowHelper._release_db_connection() + db_connection.close_if_unusable_or_obsolete.assert_not_called() + + def test_skips_a_closed_connection_without_reconnecting(self, db_connection): + db_connection.connection = None + WorkflowHelper._release_db_connection() + # get_autocommit() would ensure_connection() — i.e. open one to close it. + db_connection.get_autocommit.assert_not_called() + db_connection.ensure_connection.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() + + def test_a_failing_release_does_not_raise(self, db_connection): + db_connection.close_if_unusable_or_obsolete.side_effect = RuntimeError("boom") + WorkflowHelper._release_db_connection() # must not propagate + + def test_releases_every_initialized_alias(self, db_connections): + default, other = _open_conn("default"), _open_conn("other") + db_connections.all.return_value = [default, other] + + WorkflowHelper._release_db_connection() + + db_connections.all.assert_called_once_with(initialized_only=True) + default.close_if_unusable_or_obsolete.assert_called_once_with() + other.close_if_unusable_or_obsolete.assert_called_once_with() @pytest.mark.django_db(transaction=True) @@ -179,9 +230,23 @@ def test_orm_reconnects_with_search_path_after_release(self): assert real_connection.connection is not None # ...and custom_db re-applies the schema on that new connection. with real_connection.cursor() as cursor: - cursor.execute("SHOW search_path") + cursor.execute("SELECT current_schema()") assert cursor.fetchone()[0] == settings.DB_SCHEMA + def test_unexpired_persistent_connection_is_kept(self): + from django.db import connection as real_connection + + from workflow_manager.workflow_v2.models.execution import WorkflowExecution + + WorkflowExecution.objects.count() + # What CONN_MAX_AGE=None (persistent) sets at connect time. + real_connection.close_at = None + raw = real_connection.connection + + WorkflowHelper._release_db_connection() + + assert real_connection.connection is raw + def test_release_inside_atomic_keeps_transaction_usable(self): from django.db import connection as real_connection from django.db import transaction diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index 27e4af9c62..bce89600be 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -11,7 +11,7 @@ from celery.result import AsyncResult from configuration.enums import ConfigKey from configuration.models import Configuration -from django.db import IntegrityError, connection +from django.db import IntegrityError, connections from pg_queue.producer import DEFAULT_PRIORITY as PG_DEFAULT_PRIORITY from pg_queue.producer import enqueue_task as pg_enqueue_task from pipeline_v2.models import Pipeline @@ -489,20 +489,42 @@ def _get_execution_status( @staticmethod def _release_db_connection() -> None: - """Hand the request's DB connection back before a long sleep-poll. - - Django holds one connection for the whole request (no CONN_MAX_AGE), - and behind pgbouncer's session pool that pins a server connection for - as long as the sync wait runs — up to the API timeout, doing nothing. - Closing it here returns the slot; the next ORM call reconnects lazily - (and ``custom_db`` re-applies ``search_path`` on connect). - - Skipped inside an atomic block (e.g. ``ATOMIC_REQUESTS``): closing - there would poison the transaction for everything after the wait. + """Hand the request's DB connections back before a long sleep-poll. + + Django holds a connection for the whole request, and behind pgbouncer's + session pool that pins a server connection for as long as the sync wait + runs — up to the API timeout, doing nothing. The next ORM call + reconnects lazily (and ``custom_db`` re-applies ``search_path``). + + Uses Django's own expiry check rather than a hard close, so it closes + exactly what ``close_old_connections`` would: with no ``CONN_MAX_AGE`` + (today) every connection, and with persistent connections only expired + ones. Unlike ``close_old_connections`` it never touches a connection + with work in flight — an atomic block (e.g. ``ATOMIC_REQUESTS``) or + autocommit turned off — since closing would discard that work. """ - if connection.in_atomic_block: - return - connection.close() + for conn in connections.all(initialized_only=True): + # `.autocommit`, not get_autocommit(): that one ensure_connection()s, + # i.e. would open a connection only to close it. + if conn.connection is None: + continue + if conn.in_atomic_block or not conn.autocommit: + logger.debug( + f"Keeping DB connection '{conn.alias}' through the sync wait: " + "a transaction is open" + ) + continue + # Never let the release abort the wait it is there to serve: a raise + # here would land in the post-dispatch handler and return EXECUTING + # without waiting (the UN-3602 failure mode). + try: + conn.close_if_unusable_or_obsolete() + except Exception: + logger.warning( + f"Could not release DB connection '{conn.alias}' before the " + "sync wait; continuing with it held", + exc_info=True, + ) @staticmethod def _dispatch_orchestrator_task( From 1ce8dc9a9c0589bcb32d3a1ee77ade041dd3e9d2 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 11:09:24 +0000 Subject: [PATCH 3/4] UN-4177 [FIX] Log a skipped release once per wait; keep tests isolated Review round 2: - A release skipped for an open transaction (e.g. ATOMIC_REQUESTS) or a failed release now logs at WARNING on the first poll and DEBUG after, so the fix cannot switch itself off silently, and a repeating failure does not log every 2 s. - The persistent-connection test restores close_at; replace a vacuous mock assertion with a real-connection check that a closed connection is not reopened. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01RAwdv4B9xPDFEg9c18f5Az --- .../test_sync_wait_releases_db_connection.py | 60 +++++++++++++++++-- .../workflow_v2/workflow_helper.py | 26 ++++++-- 2 files changed, 74 insertions(+), 12 deletions(-) diff --git a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py index 4ca130f2e7..9bd3a3ec31 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py +++ b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py @@ -15,6 +15,7 @@ real Postgres, so the conftest marks it ``integration``. """ +import logging from unittest.mock import MagicMock, call, patch import pytest @@ -185,15 +186,50 @@ def test_skips_when_autocommit_is_off(self, db_connection): def test_skips_a_closed_connection_without_reconnecting(self, db_connection): db_connection.connection = None WorkflowHelper._release_db_connection() - # get_autocommit() would ensure_connection() — i.e. open one to close it. - db_connection.get_autocommit.assert_not_called() - db_connection.ensure_connection.assert_not_called() db_connection.close_if_unusable_or_obsolete.assert_not_called() def test_a_failing_release_does_not_raise(self, db_connection): db_connection.close_if_unusable_or_obsolete.side_effect = RuntimeError("boom") WorkflowHelper._release_db_connection() # must not propagate + @pytest.mark.parametrize( + ("quiet", "level"), [(False, logging.WARNING), (True, logging.DEBUG)] + ) + def test_skip_is_logged_loudly_only_when_not_quiet(self, db_connection, quiet, level): + db_connection.in_atomic_block = True + with patch(f"{_MOD}.logger") as log: + WorkflowHelper._release_db_connection(quiet=quiet) + assert log.log.call_args.args[0] == level + + +class TestReleaseLogsOncePerWait: + def test_only_the_first_poll_is_loud(self, db_connection, async_deps): + _, get_status = async_deps + get_status.side_effect = [ExecutionStatus.EXECUTING] * 2 + [ + ExecutionStatus.COMPLETED + ] + with patch.object(WorkflowHelper, "_release_db_connection") as release: + _run_async(timeout=10) + assert release.call_args_list == [ + call(quiet=False), + call(quiet=True), + call(quiet=True), + ] + + def test_wait_for_execution_only_the_first_poll_is_loud(self): + result = ExecutionResponse("wf-1", "exec-1", ExecutionStatus.EXECUTING.value) + with ( + patch(f"{_MOD}.time"), + patch.object( + WorkflowHelper, + "_get_execution_status", + side_effect=[ExecutionStatus.EXECUTING, ExecutionStatus.COMPLETED], + ), + patch.object(WorkflowHelper, "_release_db_connection") as release, + ): + WorkflowHelper.wait_for_execution(result, timeout=10) + assert release.call_args_list == [call(quiet=False), call(quiet=True)] + def test_releases_every_initialized_alias(self, db_connections): default, other = _open_conn("default"), _open_conn("other") db_connections.all.return_value = [default, other] @@ -239,13 +275,25 @@ def test_unexpired_persistent_connection_is_kept(self): from workflow_manager.workflow_v2.models.execution import WorkflowExecution WorkflowExecution.objects.count() - # What CONN_MAX_AGE=None (persistent) sets at connect time. + # What CONN_MAX_AGE=None (persistent) sets at connect time. Restored, so + # the shared connection is recyclable again for later tests. + saved_close_at = real_connection.close_at real_connection.close_at = None - raw = real_connection.connection + try: + raw = real_connection.connection + WorkflowHelper._release_db_connection() + assert real_connection.connection is raw + finally: + real_connection.close_at = saved_close_at + + def test_closed_connection_is_not_reopened(self): + from django.db import connection as real_connection + + real_connection.close() WorkflowHelper._release_db_connection() - assert real_connection.connection is raw + assert real_connection.connection is None def test_release_inside_atomic_keeps_transaction_usable(self): from django.db import connection as real_connection diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index bce89600be..b0fffb24af 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -488,7 +488,7 @@ def _get_execution_status( return execution_cache.status @staticmethod - def _release_db_connection() -> None: + def _release_db_connection(*, quiet: bool = False) -> None: """Hand the request's DB connections back before a long sleep-poll. Django holds a connection for the whole request, and behind pgbouncer's @@ -502,16 +502,25 @@ def _release_db_connection() -> None: ones. Unlike ``close_old_connections`` it never touches a connection with work in flight — an atomic block (e.g. ``ATOMIC_REQUESTS``) or autocommit turned off — since closing would discard that work. + + Args: + quiet (bool): Log a skipped or failed release at DEBUG instead of + WARNING. Callers pass it after the first poll, so one wait logs + once rather than every 2 s. """ + level = logging.DEBUG if quiet else logging.WARNING for conn in connections.all(initialized_only=True): # `.autocommit`, not get_autocommit(): that one ensure_connection()s, # i.e. would open a connection only to close it. if conn.connection is None: continue if conn.in_atomic_block or not conn.autocommit: - logger.debug( + # WARNING, not DEBUG: this silently undoes the release (e.g. under + # ATOMIC_REQUESTS), and a capped session pool would then starve. + logger.log( + level, f"Keeping DB connection '{conn.alias}' through the sync wait: " - "a transaction is open" + "a transaction is open, so it stays pinned for the whole wait", ) continue # Never let the release abort the wait it is there to serve: a raise @@ -520,7 +529,8 @@ def _release_db_connection() -> None: try: conn.close_if_unusable_or_obsolete() except Exception: - logger.warning( + logger.log( + level, f"Could not release DB connection '{conn.alias}' before the " "sync wait; continuing with it held", exc_info=True, @@ -727,10 +737,12 @@ def execute_workflow_async( execution_status = workflow_execution.status if timeout > -1: + first_poll = True while not ExecutionStatus.is_completed(execution_status) and timeout > 0: # Before every sleep, not just the first: a cache miss in # _get_execution_status reads the row and reopens one. - cls._release_db_connection() + cls._release_db_connection(quiet=not first_poll) + first_poll = False time.sleep(2) timeout -= 2 @@ -1031,8 +1043,10 @@ def wait_for_execution( workflow_id = result.workflow_id execution_id = result.execution_id if timeout > 0: + first_poll = True while not ExecutionStatus.is_completed(execution_status) and timeout > 0: - cls._release_db_connection() + cls._release_db_connection(quiet=not first_poll) + first_poll = False time.sleep(2) timeout -= 2 From ce92ab58c5295ab95f12ad41d37d5ac658c8e8e1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 11:14:41 +0000 Subject: [PATCH 4/4] UN-4177 [FIX] Pin the release to the expiry check in unit tests Review round 3: the unit fixture aliased close() to the expiry check, so a regression to a hard close passed the unit tier. The fixture now asserts close() is never called. Fix a comment that contradicted the once-per-wait logging, and move the alias test into its class. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01RAwdv4B9xPDFEg9c18f5Az --- .../test_sync_wait_releases_db_connection.py | 40 ++++++++++--------- .../workflow_v2/workflow_helper.py | 5 ++- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py index 9bd3a3ec31..9838c4346b 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py +++ b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py @@ -40,10 +40,12 @@ def db_connections(): @pytest.fixture def db_connection(db_connections): - """The single open ``default`` connection; ``.close`` is the release call.""" + """The single open ``default`` connection.""" conn = db_connections.all.return_value[0] - conn.close = conn.close_if_unusable_or_obsolete - return conn + yield conn + # The release goes through Django's expiry check, never a hard close: a + # hard close would also drop persistent (CONN_MAX_AGE) connections. + conn.close.assert_not_called() @pytest.fixture @@ -86,7 +88,7 @@ def test_connection_closed_before_every_sleep(self, db_connection, async_deps): ExecutionStatus.COMPLETED, ] order = MagicMock() - order.attach_mock(db_connection.close, "close") + order.attach_mock(db_connection.close_if_unusable_or_obsolete, "close") order.attach_mock(time_mod.sleep, "sleep") response = _run_async(timeout=10) @@ -111,7 +113,7 @@ def test_async_path_leaves_connection_alone(self, db_connection, async_deps): _run_async(timeout=-1) - db_connection.close.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() time_mod.sleep.assert_not_called() get_status.assert_not_called() @@ -123,7 +125,7 @@ def test_already_completed_skips_release(self, db_connection, async_deps): ) _run_async(timeout=10) - db_connection.close.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() get_status.assert_not_called() def test_atomic_block_is_never_closed(self, db_connection, async_deps): @@ -133,7 +135,7 @@ def test_atomic_block_is_never_closed(self, db_connection, async_deps): response = _run_async(timeout=10) - db_connection.close.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() # The wait itself still runs normally. assert get_status.called assert response.execution_status == ExecutionStatus.COMPLETED @@ -151,7 +153,7 @@ def test_connection_closed_before_every_sleep(self, db_connection): ), ): order = MagicMock() - order.attach_mock(db_connection.close, "close") + order.attach_mock(db_connection.close_if_unusable_or_obsolete, "close") order.attach_mock(time_mod.sleep, "sleep") WorkflowHelper.wait_for_execution(result, timeout=10) @@ -164,7 +166,7 @@ def test_no_timeout_leaves_connection_alone(self, db_connection): WorkflowHelper.wait_for_execution(result, timeout=None) - db_connection.close.assert_not_called() + db_connection.close_if_unusable_or_obsolete.assert_not_called() class TestReleaseDbConnection: @@ -201,6 +203,16 @@ def test_skip_is_logged_loudly_only_when_not_quiet(self, db_connection, quiet, l WorkflowHelper._release_db_connection(quiet=quiet) assert log.log.call_args.args[0] == level + def test_releases_every_initialized_alias(self, db_connections): + default, other = _open_conn("default"), _open_conn("other") + db_connections.all.return_value = [default, other] + + WorkflowHelper._release_db_connection() + + db_connections.all.assert_called_once_with(initialized_only=True) + default.close_if_unusable_or_obsolete.assert_called_once_with() + other.close_if_unusable_or_obsolete.assert_called_once_with() + class TestReleaseLogsOncePerWait: def test_only_the_first_poll_is_loud(self, db_connection, async_deps): @@ -230,16 +242,6 @@ def test_wait_for_execution_only_the_first_poll_is_loud(self): WorkflowHelper.wait_for_execution(result, timeout=10) assert release.call_args_list == [call(quiet=False), call(quiet=True)] - def test_releases_every_initialized_alias(self, db_connections): - default, other = _open_conn("default"), _open_conn("other") - db_connections.all.return_value = [default, other] - - WorkflowHelper._release_db_connection() - - db_connections.all.assert_called_once_with(initialized_only=True) - default.close_if_unusable_or_obsolete.assert_called_once_with() - other.close_if_unusable_or_obsolete.assert_called_once_with() - @pytest.mark.django_db(transaction=True) class TestReleaseAgainstRealConnection: diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index b0fffb24af..2d64b2a6ff 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -515,8 +515,9 @@ def _release_db_connection(*, quiet: bool = False) -> None: if conn.connection is None: continue if conn.in_atomic_block or not conn.autocommit: - # WARNING, not DEBUG: this silently undoes the release (e.g. under - # ATOMIC_REQUESTS), and a capped session pool would then starve. + # WARNING on a wait's first poll: this undoes the release (e.g. + # under ATOMIC_REQUESTS), and a capped session pool would then + # starve. Later polls of the same wait log at DEBUG. logger.log( level, f"Keeping DB connection '{conn.alias}' through the sync wait: "