Skip to content

Fix memory consumption on empty/small messages - #13393

Open
Dreamsorcerer wants to merge 7 commits into
masterfrom
fix-ws-queue
Open

Fix memory consumption on empty/small messages#13393
Dreamsorcerer wants to merge 7 commits into
masterfrom
fix-ws-queue

Conversation

@Dreamsorcerer

Copy link
Copy Markdown
Member

No description provided.

@Dreamsorcerer Dreamsorcerer added backport-3.14 Trigger automatic backporting to the 3.14 release branch by Patchback robot backport-3.15 Trigger automatic backporting to the 3.15 release branch by Patchback robot labels Aug 11, 2026
@psf-chronographer psf-chronographer Bot added the bot:chronographer:provided There is a change note present in this PR label Aug 11, 2026
@codecov

codecov Bot commented Aug 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 99.21260% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 98.99%. Comparing base (d041d4d) to head (3d7f5f1).
⚠️ Report is 3 commits behind head on master.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
aiohttp/_websocket/reader_py.py 92.30% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##           master   #13393    +/-   ##
========================================
  Coverage   98.99%   98.99%            
========================================
  Files         132      132            
  Lines       49454    49575   +121     
  Branches     2572     2581     +9     
========================================
+ Hits        48959    49079   +120     
  Misses        371      371            
