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..9838c4346b --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_sync_wait_releases_db_connection.py @@ -0,0 +1,310 @@ +"""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. + +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``. +""" + +import logging +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" + + +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(db_connections): + """The single open ``default`` connection.""" + conn = db_connections.all.return_value[0] + 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 +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_if_unusable_or_obsolete, "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_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 + + _run_async(timeout=-1) + + db_connection.close_if_unusable_or_obsolete.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_if_unusable_or_obsolete.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_if_unusable_or_obsolete.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_if_unusable_or_obsolete, "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_if_unusable_or_obsolete.assert_not_called() + + +class TestReleaseDbConnection: + def test_releases_outside_a_transaction(self, db_connection): + WorkflowHelper._release_db_connection() + 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_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() + 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 + + 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): + _, 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)] + + +@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("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. Restored, so + # the shared connection is recyclable again for later tests. + saved_close_at = real_connection.close_at + real_connection.close_at = None + 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 None + + 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..2d64b2a6ff 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, 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 @@ -487,6 +487,56 @@ def _get_execution_status( ) return execution_cache.status + @staticmethod + 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 + 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. + + 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: + # 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: " + "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 + # 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.log( + level, + f"Could not release DB connection '{conn.alias}' before the " + "sync wait; continuing with it held", + exc_info=True, + ) + @staticmethod def _dispatch_orchestrator_task( *, @@ -688,7 +738,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(quiet=not first_poll) + first_poll = False time.sleep(2) timeout -= 2 @@ -989,7 +1044,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(quiet=not first_poll) + first_poll = False time.sleep(2) timeout -= 2