diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/batch_consumer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/batch_consumer.py index 3d1b160ea017..71c324fa576e 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/batch_consumer.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/batch_consumer.py @@ -360,6 +360,18 @@ async def _wait_or_shutdown(self, timeout: float) -> None: except TimeoutError: pass + async def _handle_poll_timeout(self, poll_start: float) -> None: + # error, not exception: the timeout is the designed recovery + # path and its traceback carries no diagnostic value. + logger.error( # noqa: TRY400 + self._event("poll_timed_out"), + timeout_seconds=self._config.poll_timeout_seconds, + consecutive_failures=self._consecutive_poll_failures + 1, + ) + self._note_poll_failure("timeout", duration=time.monotonic() - poll_start) + await self._drop_conn("_poll_conn") + await self._wait_or_shutdown(self._poll_retry_delay()) + async def run(self) -> None: self._install_signal_handlers() @@ -399,23 +411,26 @@ async def run(self) -> None: continue poll_start = time.monotonic() + poll_timeout_ctx = asyncio.timeout(self._config.poll_timeout_seconds) try: conn = await self._ensure_poll_conn() - async with asyncio.timeout(self._config.poll_timeout_seconds): + async with poll_timeout_ctx: batches = await self._fetch_batches(conn, available=available) except TimeoutError: - # error, not exception: the timeout is the designed recovery - # path and its traceback carries no diagnostic value. - logger.error( # noqa: TRY400 - self._event("poll_timed_out"), - timeout_seconds=self._config.poll_timeout_seconds, - consecutive_failures=self._consecutive_poll_failures + 1, - ) - self._note_poll_failure("timeout", duration=time.monotonic() - poll_start) - await self._drop_conn("_poll_conn") - await self._wait_or_shutdown(self._poll_retry_delay()) + await self._handle_poll_timeout(poll_start) continue except psycopg.OperationalError as e: + if poll_timeout_ctx.expired(): + # asyncio.timeout() cancels the in-flight query on expiry, but + # psycopg's async wait loop can catch that cancellation and + # re-raise it as a generic OperationalError ("consuming input + # failed: server closed the connection unexpectedly") instead + # of letting a bare TimeoutError propagate. Timeout.expired() + # reflects whether this context's own deadline fired regardless + # of the exception type, so this is still the designed timeout + # path above, not a real queue-DB outage. + await self._handle_poll_timeout(poll_start) + continue logger.exception(self._event("poll_failed_queue_db_unreachable")) capture_exception(e) self._note_poll_failure("db_unreachable", duration=time.monotonic() - poll_start) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/postgres_queue/test_consumer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/postgres_queue/test_consumer.py index 61e2411f53d6..96158cbf5acc 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/postgres_queue/test_consumer.py +++ b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_v3/postgres_queue/test_consumer.py @@ -697,6 +697,66 @@ async def fetch(*args: Any, **kwargs: Any) -> list[PendingBatch]: consumer._shutdown.set() await asyncio.wait_for(run_task, timeout=5.0) + @pytest.mark.asyncio + async def test_poll_timeout_wrapped_as_operational_error_is_not_reported(self): + # psycopg's async wait loop can catch the CancelledError that + # asyncio.timeout() raises on expiry and re-raise it as a generic + # OperationalError ("consuming input failed: server closed the + # connection unexpectedly") instead of letting TimeoutError propagate. + # This is still the designed poll-timeout path, not a queue-DB outage, + # so it must not be reported to error tracking. + config = ConsumerConfig( + database_url="postgres://unused:unused@localhost/unused", + poll_interval_seconds=0.01, + poll_timeout_seconds=0.05, + ) + consumer = BatchConsumer(config=config, process_batch=AsyncMock()) + + second_poll_started = asyncio.Event() + fetch_calls = 0 + + async def fetch_wraps_cancellation_as_operational_error(*args: Any, **kwargs: Any) -> list[PendingBatch]: + nonlocal fetch_calls + fetch_calls += 1 + if fetch_calls >= 2: + second_poll_started.set() + return [] + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + raise psycopg.OperationalError( + "consuming input failed: server closed the connection unexpectedly" + ) from None + return [] + + with ( + patch.object( + consumer, "_connect", new_callable=AsyncMock, side_effect=lambda **kwargs: _make_healthy_conn() + ), + patch.object(consumer, "_install_signal_handlers"), + patch( + "products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.postgres_queue.consumer.BatchQueue.get_stale_executing", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.postgres_queue.consumer.BatchQueue.get_unprocessed_and_lock", + side_effect=fetch_wraps_cancellation_as_operational_error, + ), + patch( + "products.warehouse_sources.backend.temporal.data_imports.pipelines.pipeline_v3.postgres_queue.consumer.BatchQueue.release_all_owned_leases", + new_callable=AsyncMock, + ), + patch(f"{batch_consumer_module.__name__}.capture_exception") as mock_capture, + ): + run_task = asyncio.create_task(consumer.run()) + # A second poll can only start if the wrapped timeout was treated as recoverable. + await asyncio.wait_for(second_poll_started.wait(), timeout=2.0) + consumer._shutdown.set() + await asyncio.wait_for(run_task, timeout=5.0) + + mock_capture.assert_not_called() + class TestPollFailureLiveness: def test_withholds_liveness_after_threshold_consecutive_failures(self):