Skip to content

Reassemble WebSocket frames into a single buffer - #13488

Open
bdraco wants to merge 22 commits into
masterfrom
collapse-ws-fragments
Open

Reassemble WebSocket frames into a single buffer#13488
bdraco wants to merge 22 commits into
masterfrom
collapse-ws-fragments

Conversation

@bdraco

@bdraco bdraco commented Aug 18, 2026

Copy link
Copy Markdown
Member

What do these changes do?

While reassembling a WebSocket frame that arrives across many small transport reads, the reader kept one bytes object per read with no bound on the count. The reads are now collected in a list and joined once when the frame completes; if a frame dribbles in across more than max(1024, max_msg_size // 256) reads, the pending reads are folded into a single bytearray and the list is cleared, so the retained object count stays bounded. The common few-reads path is byte-identical to before, so large messages keep their throughput; the fold only fires for pathologically small reads and stays linear.

Are there changes in behavior for the user?

No functional change; a frame split across many reads is assembled the same way, with the per-read object overhead now bounded.

Is it a substantial burden for the maintainers to support this?

No; it replaces the earlier count-cap plus pause_reading() with a fold that keeps reading.

Related issue number

Follow up to #13352

Checklist

  • I think the code is well written
  • Unit tests for the changes exist
  • Documentation reflects the changes
  • If you provide code modification, please add yourself to CONTRIBUTORS.txt
    • N/A, already listed
  • Add a new news fragment into the CHANGES/ folder
    • name it <issue_or_pr_num>.<type>.rst (e.g. 588.bugfix.rst)

Drafted with Claude Code; reviewed by bdraco.

@psf-chronographer psf-chronographer Bot added the bot:chronographer:provided There is a change note present in this PR label Aug 18, 2026
@bdraco bdraco 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 18, 2026
@codecov

codecov Bot commented Aug 18, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.00%. Comparing base (8c8906a) to head (2edae9a).
✅ All tests successful. No failed tests found.

Additional details and impacted files
@@            Coverage Diff             @@
##           master   #13488      +/-   ##
==========================================
- Coverage   99.00%   99.00%   -0.01%     
==========================================
  Files         132      132              
  Lines       49635    49724      +89     
  Branches     2575     2584       +9     
==========================================
+ Hits        49141    49229      +88     
- Misses        370      371       +1     
  Partials      124      124              
Flag Coverage Δ
Autobahn 22.11% <9.52%> (+0.07%) ⬆️
CI-GHA 98.91% <100.00%> (+<0.01%) ⬆️
OS-Linux 98.69% <100.00%> (+<0.01%) ⬆️
OS-Windows 97.27% <96.19%> (+0.26%) ⬆️
OS-macOS 98.19% <96.19%> (+0.25%) ⬆️
Py-3.10 98.13% <100.00%> (+<0.01%) ⬆️
Py-3.11 98.37% <100.00%> (+<0.01%) ⬆️
Py-3.12 98.46% <100.00%> (+<0.01%) ⬆️
Py-3.13 98.44% <100.00%> (-0.01%) ⬇️
Py-3.14 98.46% <100.00%> (+<0.01%) ⬆️
Py-3.14t 97.81% <96.19%> (+0.26%) ⬆️
Py-pypy-3.11 97.41% <100.00%> (+0.01%) ⬆️
VM-macos 98.19% <96.19%> (+0.25%) ⬆️
VM-ubuntu 98.69% <100.00%> (+<0.01%) ⬆️
VM-windows 97.27% <96.19%> (+0.26%) ⬆️
cython-coverage 82.36% <100.00%> (+0.10%) ⬆️

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.

@bdraco bdraco changed the title Collapse websocket fragments instead of pausing mid-frame Collapse WebSocket fragments when a frame arrives in many small reads Aug 18, 2026
@aiolibsbot

aiolibsbot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@codspeed-hq

codspeed-hq Bot commented Aug 18, 2026

Copy link
Copy Markdown

Merging this PR will not alter performance

✅ 84 untouched benchmarks
⏩ 83 skipped benchmarks1


Comparing collapse-ws-fragments (2edae9a) with master (8c8906a)

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.

@aiolibsbot

aiolibsbot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

aiolibsbot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@bdraco
bdraco marked this pull request as ready for review August 18, 2026 22:34
@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown

Confidence Score: 5/5

The PR appears safe to merge.

The previously reported repeated-prefix-copying issue no longer remains: folding joins only newly accumulated fragments, incrementally extends the bytearray, and detaches completed payload storage before the next frame.

Reviews (4): Last reviewed commit: "Skip the empty fragment from a header-on..." | Re-trigger Greptile

Comment thread aiohttp/_websocket/reader_py.py Outdated
@bdraco
bdraco marked this pull request as draft August 18, 2026 22:56
@bdraco bdraco changed the title Collapse WebSocket fragments when a frame arrives in many small reads Reassemble WebSocket frames into a single buffer Aug 18, 2026
@aiolibsbot

aiolibsbot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

aiolibsbot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

Copy link
Copy Markdown
Contributor

@bdraco: cythonize html looks good. only hits on the slow path, never hot path so no perf concern

Confirmed independently. Rebuilt reader_c here and checked WebSocketReader.__module__.

Hot path untouched. Single-read frame keeps had_fragments falsy. Payload stays a direct slice. Buffer only grows across multi-read frames.

314 tests pass with extension. 81 pure Python. Also probed masked, unmasked, PMCE-compressed, split CLOSE, split PING, zero-length, two-frame-torn — 1/3/7-byte reads. All payloads intact.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

Comment thread aiohttp/_websocket/reader_py.py
@aiolibsbot

Copy link
Copy Markdown
Contributor

@bdraco: still if it backs up its a problem.. maybe need to redesign again to make codspeed happy

Measured on this branch. No codspeed regression.

Cython build, best-of-N, 1 MiB payload — PR vs master: 4096-byte reads 0.27 / 0.28 ms; 64-byte reads 3.54 / 3.80 ms; 2-byte reads 100.8 / 119.4 ms. Pure Python, 2-byte reads: 316 / 400 ms. Equal or faster everywhere.

Whole-message path stays byte-identical to master — had_fragments false, no buffer touched. Every benchmark in test_benchmarks_http_websocket.py feeds complete frames, so none of them enter the new code.

On backup: mid-frame accumulation capped by max_msg_size, enforced at READ_PAYLOAD_LENGTH before any payload byte lands. Between messages, WebSocketDataQueue.feed_data still pauses on _limit. Only unbounded growth would come from removing that queue-level pause, which stays.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The suggestion stands. When max_msg_size=0, _max_fragments is 0 (line 164) and the fold condition if self._max_fragments and ... (line 506) never fires, so _payload_fragments grows one bytes per transport read with no bound. That's the intended opt-out behavior, and the THREAT_MODEL row says so — but the suite has no test that touches a reader constructed with max_msg_size=0 through the fold path.

Worth adding a small pinning test since the guard is a single conditional expression — a "tidy up" refactor could flip it silently:

def test_max_msg_size_zero_disables_fold() -> None:
    out = mock.Mock()
    out.feed_data = mock.Mock()
    parser = WebSocketReader(out, 0, compress=False, decode_text=False)
    assert parser._max_fragments == 0  # no fold when limits are opted out

If you want to also assert the runtime behavior, feed a small masked frame one byte at a time and check len(parser._payload_fragments) before the final byte — it should equal frame_payload_len - 1, never zero mid-frame.

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The root cause and fix are clear. Here's the one-line correction using option 1 (feed header + first payload byte together, which is also more realistic):

    parser.feed_data(PACK_LEN2(0x80 | WSMsgType.BINARY, 126, payload_len) + b"x")
    for _ in range(payload_len - 2):
        parser.feed_data(b"x")
    # No fold: reads accumulate in the list, the buffer is untouched.
    assert len(parser._payload_fragments) == payload_len - 1

The header-only feed_data enters READ_PAYLOAD with start_pos == data_len, so f_end_pos == f_start_pos and data_cstr[4:4] (an empty bytes) gets appended. The subsequent payload_len - 1 one-byte reads then bring the list to payload_len, not payload_len - 1.

Feeding the header and first payload byte together eliminates the empty-slice artifact and leaves the assertion correct at payload_len - 1. It also more accurately reflects how a transport delivers data (header and initial payload bytes arrive together).

@aiolibsbot

aiolibsbot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Previous review — superseded by a newer review below.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

Now I have the full picture. The a1d23084 commit ("Assert the opt-out by the unbounded list, not an exact count") was pushed after the review ran against 7ac5c64ba, relaxing the assertion from == payload_len - 1 to > 1024.

Root cause of the off-by-one: when the header is fed alone via PACK_LEN2(...) with no payload bytes, _feed_data enters READ_PAYLOAD with chunk_len = 0, sets f_start_pos = f_end_pos = data_len, and unconditionally appends data_cstr[4:4] — an empty b"" — to _payload_fragments. The fragment list therefore starts with one entry before the payload loop.

At 7ac5c64ba: assertion was == payload_len - 1 (4095). Actual count was payload_len (4096) → CI red.

After a1d23084: assertion relaxed to > 1024. 4096 satisfies this → CI green. If the review summary was generated mid-push before a1d23084 landed, the "CI is red" note is stale.


That said, the cleaner fix is one line in reader_py.py rather than in the test — guard the append against zero-length slices:

if self._payload_bytes_to_read != 0:
    if f_start_pos < f_end_pos:
        self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])

This eliminates the spurious empty fragment on every header-only read, avoids a needless list entry on the hot path regardless of _max_fragments, lets test_incomplete_frame_retained_objects_stay_bounded's fold timing match its stated intent exactly, and allows the test assertion to be restored to the precise == payload_len - 1 rather than the weaker > 1024.

The a1d23084 test fix is sufficient to unblock CI. The production guard is optional polish, but it's the better fix at the source.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The root cause Kōan identifies is accurate: when _feed_data receives a buffer that is exactly consumed by header parsing, it still enters READ_PAYLOAD with chunk_len = 0, and appends data_cstr[f_start_pos:f_end_pos] — which is b"" (empty bytes) — to _payload_fragments. This means _payload_fragments has one extra empty entry after any header-only feed_data call.

For test_max_msg_size_zero_opts_out_of_the_fold: after feeding the 4-byte header, _payload_fragments = [b""]. Then payload_len - 1 single-byte reads add payload_len - 1 more entries, giving payload_len total rather than the intended payload_len - 1.

The cleaner fix is in reader_py.py — guard the append so empty slices are never stored:

                if self._payload_bytes_to_read != 0:
                    if f_start_pos < f_end_pos:
                        self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
                    if (
                        self._max_fragments
                        and len(self._payload_fragments) > self._max_fragments
                    ):
                        # Fold to bound the object count. Not a pause: nothing
                        # resumes reading until the frame is queued.
                        self._payload_buffer += b"".join(self._payload_fragments)
                        self._payload_fragments.clear()
                    break

