From e960106b670746b1d2173a10e6ce3a6add36948b Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Fri, 6 Mar 2026 09:28:44 -0400 Subject: [PATCH] fix: Include non-replica nodes in LWT PRESERVE_REPLICA_ORDER query plan The PreserveReplicaOrderIterator in TokenAwarePolicy previously returned only replica nodes in the query plan. When replicas were unavailable (e.g. prepared statements before parameter binding, or replicas going down after plan construction), the query plan could be empty or insufficient, causing "No node was available" errors. The fix adds a third pass that appends non-replica nodes from the child policy after all replicas have been returned. When no replicas are available at all (all DOWN/IGNORED), the full child policy fallback is preserved. Fixes #833 --- .../core/policies/TokenAwarePolicy.java | 76 ++++--- .../core/policies/TokenAwarePolicyTest.java | 191 +++++++++++++++++- 2 files changed, 226 insertions(+), 41 deletions(-) diff --git a/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java b/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java index 27022c235c8..30eb76ba49b 100644 --- a/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java +++ b/driver-core/src/main/java/com/datastax/driver/core/policies/TokenAwarePolicy.java @@ -38,8 +38,10 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; /** @@ -70,9 +72,9 @@ *

Lightweight Transaction (LWT) Routing

* *

For {@linkplain Statement#isLWT() lightweight transaction} queries, this policy provides - * specialized replica-only routing to optimize LWT performance and avoid contention. When LWT - * routing is enabled (the default), the query plan contains only replicas for the - * target partition, ordered by datacenter locality: + * specialized routing to optimize LWT performance and avoid contention. When LWT routing is enabled + * (the default), the query plan prioritizes replicas for the target partition, ordered by + * datacenter locality, followed by non-replica nodes for failover: * *

* *

Rack awareness is intentionally not applied to LWT replica ordering. @@ -243,36 +243,38 @@ protected Host computeNext() { /** * An iterator that returns replicas first, with local replicas prioritized (preserving primary - * replica order), then remote replicas. Used for LWT queries to ensure replica-only routing and - * minimize coordinator forwarding overhead. DOWN and IGNORED hosts are filtered out. + * replica order), then remote replicas, then non-replica nodes from the child policy. DOWN and + * IGNORED hosts are filtered out from replicas. * - *

Query plan follows a three-pass strategy: + *

Query plan follows a four-pass strategy: * *

    *
  1. Local replicas: Returns UP replicas marked as LOCAL by the child policy, * in the order provided by cluster metadata (preserving primary replica order). *
  2. Remote replicas: Returns UP replicas marked as REMOTE by the child * policy. - *
  3. Child policy fallback: If no suitable replicas are available (for - * example, all are DOWN or IGNORED and thus none are returned), falls back to the child - * policy's query plan for the remaining hosts. The child policy's plan is used as-is and - * may include hosts that were already considered by this iterator. + *
  4. Non-replica nodes: Returns remaining nodes from the child policy's query + * plan, skipping any hosts already returned as replicas. This ensures all available nodes + * are included in the query plan for failover. + *
  5. Child policy fallback: If no suitable replicas were returned at all (for + * example, all are DOWN or IGNORED), falls back to the child policy's full query plan. *
*/ private class PreserveReplicaOrderIterator extends AbstractIterator { private final Iterator replicasIterator; + private final List replicas; private final String keyspace; private final Statement statement; private List nonLocalReplicas; private Iterator nonLocalReplicasIterator; - private boolean hasReturnedReplicas; + private Set returnedHosts; private Iterator childIterator; - public PreserveReplicaOrderIterator( - String keyspace, Statement statement, Iterator replicasIterator) { + public PreserveReplicaOrderIterator(String keyspace, Statement statement, List replicas) { this.keyspace = keyspace; this.statement = statement; - this.replicasIterator = replicasIterator; + this.replicas = replicas; + this.replicasIterator = replicas.iterator(); } @Override @@ -289,7 +291,8 @@ protected Host computeNext() { switch (distance) { case LOCAL: - hasReturnedReplicas = true; + if (returnedHosts == null) returnedHosts = new HashSet<>(); + returnedHosts.add(host); return host; case REMOTE: // Collect remote replicas for second pass @@ -307,21 +310,31 @@ protected Host computeNext() { if (nonLocalReplicasIterator == null) { nonLocalReplicasIterator = nonLocalReplicas.iterator(); } - if (nonLocalReplicasIterator.hasNext()) { - hasReturnedReplicas = true; - return nonLocalReplicasIterator.next(); + while (nonLocalReplicasIterator.hasNext()) { + Host host = nonLocalReplicasIterator.next(); + if (returnedHosts == null) returnedHosts = new HashSet<>(); + returnedHosts.add(host); + return host; } } - // Third pass: fallback to child policy if no suitable replicas were returned - // This handles cases where all replicas are empty, DOWN or IGNORED - if (!hasReturnedReplicas) { - if (childIterator == null) { - childIterator = childPolicy.newQueryPlan(keyspace, statement); + // Third pass: return remaining nodes from child policy + if (childIterator == null) { + childIterator = childPolicy.newQueryPlan(keyspace, statement); + } + while (childIterator.hasNext()) { + Host host = childIterator.next(); + // Skip hosts we already returned as replicas + if (returnedHosts != null && returnedHosts.contains(host)) { + continue; } - if (childIterator.hasNext()) { - return childIterator.next(); + // If we returned some replicas, skip remaining replicas from child policy + // to avoid duplicates. If no replicas were returned (all DOWN/IGNORED), + // allow full child policy fallback including replica hosts. + if (returnedHosts != null && replicas.contains(host)) { + continue; } + return host; } return endOfData(); @@ -477,7 +490,10 @@ private Iterator newQueryPlanRegular( private Iterator newQueryPlanPreserveReplicaOrder( String keyspace, Statement statement, List replicas) { - return new PreserveReplicaOrderIterator(keyspace, statement, replicas.iterator()); + if (replicas.isEmpty()) { + return childPolicy.newQueryPlan(keyspace, statement); + } + return new PreserveReplicaOrderIterator(keyspace, statement, replicas); } @Override diff --git a/driver-core/src/test/java/com/datastax/driver/core/policies/TokenAwarePolicyTest.java b/driver-core/src/test/java/com/datastax/driver/core/policies/TokenAwarePolicyTest.java index 80a0dd66ff5..eeb57e91b2b 100644 --- a/driver-core/src/test/java/com/datastax/driver/core/policies/TokenAwarePolicyTest.java +++ b/driver-core/src/test/java/com/datastax/driver/core/policies/TokenAwarePolicyTest.java @@ -75,6 +75,7 @@ public class TokenAwarePolicyTest { private Host host2 = mock(Host.class); private Host host3 = mock(Host.class); private Host host4 = mock(Host.class); + private Host host5 = mock(Host.class); private LoadBalancingPolicy childPolicy; private Cluster cluster; @@ -109,6 +110,7 @@ public void initMocks() { when(host2.isUp()).thenReturn(true); when(host3.isUp()).thenReturn(true); when(host4.isUp()).thenReturn(true); + when(host5.isUp()).thenReturn(true); } @DataProvider(name = "shuffleProvider") @@ -163,6 +165,8 @@ public void should_prioritize_local_replicas_for_lwt(TokenAwarePolicy.ReplicaOrd when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); when(childPolicy.distance(host1)).thenReturn(HostDistance.REMOTE); when(childPolicy.distance(host2)).thenReturn(HostDistance.LOCAL); + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4, host3, host2, host1).iterator()); TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); @@ -170,8 +174,8 @@ public void should_prioritize_local_replicas_for_lwt(TokenAwarePolicy.ReplicaOrd // when Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); - // then: local replica first, then remaining replicas only - assertThat(queryPlan).containsExactly(host2, host1); + // then: local replica first, remote replica, then non-replicas from child policy + assertThat(queryPlan).containsExactly(host2, host1, host4, host3); } @Test(groups = "unit", dataProvider = "shuffleProvider") @@ -184,6 +188,8 @@ public void should_preserve_replica_order_for_lwt(TokenAwarePolicy.ReplicaOrderi when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); when(metadata.getReplicasList(Metadata.quote(KEYSPACE), null, null, routingKey)) .thenReturn(Lists.newArrayList(host2, host3, host1)); + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4, host3, host2, host1).iterator()); TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); @@ -191,8 +197,8 @@ public void should_preserve_replica_order_for_lwt(TokenAwarePolicy.ReplicaOrderi // when Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); - // then: replica order preserved and only replicas returned - assertThat(queryPlan).containsExactly(host2, host3, host1); + // then: replica order preserved, then non-replicas from child policy + assertThat(queryPlan).containsExactly(host2, host3, host1, host4); } @Test(groups = "unit") @@ -241,14 +247,18 @@ public void should_filter_down_replicas_for_lwt(TokenAwarePolicy.ReplicaOrdering when(childPolicy.distance(host3)).thenReturn(HostDistance.REMOTE); when(host3.isUp()).thenReturn(false); + // host4 is a non-replica available via child policy + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4).iterator()); + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); // when Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); - // then: only UP replicas are returned (host1 and host3 are DOWN so excluded) - assertThat(queryPlan).containsExactly(host2); + // then: UP replicas first, then non-replicas from child policy + assertThat(queryPlan).containsExactly(host2, host4); } @Test(groups = "unit", dataProvider = "shuffleProvider") @@ -274,20 +284,24 @@ public void should_filter_ignored_replicas_for_lwt(TokenAwarePolicy.ReplicaOrder when(childPolicy.distance(host3)).thenReturn(HostDistance.REMOTE); when(host3.isUp()).thenReturn(true); + // host4 is a non-replica available via child policy + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4).iterator()); + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); // when Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); - // then: IGNORED replicas are excluded (host2), local first then remote - assertThat(queryPlan).containsExactly(host1, host3); + // then: IGNORED replicas excluded, local first, remote, then non-replicas + assertThat(queryPlan).containsExactly(host1, host3, host4); } @Test(groups = "unit", dataProvider = "shuffleProvider") public void should_filter_down_and_ignored_replicas_for_lwt( TokenAwarePolicy.ReplicaOrdering ordering) { - // given: LWT statement with mixed replica states + // given: LWT statement with mixed replica states (all 4 hosts are replicas) Statement lwtStatement = mock(Statement.class); when(lwtStatement.isLWT()).thenReturn(true); when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) @@ -312,6 +326,10 @@ public void should_filter_down_and_ignored_replicas_for_lwt( when(childPolicy.distance(host4)).thenReturn(HostDistance.REMOTE); when(host4.isUp()).thenReturn(true); + // child policy returns empty since all hosts are replicas + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList().iterator()); + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); @@ -390,14 +408,18 @@ public void should_return_all_local_replicas_when_all_replicas_are_local( when(childPolicy.distance(host3)).thenReturn(HostDistance.LOCAL); when(host3.isUp()).thenReturn(true); + // host4 is a non-replica available via child policy + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4).iterator()); + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); policy.init(cluster, null); // when Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); - // then: should return all local replicas without NPE (nonLocalReplicas remains null) - assertThat(queryPlan).containsExactly(host1, host2, host3); + // then: all local replicas first, then non-replicas + assertThat(queryPlan).containsExactly(host1, host2, host3, host4); } @Test(groups = "unit", dataProvider = "shuffleProvider") @@ -433,6 +455,153 @@ public void should_allow_child_policy_to_retry_down_replicas_in_fallback( assertThat(queryPlan).containsExactly(host1, host2, host3, host4); } + @Test(groups = "unit", dataProvider = "shuffleProvider") + public void should_preserve_replica_order_with_all_remote_replicas( + TokenAwarePolicy.ReplicaOrdering ordering) { + // given: LWT statement where all replicas are in a remote DC + Statement lwtStatement = mock(Statement.class); + when(lwtStatement.isLWT()).thenReturn(true); + when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) + .thenReturn(routingKey); + when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); + when(metadata.getReplicasList(Metadata.quote(KEYSPACE), null, null, routingKey)) + .thenReturn(Lists.newArrayList(host1, host2)); + + // Both replicas are REMOTE + when(childPolicy.distance(host1)).thenReturn(HostDistance.REMOTE); + when(childPolicy.distance(host2)).thenReturn(HostDistance.REMOTE); + when(childPolicy.distance(host3)).thenReturn(HostDistance.LOCAL); + when(childPolicy.distance(host4)).thenReturn(HostDistance.LOCAL); + + // Child policy returns non-replica local nodes + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host3, host4).iterator()); + + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); + policy.init(cluster, null); + + // when + Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); + + // then: remote replicas first (preserving order), then local non-replicas from child policy + assertThat(queryPlan).containsExactly(host1, host2, host3, host4); + } + + @Test(groups = "unit", dataProvider = "shuffleProvider") + public void should_fallback_to_child_policy_when_lwt_has_no_routing_key( + TokenAwarePolicy.ReplicaOrdering ordering) { + // given: LWT statement with no routing key + Statement lwtStatement = mock(Statement.class); + when(lwtStatement.isLWT()).thenReturn(true); + when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) + .thenReturn(null); + when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); + + // Child policy returns all hosts + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host1, host2, host3, host4).iterator()); + + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); + policy.init(cluster, null); + + // when + Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); + + // then: falls back to child policy since no routing key is available + assertThat(queryPlan).containsExactly(host1, host2, host3, host4); + } + + @Test(groups = "unit", dataProvider = "shuffleProvider") + public void should_fallback_to_child_policy_when_lwt_has_no_keyspace( + TokenAwarePolicy.ReplicaOrdering ordering) { + // given: LWT statement with no keyspace and no logged keyspace + Statement lwtStatement = mock(Statement.class); + when(lwtStatement.isLWT()).thenReturn(true); + when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) + .thenReturn(routingKey); + when(lwtStatement.getKeyspace()).thenReturn(null); + + // Child policy returns all hosts + when(childPolicy.newQueryPlan(null, lwtStatement)) + .thenReturn(Lists.newArrayList(host1, host2, host3, host4).iterator()); + + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); + policy.init(cluster, null); + + // when: both statement keyspace and logged keyspace are null + Iterator queryPlan = policy.newQueryPlan(null, lwtStatement); + + // then: falls back to child policy since keyspace is unknown + assertThat(queryPlan).containsExactly(host1, host2, host3, host4); + } + + @Test(groups = "unit", dataProvider = "shuffleProvider") + public void should_produce_deterministic_query_plan_for_lwt( + TokenAwarePolicy.ReplicaOrdering ordering) { + // given: LWT statement with routing key and replicas + Statement lwtStatement = mock(Statement.class); + when(lwtStatement.isLWT()).thenReturn(true); + when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) + .thenReturn(routingKey); + when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); + when(metadata.getReplicasList(Metadata.quote(KEYSPACE), null, null, routingKey)) + .thenReturn(Lists.newArrayList(host2, host1)); + when(childPolicy.distance(host1)).thenReturn(HostDistance.LOCAL); + when(childPolicy.distance(host2)).thenReturn(HostDistance.LOCAL); + + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); + policy.init(cluster, null); + + // when: query plan is generated multiple times + for (int i = 0; i < 3; i++) { + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host3, host4).iterator()); + Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); + + // then: replicas always appear first in the same order, followed by non-replicas + List plan = Lists.newArrayList(queryPlan); + assertThat(plan).hasSize(4); + assertThat(plan.get(0)).isEqualTo(host2); + assertThat(plan.get(1)).isEqualTo(host1); + } + } + + @Test(groups = "unit", dataProvider = "shuffleProvider") + public void should_order_local_replicas_then_remote_replicas_then_non_replicas( + TokenAwarePolicy.ReplicaOrdering ordering) { + // given: LWT statement with mixed local/remote replicas and non-replica nodes + Statement lwtStatement = mock(Statement.class); + when(lwtStatement.isLWT()).thenReturn(true); + when(lwtStatement.getRoutingKey(any(ProtocolVersion.class), any(CodecRegistry.class))) + .thenReturn(routingKey); + when(lwtStatement.getKeyspace()).thenReturn(KEYSPACE); + when(metadata.getReplicasList(Metadata.quote(KEYSPACE), null, null, routingKey)) + .thenReturn(Lists.newArrayList(host1, host2, host3)); + + // host1 is LOCAL replica + when(childPolicy.distance(host1)).thenReturn(HostDistance.LOCAL); + // host2 is REMOTE replica + when(childPolicy.distance(host2)).thenReturn(HostDistance.REMOTE); + // host3 is LOCAL replica + when(childPolicy.distance(host3)).thenReturn(HostDistance.LOCAL); + // host4 and host5 are non-replica nodes + when(childPolicy.distance(host4)).thenReturn(HostDistance.LOCAL); + when(childPolicy.distance(host5)).thenReturn(HostDistance.REMOTE); + + when(childPolicy.newQueryPlan(KEYSPACE, lwtStatement)) + .thenReturn(Lists.newArrayList(host4, host5).iterator()); + + TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering); + policy.init(cluster, null); + + // when + Iterator queryPlan = policy.newQueryPlan(KEYSPACE, lwtStatement); + + // then: local replicas first (host1, host3), then remote replica (host2), + // then non-replicas from child policy (host4, host5) + assertThat(queryPlan).containsExactly(host1, host3, host2, host4, host5); + } + /** * Ensures that {@link TokenAwarePolicy} will shuffle discovered replicas depending on the value * of shuffleReplicas used when constructing with {@link