code: unify BytesIO allocation in encode_message compression path (no measurable end-to-end perf effect)

[mykaul]

Summary

Unify the two BytesIO allocations in encode_message into a single buff = io.BytesIO() declared once before the compression branch. In the compression path (protocol v4 and below), the body is written into buff, extracted via getvalue(), compressed, and returned via direct bytes concatenation (header + compressed_body). The non-compression path reuses the same buff with the existing seek-based header reservation pattern.

Before (master): 2 BytesIO() allocations on the compression path (buff + body), 1 on the non-compression path.
After: 1 BytesIO() allocation on both paths.

Motivation

For protocol v4 and below, compression happens at the message level (v5+ uses segment-level compression with checksumming). The original code created a separate body = io.BytesIO() for the uncompressed payload, then copied the result into a second buff = io.BytesIO() before writing the header. This second buffer is unnecessary — we can write the body into buff directly, extract it, compress, and concatenate the header as bytes.

Benchmark (updated — see note below)

Original PR description quoted a two-run, per-body-size microbenchmark showing 17-32% savings. That measurement did not hold up under closer scrutiny.

A more rigorous re-measurement (_ProtocolHandler.encode_message end-to-end, RegisterMessage + lz4 compressor, protocol v4, 10 independent fresh-process repeats per branch, each the median of 7 timeit reps over 1M iterations, run back-to-back on an otherwise idle machine) found:

	mean	median	min	max	stddev
master	1439.46ns	1432.35ns	1408.37ns	1474.33ns	22.32ns
this PR	1430.51ns	1428.91ns	1408.65ns	1475.65ns	18.39ns

Delta: ~3ns / 0.24% — not statistically distinguishable from noise (both distributions fully overlap, well inside each side's own stddev). An earlier single-process, 3-repeat pass had shown ~150ns/~10%, but that appears to have been an artifact of run ordering/warm-up rather than a real effect — it did not replicate.

Removing one BytesIO() allocation is a real, structurally sound simplification, but at this call's scale the fixed costs (lz4 compression, string-list encoding, dict overhead) dominate and the saved allocation isn't measurable in the overall encode_message cost. Landing this on code-quality/simplicity grounds rather than a performance claim.

Changes

• cassandra/protocol.py: Move buff = io.BytesIO() before the if branch. In the compression path, write body into buff instead of a separate body = io.BytesIO(), extract via getvalue(), compress, and return header + body via bytes concat. Non-compression path uses the same buff with seek(9) header reservation as before.

Testing

Unit tests pass (645 passed, 43 skipped). The non-compression path is structurally unchanged.

Isolated allocation cost (supplementary)

In a microbenchmark isolating only the header-pack + body-write logic (no compressor in the loop), removing the second BytesIO() saves ~159-185ns (~51-59%) per call. This confirms the allocation itself is meaningfully cheaper; it's just not visible at the full encode_message scale where lz4 compression dominates, per the end-to-end measurement above.

[mykaul]

Benchmark results (CPython 3.14, 500k iterations)

Body size	Original	Optimized	Δ per call
64B	352ns	180ns	-172ns
1KB	384ns	230ns	-154ns
16KB	658ns	446ns	-212ns

Applies to protocol v4 and below where compression is at the message level. Eliminates one BytesIO allocation per compressed message by using direct bytes concatenation for the header + compressed body.

[mykaul]

Follow-up commit: inline has_checksumming_support check

Commit: 04e4ba5d7 — perf: inline has_checksumming_support check to avoid classmethod call overhead

Change

Replaced ProtocolVersion.has_checksumming_support(protocol_version) calls in encode_message and decode_message with inline integer comparisons using pre-computed module-level constants (_CHECKSUMMING_MIN_VERSION, _CHECKSUMMING_MAX_VERSION). This avoids the classmethod dispatch overhead on every encode/decode call.

Semantics are identical: V5 <= version < DSE_V1.

Benchmark results (benchmarks/bench_checksumming_inline.py)

Python 3.14.3
  classmethod call: 61.3 ns
  inline compare:   25.5 ns
  saving: 35.8 ns/call (2.4x)

This is called twice per message (once on encode, once on decode), saving ~72 ns per round-trip.

Tests

607 unit tests passed, 0 failures.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mykaul]

Rebased onto origin/master

Master independently added a protocol_features parameter to encode_message/send_body (ea1ff3390, DRIVER-153/SCYLLA_USE_METADATA_ID). This PR's compression-path rewrite touched the same lines, so the rebase produced a real conflict in cassandra/protocol.py (both sides edited the if msg.custom_payload: write_bytesmap(...); msg.send_body(...); body = ...getvalue() block inside the compression branch).

Resolution: kept master's signature (protocol_features required kwarg, forwarded to send_body) while applying this PR's allocation optimization (single buff = io.BytesIO() reused across both branches, getvalue() + compress + header + body bytes-concat return instead of the old two-buffer copy). The second commit's inlined _CHECKSUMMING_MIN_VERSION/_CHECKSUMMING_MAX_VERSION check rebased cleanly on top with no further conflicts.

Byte-identity verification (lz4 + snappy)

Since this touches the compression path directly, I wrote a standalone check that loads origin/master's protocol.py (baseline: protocol_features threading, original two-BytesIO compression path) and this branch's rebased protocol.py as two independent module instances, and compares ProtocolHandler.encode_message() output byte-for-byte with real lz4 and snappy compressors:

• 192 cases: protocol v3/v4 × body sizes {0, 1, 64B, 1KB, 16KB, 100KB} × {lz4, snappy} × custom_payload on/off × tracing on/off × allow_beta on/off
• 8 additional cases using real QueryMessage/ExecuteMessage with lz4/snappy

All 200 comparisons produced byte-identical frames. The allocation optimization is confirmed to not change wire output for either supported compression codec.

Tests

• tests/unit/test_protocol.py: 25 passed
• tests/unit/test_connection.py: 32 passed
• Full tests/unit/: 720 passed, 88 skipped, 0 failed

Self-review fix

Found and fixed a stale path in the benchmark script's docstring (benchmarks/micro/bench_checksumming_inline.py — the "Run:" comment pointed at benchmarks/bench_checksumming_inline.py, missing the micro/ subdirectory). Amended into the same commit; verified the script still runs and reproduces the claimed ~2.2x / ~35ns saving.

No unresolved review threads on the PR. Force-pushed the rebased branch (2 commits, both amended in place — no new commits added) to keep this a clean rebase rather than a merge commit.

[Copilot]

Pull request overview

Note

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

This PR optimizes protocol frame encoding/decoding hot paths by reducing BytesIO allocations in the pre-v5 compression path and inlining the “checksumming supported” check to avoid classmethod overhead.

Changes:

• Reworked encode_message to use a single io.BytesIO() for both compressed and non-compressed paths, returning header + body directly for compressed frames.
• Inlined checksumming-support detection via module-level constants and integer comparisons in both encode_message and decode_message.
• Added a micro-benchmark to measure the overhead of the classmethod call vs inline comparison.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
cassandra/protocol.py	Reduces allocations in encode_message compression path; replaces classmethod checksumming test with inline range check in encode/decode.
benchmarks/micro/bench_checksumming_inline.py	Adds a micro-benchmark to quantify the classmethod vs inline checksumming check overhead.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.