This is one guarding line, touches only the fragment-list management, and makes the count invariant exact: len(_payload_fragments) now equals the number of non-empty payload chunks received, which is what the test's payload_len - 1 loop is counting on.

The had_fragments branch at frame completion doesn't need the same guard — it only fires when _frame_payload_len != 0, so the frame had actual payload bytes.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The root cause is correct. When feed_data(PACK_LEN2(...)) is called with a header-only buffer, _feed_data reaches READ_PAYLOAD with start_pos == data_len == 4, so chunk_len == 0. The frame isn't complete (_payload_bytes_to_read != 0), so it appends data_cstr[4:4] — an empty b"" — to _payload_fragments. That extra item is what produces payload_len fragments instead of payload_len - 1 after the loop.

The cleanest one-line fix is to not let the header feed reach READ_PAYLOAD with zero payload bytes — feed the header and first payload byte together:

    parser.feed_data(PACK_LEN2(0x80 | WSMsgType.BINARY, 126, payload_len) + b"x")
    for _ in range(payload_len - 2):
        parser.feed_data(b"x")

That keeps the loop structure identical, eliminates the phantom empty fragment, and means after payload_len - 2 additional single-byte reads exactly payload_len - 1 fragments are queued (one per actual payload read) with _payload_bytes_to_read == 1 — exactly what the assertions below are designed to verify.

