Fix LWT routing: preserve Paxos leader order in TokenAwarePolicy

[mykaul]

Summary

LWT (Lightweight Transaction) queries rely on Paxos consensus, where the first natural replica in the token ring acts as the Paxos leader. Routing LWT queries directly to the Paxos leader avoids an extra network hop and reduces Paxos round-trips from 4 to 3, significantly improving latency.

TokenAwarePolicy.make_query_plan() currently passes all replicas through yield_in_order(), which re-sorts them by distance (LOCAL_RACK → LOCAL → REMOTE). This is correct for regular queries, but breaks Paxos leader routing for LWT queries — the leader may be demoted if it's in a different rack than the client (with RackAwareRoundRobinPolicy).

Additionally, the tablet code path constructs replicas from the child policy's round-robin order (child.make_query_plan()), which completely loses the natural token-ring order for LWT queries regardless of child policy.

This is modeled after gocql's pickLWTReplicas() which yields replicas in natural order without distance reordering for LWT queries.

Changes

When query.is_lwt() returns True:

1. Tablet path: Resolve replicas from tablet.replicas in natural order using get_host_by_host_id() instead of filtering through the child policy's round-robin output
2. Non-tablet path: Yield replicas in their natural token-ring order (from get_replicas()), skipping only down/IGNORED hosts — do NOT pass through yield_in_order() distance bucketing
3. Non-replica fallback hosts still use distance-based ordering (unchanged)
4. Non-LWT queries: Behavior is completely unchanged

Related Issues

• FixesLWT routing: RackAwareRoundRobinPolicy demotes Paxos leader #780 (RackAwareRoundRobinPolicy demotes Paxos leader)
• FixesLWT routing: Tablet path loses natural token-ring order (Paxos leader not prioritized) #781 (Tablet path loses ring order for LWT)

Tests

Added LWTTokenAwareRoutingTest class with 11 new tests covering:

• Basic LWT preserves ring order (Paxos leader first)
• Non-LWT still uses distance-based ordering
• LWT with RackAwareRoundRobinPolicy preserves leader even when in different rack
• LWT skips down hosts
• LWT skips IGNORED hosts (remote DC)
• LWT with tablet routing preserves natural order
• Non-LWT tablet routing preserves round-robin (child policy) ordering
• LWT with shuffle disabled still preserves order
• Non-LWT with shuffle enabled randomizes replicas
• Non-LWT with DCAwareRoundRobinPolicy preserves behavior
• Non-LWT with RackAwareRoundRobinPolicy preserves rack-aware ordering

All 93 tests in tests/unit/test_policies.py pass.

Note

This fix is against master. A follow-up will apply the same logical fix on top of PR #651 (query plan optimization), which has the same bugs in its refactored code structure.

[mykaul]

CC @calebxyz

This is very surprising, I did not put enough efforts into investigating this, I should.

[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: 67df2ada-6bd6-42de-962e-e3063fa8b4f1

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.}

[coderabbitai]

🧹 Nitpick comments (2)

tests/unit/test_policies.py (2)

1369-1378: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Conditional assertion can silently pass.

The setup guarantees both rack groups are present in the fallback, so the if should be an assertion instead — otherwise a regression that drops fallback hosts entirely still passes.

♻️ Suggested change

 fallback = qplan[1:]
rack1_hosts = [h for h in fallback if h.rack == "rack1" and h.datacenter == "dc1"]
rack2_hosts = [h for h in fallback if h.rack == "rack2" and h.datacenter == "dc1"]
- if rack1_hosts and rack2_hosts:- # LOCAL_RACK hosts should appear before LOCAL hosts- first_rack1_idx = fallback.index(rack1_hosts[0])- first_rack2_idx = fallback.index(rack2_hosts[0])- assert first_rack1_idx < first_rack2_idx, (- "Fallback hosts should use distance order: LOCAL_RACK before LOCAL"- )+ assert rack1_hosts and rack2_hosts, "Expected both rack1 and rack2 fallback hosts"+ # LOCAL_RACK hosts should appear before LOCAL hosts+ assert fallback.index(rack1_hosts[0]) < fallback.index(rack2_hosts[0]), (+ "Fallback hosts should use distance order: LOCAL_RACK before LOCAL"+ )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/test_policies.py` around lines 1369 - 1378, Replace the
conditional guard around the rack-order check in the fallback section of the
policy test with an unconditional assertion that both rack1_hosts and
rack2_hosts are present, then retain the existing index-order assertion. This
ensures missing fallback groups fail the test instead of silently passing.