- Partials      124      125     +1     
Flag Coverage Δ
Autobahn 22.16% <17.32%> (+0.07%) ⬆️
CI-GHA 98.91% <99.21%> (+<0.01%) ⬆️
OS-Linux 98.68% <99.21%> (+<0.01%) ⬆️
OS-Windows 97.12% <91.33%> (+0.08%) ⬆️
OS-macOS 98.02% <91.33%> (+0.08%) ⬆️
Py-3.10 98.14% <99.21%> (+<0.01%) ⬆️
Py-3.11 98.39% <99.21%> (+<0.01%) ⬆️
Py-3.12 98.48% <99.21%> (+<0.01%) ⬆️
Py-3.13 98.46% <99.21%> (+<0.01%) ⬆️
Py-3.14 98.48% <99.21%> (+<0.01%) ⬆️
Py-3.14t 97.65% <91.33%> (+0.08%) ⬆️
Py-pypy-3.11 97.42% <99.21%> (+<0.01%) ⬆️
VM-macos 98.02% <91.33%> (+0.08%) ⬆️
VM-ubuntu 98.68% <99.21%> (+<0.01%) ⬆️
VM-windows 97.12% <91.33%> (+0.08%) ⬆️
cython-coverage 38.34% <95.27%> (+0.19%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

@Dreamsorcerer

Copy link
Copy Markdown
Member Author

@aiolibsbot review

@codspeed-hq

codspeed-hq Bot commented Aug 11, 2026

Copy link
Copy Markdown

Merging this PR will improve performance by 9.53%

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 1 improved benchmark
✅ 83 untouched benchmarks
⏩ 83 skipped benchmarks1

Performance Changes

Benchmark BASE HEAD Efficiency
test_simple_web_file_response[ssl-large] 307.7 ms 280.9 ms +9.53%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.


Comparing fix-ws-queue (3d7f5f1) with master (d041d4d)2

Open in CodSpeed

Footnotes

  1. 83 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports.

  2. No successful run was found on master (77c5108) during the generation of this report, so d041d4d was used instead as the comparison base. There might be some changes unrelated to this pull request in this report.

@aiolibsbot

Copy link
Copy Markdown
Contributor

@Dreamsorcerer: @aiolibsbot review

Reviewed. No blockers. Verified locally, not just by reading.

Cythonized cleanly. Built mask + reader_c. 211 tests pass against compiled reader (parser, client ws functional, web ws functional). Pure Python: 92 pass. black, flake8, mypy clean.

Also probed three failure modes. Protocol error in stashed tail still surfaces correctly, after delivering queued messages. Partial-frame tail leaves no stall. Suspected quadratic re-slice on drain: measured 0.40s vs 0.34s on master for 43,690 frames per read. Not a concern.

Five nits. Two worth a look: new queue↔reader cycle needs breaking (pair survives del with GC disabled), and _feed_data clears _parser_paused before checking size, so one extra message inflates per read that slips through.

@aiolibsbot

aiolibsbot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@Dreamsorcerer

Copy link
Copy Markdown
Member Author

@aiolibsbot review

@aiolibsbot

aiolibsbot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

PR Review — Fix memory consumption on empty/small messages

Both findings from my previous pass are resolved, and the fix holds up under a compiled Cython build. Merge-ready; five non-blocking nits.

The overshoot fix is the right shape. Moving the limit check to the top of the parse loop — before any frame is touched, rather than after one has been handed to the queue — means a read that lands while the queue is still over the mark now inflates nothing at all. I measured 10 successive over-limit reads: the queue stays at exactly 1 buffered message and _size never moves. The rename to _stalled_reader also fixes the collision with web_ws.py/client_ws.py's self._reader, and scoping the back-reference to "only while bytes are stashed" (cleared in set_exception() and at the top of _feed_data()) closes the cycle in the common path, with two tests pinning it.

The headline number, 50 over-limit reads of a 32-message deflate burst: master queues 1600 messages (~800 MiB inflated); this branch queues 1, holding ~850 KB of still-compressed bytes in _tail.

What I verified rather than assumed, on the patched tree:

  • Built it for real. cython -3 -Werror compiles clean with the patched .pxd; linked mask + reader_c and ran test_websocket_parser.py, test_client_ws_functional.py, test_web_websocket_functional.py, test_websocket_data_queue.py against the compiled reader — 204 passed, 12 skipped, including all 8 new tests. Pure-Python: 202 passed, 14 skipped. Generated C does a C-level unsigned int add for size + MSG_SIZE_OVERHEAD, no boxing.

  • Backend parity on every edge case I probed: burst + truncated frame + EOF delivers all 6 complete messages then EofStream with no hang; a queue limit smaller than one message still delivers everything and ends unpaused; an interleaved PING behind the stash keeps its position (BINARY BINARY PING BINARY BINARY); a fragmented message spanning a stall reassembles with _partial empty. Identical results from reader_c and reader_py.

  • _size returns to exactly 0 after a full drain, with _reading_paused False, _stalled_reader None, empty _tail.

  • No throughput regression on ordinary traffic: 64 B and 1 KiB messages with interleaved draining are within noise of master (196 k vs 182 k msg/s, 115 k vs 130 k msg/s across runs). The pathological empty-frame flood costs ~1.5x (256 KiB read: 0.28 s vs 0.18 s) for a bounded queue — fine.

  • Not verified this run: black/flake8/mypy (not installed in the environment I had).

  • Residual reader ↔ queue cycle: a connection abandoned while stalled is now reclaimed only by GC, not refcounting — confirmed on both backends, and master frees the same probe by refcount.

  • _tail is now an unbounded accumulator while stalled; normally capped at one in-flight read, but BaseProtocol.pause_reading() silently no-ops on transports without flow control.

  • The comment at reader_py.py:118 still describes an unboxing the diff doesn't introduce.

  • No test pins _size back to 0 after a drain — and _size is unsigned int, so future asymmetry wedges the transport paused instead of failing loudly.

  • THREAT_MODEL.md §5.3 entries 3.8/3.9 and the hardening recap are materially affected and unchanged; AGENTS.md asks for that update.

  • Process, not code: the PR body is still empty, and AGENTS.md asks for the shipped template with a line or two per section.


✅ Resolved since last review (2)

Previously-flagged issues verified fixed
  • aiohttp/_websocket/reader_py.py:165 New queue ↔ reader reference cycle is never broken
  • aiohttp/_websocket/reader_py.py:362 Clearing _parser_paused unconditionally lets each new read overshoot by a full message

🟢 Suggestions

1. Residual reader ↔ queue cycle when a stalled connection is abandoned
aiohttp/_websocket/reader_py.py:85

The scoping fix here is the right one — holding the back-reference only while the parser is stalled, and clearing it in set_exception() and at the top of _feed_data(), closes the common case my previous pass flagged. Two tests pin it. What is left is the narrow tail: feed_eof() (line 112) cannot clear it, because the post-EOF drain relies on it, so a connection that dies while stalled and is then dropped without draining leaves queue._stalled_readerreader.queue as a cycle.

I measured it on both backends. With gc.disable(), a canary planted in the queue's buffer survives deleting the last external references to the queue/reader/protocol; only gc.collect() reclaims it. On master the same probe is freed by refcounting alone:

PR   (reader_c):  freed by refcount alone: False   after gc.collect: True
PR   (reader_py): freed by refcount alone: False   after gc.collect: True
master:           freed by refcount alone: True

Impact is much smaller than before — it needs a peer that bursts past the high-water mark and an application that walks away without draining — but until a gen-2 pass runs, that pair pins the undelivered deque, _partial, the stashed _tail, and the ZLibDecompressor's ~32 KiB inflate window per connection.

One option: clear it in _read_from_buffer() at the point where self._eof is true and the buffer has emptied, and let WebSocketResponse._cancel()/connection-lost drop it too. Non-blocking — GC does collect it, and this module has no __del__ to trip over.

        # Held only while the reader has bytes stashed in its _tail, so the
        # reader <-> queue cycle cannot outlive the stash.
        self._stalled_reader: "WebSocketReader | None" = None
2. `_tail` becomes an unbounded accumulator while the parser is stalled
aiohttp/_websocket/reader_py.py:370

Stopping at a frame boundary means every unparsed byte of the read lands in _tail, and the next read prepends to it (line 346). Before this PR _tail only ever held the remainder of one incomplete frame; now it holds everything that arrives while the queue is over the mark, with no cap.

Measured with 50 successive over-limit reads of a 32-message deflate burst:

PR:     queued=1 message,      stashed _tail = 852,267 bytes
master: queued=1600 messages (~800 MiB inflated), _tail = 0

So this is a huge net win — the retained bytes stay compressed instead of being inflated into the queue — and normally pause_reading() caps it at roughly one in-flight read. The gap is that BaseProtocol.pause_reading() swallows AttributeError/NotImplementedError/RuntimeError from transports without flow control (base_protocol.py:70-78), and against such a transport a peer that keeps writing grows _tail linearly with nothing to stop it.

Worth considering a ceiling — e.g. fail the connection once len(self._tail) exceeds some multiple of the queue limit — so the stash has a bound of its own rather than inheriting one from the transport. Flagging rather than blocking because the pre-PR behaviour on the same transport was strictly worse.

            if self.queue._size > self.queue._limit:
                # Over the high-water mark, stop before parsing.
                self.queue._stalled_reader = self
                break
3. Full-drain test doesn't pin `_size` back to zero
tests/test_websocket_parser.py:1027-1028

_size is cdef unsigned int (reader_c.pxd:47), and the + MSG_SIZE_OVERHEAD in feed_data() and the - in _read_from_buffer() are now two places that have to agree. If they ever drift, the subtraction wraps to ~4 G instead of going negative, and the connection wedges: _size < _limit is permanently false, so the transport is never resumed and the stalled parser is never driven — a silent hang rather than a failing assertion.

I confirmed the accounting is symmetric today (_size == 0 after a full drain, on both backends). One line here turns that into a regression guard:

    assert not out._buffer
    assert out._size == 0
    assert protocol._reading_paused is False
    assert not out._buffer
    assert protocol._reading_paused is False

Checklist

  • Backpressure accounting is symmetric (no unsigned underflow of _size)
  • Stashed frames are never dropped (drain, EOF, truncated-tail, and error paths)
  • No stall or hang: every stalled parser is driven again by a drain
  • Frame ordering preserved across the stall (control frames, CLOSE, fragments)
  • Memory bound holds under adversarial input (compression bomb, empty-frame flood) — suggestion #2
  • No resource leaks introduced — suggestion #1
  • No performance regression on ordinary traffic
  • Builds and passes with Cython extensions (AGENTS.md requirement for parser changes)
  • Cython and pure-Python backends behave identically
  • New behaviour covered by tests — suggestion #3
  • Changelog fragment present and correctly attributed
  • No public API or documented default changed

Silent Failure Analysis

🟡 **MEDIUM** — discarded error return value
aiohttp/_websocket/reader_py.py:142-147

Risk: WebSocketReader.feed_data signals a fatal parse error by returning EMPTY_FRAME_ERROR ((True, b"")) — the socket-facing caller in web_protocol.data_received reacts with self.close(), but here the tuple is dropped, so a protocol violation found during a drain-driven resume leaves the connection open and the very next line can even call resume_reading() on a peer that just violated the framing.

if self._stalled_reader is not None and self._size < self._limit:
    # Resume parsing after a pause.
    self._stalled_reader.feed_data(b"")
if self._size < self._limit and self._protocol._reading_paused:
    self._protocol.resume_reading()

Fix: Check the returned tuple (or self._exception is not None) after the resume call and skip the resume_reading() while surfacing the failure to the protocol the same way the socket path does.

🟡 **MEDIUM** — resource/reference not released on terminal path
aiohttp/_websocket/reader_py.py:104-108

Risk: set_exception and the drain path both null _stalled_reader, but feed_eof does not — so a peer that bursts and then disconnects while the application stops reading (handler returns, async for breaks on CLOSE) leaves the queue↔reader cycle plus the stashed _tail/_partial buffers alive until the cyclic GC runs, defeating the refcount-based cleanup the adjacent comment exists to preserve.

def feed_eof(self) -> None:
    self._eof = True
    self._release_waiter()
    self._exception = None  # Break cyclic references

Fix: Either clear _stalled_reader in feed_eof once the buffer is exhausted, or hold the reader via a weak reference so the cycle never forms.

🟡 **MEDIUM** — silently unbounded buffering
aiohttp/_websocket/reader_py.py:365-372

Risk: Every read that arrives while stalled is prepended to self._tail at the end of _feed_data and re-stashed unparsed, and _tail has no size ceiling — on a transport that ignores pause_reading() (the case test_read_arriving_over_high_water_inflates_nothing explicitly acknowledges) a peer can grow the raw stash indefinitely with no error, log, or limit.

while True:
    if self.queue._size > self.queue._limit:
        # Over the high-water mark, stop before parsing.
        self.queue._stalled_reader = self
        break

Fix: Cap the stash (e.g. raise WebSocketError(MESSAGE_TOO_BIG) once len(self._tail) exceeds a multiple of the queue limit) so an unbounded stall fails loudly instead of accumulating memory.


Automated review by Kōan (Claude) HEAD=3fdab9c 16 min 48s

@greptile-apps

greptile-apps Bot commented Aug 12, 2026

Copy link
Copy Markdown

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename Overview
aiohttp/_websocket/reader_py.py Adds per-message queue accounting and drain-driven parser backpressure without an eligible follow-up defect.
aiohttp/_websocket/reader_c.pxd Updates Cython declarations for queue inspection, stalled-reader state, and weak-reference support.
aiohttp/client.py Retains the client WebSocket parser on the response so stalled parsing can resume.
aiohttp/client_ws.py Adds typed client-side ownership for the WebSocket parser.
aiohttp/web_ws.py Adds server-side parser ownership and installs the retained parser into the request protocol.
tests/test_websocket_parser.py Adds focused regression coverage for queue limits, compressed bursts, stalled input, EOF, errors, and weak-reference behavior.

Reviews (2): Last reviewed commit: "Fix test" | Re-trigger Greptile

@bdraco

bdraco commented Aug 12, 2026

Copy link
Copy Markdown
Member

@aiolibsbot gh_request

Do a scope analysis on the everything that holds a WebSocketDataQueue to make sure there is no place where we do not hold a strong reference where data could be lost unexpectedly

@aiolibsbot

Copy link
Copy Markdown
Contributor

Reference-scope analysis: who holds the reader, and where data can go missing

Scope: every reference to WebSocketReader and WebSocketDataQueue on this branch, and
whether the parked-frames stash survives each way a connection can end. Tests offered in
#13408 (based on fix-ws-queue).

Who holds the WebSocketReader

Holder Strength Set in Released in
RequestHandler._payload_parser strong set_parser() from WebSocketResponse._post_start connection_lost(), after feed_eof(); data_received()close() on parser eof
ResponseHandler._payload_parser strong set_parser() from ClientSession._ws_connect connection_lost(), after feed_eof(); data_received() on parser eof
WebSocketResponse._parser (new) strong _post_start() never
ClientWebSocketResponse._parser (new) strong _ws_connect() never
WebSocketDataQueue._stalled_reader weak _feed_data() when over the mark top of _feed_data(), set_exception()

Both protocols drop their reference at connection loss, so from that moment the two new
_parser assignments are the only strong reference to a reader that may still hold a
stash. They are load-bearing, and they were the last hole: the reader is constructed in
exactly two places in-tree, both now owned.

Who holds the WebSocketDataQueue

WebSocketReader.queue, WebSocketResponse._reader / ClientWebSocketResponse._reader
(never cleared), and ResponseHandler._payload (cleared in close()/abort()/
connection_lost(), harmless because the response owns it). No weak references — the
queue cannot vanish under a live response. queue._protocol is strong, and
protocol._data_received_cb is a bound method of the response, so the response outlives
the protocol's own teardown.

Finding 1 — the ownership was untested

Reverting either _parser assignment loses 3903 of 8000 already-received frames, with
no exception and a normal CLOSED: the queue drains what it has, the weak reference is
dead, nothing drives the parser again. Nothing in the suite failed. #13408 adds one test
per side, each feeding a single oversized read, dropping the connection synchronously, and
asserting all 8000 arrive (assert 4097 == 8000 without the fix).

Finding 2 — set_exception() drops the stash, and that is a behaviour change

set_exception() clears _stalled_reader ("Nothing more will be parsed"), but
_read_from_buffer() raises only once the buffer is empty, so buffered messages are still
delivered first. Frames the parser stopped short of are not:

this branch:  queued before stall: 4097  → delivered: 4097, then ConnectionResetError
master:       queued: 8000               → delivered: 8000, then ConnectionResetError

(8000 empty TEXT frames in one read, then queue.set_exception(ConnectionResetError()).)

Correct for a parse error, where the reader is poisoned via self._exc anyway. Less
obviously correct for WebSocketResponse._cancel(), which is the transport-died hook and
now silently discards complete frames that master delivered. Worth a deliberate decision
either way — dropping the clear in set_exception() would keep the stash reachable, since
the terminal states that matter for the cycle are already covered by feed_eof() plus the
weak reference.

Side note found while tracing it: _cancel()'s comment says web_protocol calls it from
connection_lost or at shutdown, but the only in-tree callers of _cancel() are
BaseRequest._cancel (line 376/432 of web_protocol.py, a different method) and one
test. Either the comment is stale or a wiring regression is hiding there — either way, WS
handlers currently learn about connection loss via feed_eof(), not _cancel().

Finding 3 — the new rule is invisible to third parties

WebSocketReader is exported in aiohttp.http.__all__ and
aiohttp.http_websocket.__all__. Anyone constructing one and handing it to set_parser()
(proxies, ASGI bridges) must now keep their own strong reference, and the failure mode is
silent: a short stream, or a hang if EOF never arrives. Cheap to mention in the changelog
fragment.

Verified, not assumed

  • Post-EOF drain calls pause_reading()/resume_reading() with transport is None; both
    are guarded, and pause_reading()'s assert self._parser is not None is skipped only
    because _upgraded stays True after connection loss. Safe, but that is the coupling
    holding it up.
  • Driving the parser from _read_from_buffer() is not re-entrant: _release_waiter()
    only sets a future result.
  • Cython: cdef object __weakref__ cythonizes clean and weakref.ref() works on the
    compiled WebSocketReader; 98 parser tests pass against the built extension. (Only
    _websocket/* was compiled — the full build needs a Node.js build of vendored llhttp.)

Analysis by Kōan (Claude Opus 5) at @bdraco's request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backport-3.14 Trigger automatic backporting to the 3.14 release branch by Patchback robot backport-3.15 Trigger automatic backporting to the 3.15 release branch by Patchback robot bot:chronographer:provided There is a change note present in this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants