(improvement) Eliminate extra BytesIO allocation in encode_message compression path (80-300ns savings, 17-32%)

[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 (CPython 3.14, per-call, two runs on quiet machine)

Body size	Original	Optimized	Savings per call
64B	403-416ns	305-317ns	86-111ns (21-27%)
1KB	556-561ns	422-437ns	119-138ns (21-25%)
16KB	865-1122ns	718-761ns	148-361ns (17-32%)

Applies to every compressed message on protocol v4 and below.

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.

[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

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 30d52945-aad1-4dd0-944f-3cedabc5fb9b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{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.