The same header-only feed pattern affects test_incomplete_frame_retained_objects_stay_bounded (starts with one empty fragment before the fold loop), but the assertion there is <= _max_fragments so it still passes; the two-byte-per-read test probably warrants the same fix for consistency, though it's not CI-blocking.

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The root cause is clear. In reader_py.py, when PACK_LEN2(...) is fed alone, the parser transitions from READ_PAYLOAD_LENGTH into READ_PAYLOAD with start_pos == data_len == 4. That gives chunk_len = 0 and f_start_pos == f_end_pos, so it appends data_cstr[4:4] — an empty b"" — to _payload_fragments before breaking. After that, each feed_data(b"x") adds one more fragment, so the count after payload_len - 1 byte-reads is 1 + (payload_len - 1) = payload_len, not payload_len - 1.

The final data is still correct (b"".join([b"", b"x", ...])b"x" * payload_len), so this is purely a test-assertion mismatch.

There are two ways to fix it:

Option A — one-line test fix (what you called out):

    assert len(parser._payload_fragments) == payload_len

Option B — reader fix (skip empty slices) — slightly more invasive but fixes the root cause so the count reflects "one fragment per non-empty read":

In reader_py.py, guard the append:

if self._payload_bytes_to_read != 0:
    if f_end_pos > f_start_pos:
        self._payload_fragments.append(data_cstr[f_start_pos:f_end_pos])
        if (
            self._max_fragments
            and len(self._payload_fragments) > self._max_fragments
        ):
            self._payload_buffer += b"".join(self._payload_fragments)
            self._payload_fragments.clear()
    break

Option B also needs the same change in the Cython path to stay in sync.

Given the review calls it "test-only, one-line fix", Option A unblocks CI with minimal churn. Option B is cleaner semantically — the fragment count then means "number of non-empty reads" — but touches two files and requires a Cython rebuild.

@aiolibsbot

Copy link
Copy Markdown
Contributor

PR Review — Reassemble WebSocket frames into a single buffer

Merge-ready. Both findings from my previous pass are resolved at 008d3ad2b, and everything I could verify at this head is green.

What is specifically strong here:

  • Removing the mid-frame pause_reading() is a genuine bugfix, and I can now show it. On master I reproduced a hard stall: 4-byte frames fed as header-read + payload-read stall at frame 1025 with the transport paused and the queue empty. That read shape is exactly what writer.py:166-167 emits (two transport.write() calls) for unmasked messages over MSG_SIZE. The same 5000-frame run on this branch never pauses. The PR description undersells this as "no functional change".
  • The fold is amortized linear, not quadratic — b"".join() only ever covers the ≤ cap pending fragments and bytearray.__iadd__ extends in place. @greptile-apps' 500 MiB estimate does not apply to this shape, and it walked the claim back itself.
  • The detach (= bytearray(), not .clear()) is the correct handoff, and test_consecutive_folded_frames_do_not_bleed pins precisely the aliasing bug .clear() would introduce.
  • THREAT_MODEL 3.8 is accurate across the configuration range, including the max_msg_size > 0 precondition, the 1024 floor caveat, and the #13350#13352 recap correction.

Verification this pass (fresh worktree at 008d3ad2b): cythonized and compiled reader_c + mask, confirmed WebSocketReader.__module__ == aiohttp._websocket.reader_c, then tests/test_websocket_parser.py, test_websocket_data_queue.py, test_web_websocket.py, test_websocket_writer.py → 168 passed; test_web_websocket_functional.py + test_client_ws_functional.py → 116 passed; benchmarks 6 passed. Pure Python: 86 passed. Differential fuzz of C vs pure-Python readers (400 trials: 1/2/3/17/512-byte chunkings, masked and unmasked, max_msg_size 0 / 4 KiB / 64 KiB, 1–3 frame messages) → 0 mismatches. black and flake8 clean on both changed files; mypy is not installed here, so type-checking is unverified. CHANGES/13488.bugfix.rst resolves as a symlink to the still-unreleased 13352.bugfix.rst, per AGENTS.md. Scope is tight: 6 files, nothing unrelated.

