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