---

1300-1304: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assertion doesn't actually distinguish the non-LWT tablet path.

Checking that the first three hosts are the tablet replicas passes under the LWT path too, so this wouldn't catch a regression where non-LWT starts using tablet.replicas order. Assert that the tablet-natural-order lookup was never used.

♻️ Suggested strengthening

- # Non-LWT should use child policy's order (round-robin filtered by tablet replicas)- # The important thing is that it does NOT necessarily follow tablet order
replica_ids = {hosts[0].host_id, hosts[1].host_id, hosts[2].host_id}
first_three = set(qplan[:3])
- assert all(h.host_id in replica_ids for h in first_three), "First 3 should be the tablet replicas"+ assert all(h.host_id in replica_ids for h in first_three), \+ "First 3 should be the tablet replicas"+ # Non-LWT must go through the child policy plan, not tablet.replicas order+ cluster.metadata.get_host_by_host_id.assert_not_called()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/test_policies.py` around lines 1300 - 1304, The non-LWT policy
test only verifies replica membership and does not distinguish child-policy
ordering from tablet-replica ordering. Strengthen the test around the non-LWT
query-plan setup to track or inspect the tablet-natural-order lookup, and assert
that it is never used while retaining the existing replica-membership checks.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unit/test_policies.py`:
- Around line 1369-1378: Replace the conditional guard around the rack-order
check in the fallback section of the policy test with an unconditional assertion
that both rack1_hosts and rack2_hosts are present, then retain the existing
index-order assertion. This ensures missing fallback groups fail the test
instead of silently passing.
- Around line 1300-1304: The non-LWT policy test only verifies replica
membership and does not distinguish child-policy ordering from tablet-replica
ordering. Strengthen the test around the non-LWT query-plan setup to track or
inspect the tablet-natural-order lookup, and assert that it is never used while
retaining the existing replica-membership checks.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 683c9eea-1bf4-41c4-8c81-e8192a8d53e8

📥 Commits

Reviewing files that changed from the base of the PR and between bcc2d3d and 733c63c.

📒 Files selected for processing (2)

• cassandra/policies.py
• tests/unit/test_policies.py

[Copilot]

Pull request overview

Updates token-aware routing to preserve Paxos leader order for LWT queries.

Changes:

• Preserves natural replica order for vnode and tablet LWT routing.
• Retains distance-based fallback routing.
• Adds comprehensive LWT routing tests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
cassandra/policies.py	Implements LWT-specific replica ordering.
tests/unit/test_policies.py	Adds LWT routing regression tests.

---

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

[calebxyz]

CC @calebxyz

This is very surprising, I did not put enough efforts into investigating this, I should.

so leader aware paxos?

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (2)

tests/unit/test_policies.py:1312

• This test never verifies the child policy's order: converting the first three hosts to a set and checking only tablet membership also passes if the implementation switches to tablet.replicas order. Use a deterministic child plan with replica shuffling disabled and assert the exact prefix so the advertised non-LWT regression is actually covered.

 # Non-LWT should use child policy's order (round-robin filtered by tablet replicas)
# The important thing is that it does NOT necessarily follow tablet order
replica_ids = {hosts[0].host_id, hosts[1].host_id, hosts[2].host_id}
first_three = set(qplan[:3])
assert all(h.host_id in replica_ids for h in first_three), "First 3 should be the tablet replicas"

cassandra/policies.py:514

• This comment overstates the guarantee after the new REMOTE deferral below: when the first tablet replica is remote, healthy local replicas are intentionally tried first, so the first tablet replica is not selected by every client. Describe this as preserving tablet order for the locality-aware LWT path instead.

 # For LWT queries, preserve the tablet's natural replica order
# so that the first replica is tried first by every client.