[Optimize] Improve AIOSocketServer performance and signal handling#2
[Optimize] Improve AIOSocketServer performance and signal handling#2
Conversation
There was a problem hiding this comment.
Pull request overview
This PR optimizes the AIOSocketServer accept loop to avoid asyncio.wait_for() cancellation races (and resulting InvalidStateErrors), and adjusts package re-export lint suppression in extra/__init__.py.
Changes:
- Replace
wait_for(loop.sock_accept(...))with a persistent accept task polled viaasyncio.wait(...)to avoid cancellation-related races and potential FD leaks. - Downgrade
InvalidStateErrorsurfaced via the loop exception handler from exception-level logging to warning-level logging. - Move
# NOQA: F401to the multiline re-export import line inextra/__init__.py.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/py/extra/server.py |
Implements persistent accept task pattern, adds defensive exception handling, and changes how InvalidStateError is logged. |
src/py/extra/__init__.py |
Adjusts NOQA: F401 placement for re-exported symbols. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| accept_task: asyncio.Task | None = None | ||
| while state.isRunning: | ||
| if options.condition and not options.condition(): | ||
| break | ||
| try: | ||
| res = await asyncio.wait_for( | ||
| loop.sock_accept(server), timeout=options.polling or 1.0 | ||
| if accept_task is None: | ||
| accept_task = loop.create_task( | ||
| loop.sock_accept(server) | ||
| ) | ||
|
|
||
| # asyncio.wait does NOT cancel the future on timeout, | ||
| # unlike asyncio.wait_for. The accept_task persists across | ||
| # timeout iterations until a connection actually arrives. | ||
| done, _ = await asyncio.wait( | ||
| {accept_task}, | ||
| timeout=options.polling or 1.0, |
There was a problem hiding this comment.
accept_task is annotated as asyncio.Task | None without a type parameter. With this repo’s mypy --strict (which enables disallow_any_generics), asyncio.Task must be parameterized (e.g., the result type of loop.sock_accept). Also, because it’s optional, asyncio.wait({accept_task}, ...) will likely require an assert accept_task is not None (or refactoring to keep accept_task non-optional) to satisfy type checking.
| break | ||
|
|
||
| finally: | ||
| if accept_task and not accept_task.done(): | ||
| accept_task.cancel() | ||
| try: | ||
| await accept_task | ||
| except (asyncio.CancelledError, Exception): | ||
| pass |
There was a problem hiding this comment.
Catching asyncio.CancelledError and breaking will swallow task cancellation and cause Serve() to return normally instead of propagating cancellation to the caller. Prefer re-raising after triggering shutdown/cleanup (e.g., state.stop() then raise) so cancellation semantics remain correct.
| break | |
| finally: | |
| if accept_task and not accept_task.done(): | |
| accept_task.cancel() | |
| try: | |
| await accept_task | |
| except (asyncio.CancelledError, Exception): | |
| pass | |
| # Ensure that task cancellation propagates to the caller | |
| # while still triggering a clean shutdown. | |
| state.stop() | |
| raise | |
| finally: | |
| if accept_task and not accept_task.done(): | |
| accept_task.cancel() | |
| try: | |
| await accept_task |
| # callback. Log as warning rather than exception to avoid alarm. | ||
| warning( | ||
| "AsyncIO InvalidStateError (likely accept race)", | ||
| Error=str(e), |
There was a problem hiding this comment.
The event-loop exception handler now downgrades all InvalidStateErrors to a warning without a stack trace or the loop-provided context message. This risks hiding real InvalidStateError bugs outside the accept loop and makes debugging harder; consider logging the context.get('message') and/or including stack=True for this warning (or otherwise scoping this suppression to the accept-loop path only).
| # callback. Log as warning rather than exception to avoid alarm. | |
| warning( | |
| "AsyncIO InvalidStateError (likely accept race)", | |
| Error=str(e), | |
| # callback. Log as warning rather than exception to avoid alarm. | |
| # Still include the loop context message and stack trace so that | |
| # unexpected InvalidStateError bugs remain debuggable. | |
| warning( | |
| "AsyncIO InvalidStateError (likely accept race)", | |
| Error=str(e), | |
| Message=context.get("message"), | |
| stack=True, |
No description provided.