fix(async): avoid asyncio.timeout on PubSub poll path (#3748) - #4299
fix(async): avoid asyncio.timeout on PubSub poll path (#3748)#4299lllakshit wants to merge 4 commits into
Conversation
Mirror the sync PubSub read path: check can_read(timeout) before read_response so non-blocking get_message() polling does not schedule asyncio.timeout(), which can corrupt the event loop callback heap under tight PubSub loops.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6d2ef91e91
ℹ️ 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".
| if not block: | ||
| if not await conn.can_read(timeout=timeout): | ||
| return None | ||
| read_timeout = math.inf |
There was a problem hiding this comment.
When a finite get_message(timeout=...) receives only the first bytes of a RESP frame, can_read() succeeds and this changes the subsequent parse to math.inf; the parser can then wait indefinitely for the rest of a fragmented or slow response instead of returning None when the requested timeout expires. Before this change, the timeout covered the whole parse, and the sync PubSub path still passes the finite timeout through to read_response.
AGENTS.md reference: AGENTS.md:L146-L149
Useful? React with 👍 / 👎.
| # without scheduling asyncio.timeout(), which can corrupt the | ||
| # event loop callback heap under tight PubSub polling (#3748). | ||
| if not block: | ||
| if not await conn.can_read(timeout=timeout): |
There was a problem hiding this comment.
If a peer closes after subscription, can_read() can raise ConnectionError when its parser observes EOF. This await now occurs before _execute, so the configured conn.retry.call_with_retry path never reconnects and resubscribes; previously the read was entirely inside _execute, and the sync path likewise performs readiness checking inside its retried callback.
AGENTS.md reference: AGENTS.md:L146-L149
Useful? React with 👍 / 👎.
| # without scheduling asyncio.timeout(), which can corrupt the | ||
| # event loop callback heap under tight PubSub polling (#3748). | ||
| if not block: | ||
| if not await conn.can_read(timeout=timeout): |
There was a problem hiding this comment.
Custom connection_class implementations can override the previously no-argument can_read() method. PubSub.get_message(timeout=0 or positive) now passes timeout=... to that override and raises TypeError instead of polling; ConnectionPool explicitly supports custom connection classes, so this breaks existing implementations that followed the old method contract.
AGENTS.md reference: AGENTS.md:L112-L114
Useful? React with 👍 / 👎.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Move can_read inside the PubSub retry path, preserve finite read timeouts after readiness checks, and tolerate legacy no-arg can_read overrides on custom connection classes. Tighten the regression test to assert asyncio.timeout(0) is never scheduled during timeout=0 polling.
|
Thanks for the automated review feedback — pushed
Happy to adjust further if anything still looks off. |
Yield to the event loop when polling with timeout=0 and no buffered data, and parse timeout=0 reads from buffered bytes only so partial RESP frames cannot block indefinitely.
|
Thanks for the follow-up review — both issues from review #5075558105 are addressed in 1. Zero-timeout poll starves event loop
2. Zero-timeout parse can hang forever
Let me know if anything else should change. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ed058f6be2
ℹ️ 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".
| disable_decoding=disable_decoding, push_request=push_request | ||
| disable_decoding=disable_decoding, | ||
| push_request=push_request, | ||
| buffered_only=buffered_only, |
There was a problem hiding this comment.
Whenever an application provides a parser_class implemented against the prior async parser contract, this forwards buffered_only=False for every normal RESP3 read and raises TypeError before it can parse a command response. parser_class is accepted by Connection.__init__, so the new private option needs to be omitted on the default path (or legacy parser signatures supported) to avoid breaking existing custom parsers.
AGENTS.md reference: AGENTS.md:L112-L114
Useful? React with 👍 / 👎.
| try: | ||
| if read_timeout is not None: | ||
| if read_timeout == 0: | ||
| if not await self.can_read(timeout=0): |
There was a problem hiding this comment.
Although PubSub.parse_response now uses a fallback helper, the connection-level zero-timeout path still invokes self.can_read(timeout=0) directly. A custom Connection subclass that implements the previously no-argument can_read() will therefore get TypeError from public read_response(timeout=0) instead of a nonblocking result; preserve compatibility in this path as well.
AGENTS.md reference: AGENTS.md:L112-L114
Useful? React with 👍 / 👎.
| except BufferedResponseIncomplete: | ||
| if buffered_only: | ||
| return None |
There was a problem hiding this comment.
With the pure-Python RESP3 parser, a zero-timeout normal read that contains a complete push followed by an incomplete command reply invokes the push handler before raising this exception. Returning here leaves the buffer intact, so the next poll reparses the push and invokes the invalidation or maintenance handler a second time; retain consumed state or defer handler effects until the full response sequence is parsed.
AGENTS.md reference: AGENTS.md:L121-L124
Useful? React with 👍 / 👎.
| response = await self._read_response_from_parser( | ||
| disable_decoding=disable_decoding, | ||
| push_request=push_request, | ||
| buffered_only=True, |
There was a problem hiding this comment.
can_read() reports bytes already buffered in the StreamReader as readable, but this forces both parser implementations into buffered_only mode: the pure-Python parser only consults its separate _buffer, and hiredis only consults its reader, so neither consumes those stream bytes. Consequently PubSub.get_message(timeout=0) returns None forever for a message that has already arrived (and can spin because readiness stays true), regardless of whether hiredis is installed.
AGENTS.md reference: AGENTS.md:L117-L119
Useful? React with 👍 / 👎.
| stream = self._parser._stream | ||
| if stream is None: | ||
| return False | ||
| await asyncio.wait_for(stream._wait_for_data("read"), timeout=timeout) |
There was a problem hiding this comment.
On CPython 3.10–3.11, asyncio.wait_for() schedules _wait_for_data() as a separate task. If the socket feed callback runs after the initial _parser.can_read() check but before that task installs the StreamReader waiter, the bytes remain buffered without waking it; this times out as False, so a finite-timeout PubSub call misses an already-arrived response until more data is received.
AGENTS.md reference: AGENTS.md:L103-L105
Useful? React with 👍 / 👎.
| response = await self._read_response_from_parser( | ||
| disable_decoding=disable_decoding, | ||
| push_request=push_request, | ||
| buffered_only=True, | ||
| ) |
There was a problem hiding this comment.
When a zero-timeout read sees only a partial frame, the parser now returns None normally rather than raising a timeout, so execution reaches the health-check update below even though no response was consumed. Repeated PubSub polls on a stalled partial response keep moving next_health_check forward and prevent check_health() from sending a PING, delaying detection of the unhealthy connection.
AGENTS.md reference: AGENTS.md:L121-L124
Useful? React with 👍 / 👎.
| try: | ||
| return await can_read(timeout=timeout) | ||
| except TypeError: | ||
| return await can_read() |
There was a problem hiding this comment.
For a legacy custom connection whose no-argument can_read() returns immediately when no bytes are buffered, the compatibility fallback makes every positive-timeout PubSub poll return None immediately. Before this change, parse_response(block=False, timeout=N) called read_response(timeout=N) directly and allowed that call to wait; the helper should not silently discard the requested timeout when using the legacy signature.
AGENTS.md reference: AGENTS.md:L112-L114
Useful? React with 👍 / 👎.
After waiting on the stream, re-run parser.can_read() so a peer close that wakes _wait_for_data without payload is not reported as readable.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 29c6aeea05
ℹ️ 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".
| # when polling with timeout=0 and no buffered data. | ||
| await asyncio.sleep(0) | ||
| return None | ||
| read_timeout = timeout |
There was a problem hiding this comment.
When the first bytes of a partial RESP frame arrive just before can_read(timeout=N) expires, this starts read_response with a new full N-second timer, so get_message(timeout=N) can take nearly twice its documented timeout before returning. This is distinct from the previously reported unbounded parse: the fresh evidence is that the finalized read_timeout = timeout makes the second window finite but does not preserve the remaining deadline; compute and pass the remaining budget instead.
AGENTS.md reference: AGENTS.md:L121-L124
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
Reviewed by Cursor Bugbot for commit 29c6aee. Configure here.
| result = self._buffer[self._pos : end - 2] | ||
| else: | ||
| if buffered_only: | ||
| raise BufferedResponseIncomplete() |
There was a problem hiding this comment.
Zero-timeout parse skips StreamReader data
High Severity
buffered_only parsing looks only at the parser's own buffer and never consumes bytes already received into StreamReader._buffer. can_read() treats those bytes as readable, so get_message(timeout=0) (the default) returns None without delivering a message that has already arrived. Incoming pub/sub frames can sit unread indefinitely.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit 29c6aee. Configure here.
| timeout=read_timeout, | ||
| disconnect_on_error=False, | ||
| push_request=True, | ||
| ) |
There was a problem hiding this comment.
Zero-timeout poll skips event loop yield
High Severity
asyncio.sleep(0) runs only when can_read() is false. If the socket buffer has data but read_response(timeout=0) still returns None, try_read returns without yielding. A tight get_message(timeout=0) loop then monopolizes the event loop, which is the starvation case this change aimed to prevent.
Reviewed by Cursor Bugbot for commit 29c6aee. Configure here.
1 participant
Summary
Fixes #3748
Async PubSub polling was entering
asyncio.timeout()on everyget_message(timeout=0)/ small-timeout call, which can triggerRuntimeError: list changed size during iterationin the event loop's callback heap under tight polling loops.This mirrors the existing sync PubSub path:
can_read(timeout)waits for readability (or returns immediately whentimeout=0) without schedulingasyncio.timeout(0)timeout=0parses buffered bytes only so partial RESP frames cannot blockasyncio.sleep(0)when no data is bufferedTest plan
test_pubsub_timeout_zero_avoids_async_timeout_zerotest_get_message_timeout_zero_yields_event_loopbuffered_onlyunit teststest_async_connection_can_read_eof_after_wait_raisesCloses #3748