fix: Use asyncio.timeout in timeout_async to prevent DB timeouts (closes #16299) - #22739
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 755ae50691
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @asynccontextmanager | ||
| async def timeout_async( |
There was a problem hiding this comment.
Keep timeout_async compatible with existing with callers
After this decorator change, timeout_async() returns an _AsyncGeneratorContextManager, but existing callers still enter it with the synchronous with protocol (src/prefect/flow_engine.py:1988, src/prefect/task_engine.py:1622, src/prefect/concurrency/services.py:67). Any async flow/task or concurrency acquisition now raises TypeError before reaching the timeout logic; preserve the synchronous context-manager contract or update all callers together.
AGENTS.md reference: src/prefect/AGENTS.md:L5-L8
Useful? React with 👍 / 👎.
| ) | ||
|
|
||
| import asyncio | ||
| from contextlib import contextmanager, suppress |
There was a problem hiding this comment.
Retain sync timeout cancellation imports
This import rewrite drops cancel_sync_after and CancelledError, but the unchanged timeout() helper below still references both. Any sync flow/task that sets timeout_seconds now raises NameError when entering the timeout block instead of running with a timeout, so the sync cancellation imports need to remain while the sync helper still uses them.
Useful? React with 👍 / 👎.
| if timeout_exc_type is TimeoutError: | ||
| # Use asyncio's built-in timeout, which is robust and raises TimeoutError. | ||
| try: | ||
| async with asyncio.timeout(seconds): |
There was a problem hiding this comment.
Use a Python 3.10-compatible async timeout
asyncio.timeout is only available starting in Python 3.11, while this package still supports Python 3.10 (pyproject.toml:10). On Python 3.10, any non-None async timeout will raise AttributeError: module 'asyncio' has no attribute 'timeout' instead of enforcing the timeout, so this needs a compatibility path for 3.10.
AGENTS.md reference: AGENTS.md:L83-L86
Useful? React with 👍 / 👎.
What
Server services (e.g., FlowRunNotifications, RecentDeploymentsScheduler) are experiencing timeouts due to slow database communication. The current implementation of
timeout_asyncrelies oncancel_async_afterwhich may not properly cancel underlying asyncpg operations, causingTimeoutErrorin SQLAlchemy and service overruns.Fix
Replace the cancellation-based timeout in
timeout_asyncwithasyncio.timeout(Python 3.11+), which more reliably raisesTimeoutErrorwhen the operation exceeds the limit. The adapter supports customTimeoutErrorsubclasses by catching the built-in timeout and re-raising with the custom type. The synctimeoutfunction remains unchanged to minimize risk.Closes #16299