Non-blocking:

  • 🟢 [Pre-Existing Issue] A header-boundary read appends an empty fragment that the single-slice completion path never clears — 3000 stale list slots after 3000 frames with max_msg_size=0 (bounded by the fold otherwise). One-line guard; same fix @greptile-apps proposed.
  • 🟢 Folded unmasked frames pay one extra bytes(bytearray) copy in _handle_frame that the list-join path avoids. Marginal, on an already-degraded path.


Checklist

  • Reassembly preserves payload order and integrity (masked, unmasked, PMCE, C and pure Python)
  • No aliasing between the delivered payload and the reused buffer across consecutive folded frames
  • Retained bytes bounded by max_msg_size (READ_PAYLOAD_LENGTH pre-check precedes any buffering)
  • Retained object count bounded by the fold cap
  • had_fragments / _payload_buffer invariant holds on every completion path
  • Fold path is amortized linear, not quadratic
  • Removing the mid-frame pause_reading() is safe (and fixes a reproducible stall on master)
  • Delivered msg.data type unchanged (bytearray normalized to bytes in _handle_frame)
  • Cython .pxd declares the new attribute; extension cythonized, compiled and exercised at this head
  • Multi-frame (fin=0 + continuation) messages survive the fold — covered by the new test
  • Full websocket suites pass, C and pure Python; differential fuzz clean
  • black + flake8 clean on changed files (mypy unverified: not installed)
  • THREAT_MODEL claims precise across the full configuration range
  • Changelog follows the AGENTS.md symlink convention; Limit number of queued WS fragments #13352 still unreleased
  • No scope creep beyond the stated fix
ℹ️ Triage summary

1 pre-existing finding(s) on unchanged code suppressed (freeze).

Important

The branch moved during review. This review was performed against HEAD=008d3ad, but the PR branch now points at 2edae9a. Commits pushed after the review started are not reflected below — re-run /review to cover them.


Automated review by Kōan (Claude) HEAD=008d3ad 8 min 45s

@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown

The Koan review is thorough and I agree with its assessment. A few observations from reviewing the diff directly:

The pre-existing "empty fragment" issue Koan flagged is already fixed in this PR. The guard if f_start_pos < f_end_pos: skips the header-only read append, so the 3000-stale-slots scenario doesn't apply at this head. The non-blocking note in the checklist can be resolved.

The unchecked checklist item ("Retained object count bounded by the fold cap") looks like an oversight — the fold at len(self._payload_fragments) > self._max_fragments combined with .clear() does bound the list to _max_fragments entries at any point.

Two minor notes:

  1. In the folded completion path, payload = self._payload_buffer before self._payload_buffer = bytearray() hands the old bytearray directly to the caller (no copy). This is correct and intentional — the detach pattern is exactly what test_consecutive_folded_frames_do_not_bleed pins — but worth confirming _handle_frame doesn't mutate payload after dispatching it, since the consumer holds a live reference.

  2. CHANGES/13488.bugfix.rst as a symlink to the still-unreleased 13352.bugfix.rst is correct per AGENTS.md, and the updated changelog text covers both PRs accurately.

Overall the core correctness properties are well-established: buffer detach prevents aliasing, object count stays bounded at the fold cap, the unresumable pause_reading() is removed, and the fold is linear. The test suite covers the critical paths (masked/unmasked, folded/unfolded, multi-frame, PMCE, consecutive frames, max_msg_size=0 opt-out). Merge-ready from my perspective as well.

@bdraco
bdraco marked this pull request as ready for review August 19, 2026 02:49
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.

2 participants