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

Author: dkropachev

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,20 +72,18 @@
  * <h3>Lightweight Transaction (LWT) Routing</h3>
  *
  * <p>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 <strong>only replicas</strong> 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:
  *
  * <ul>
  *   <li>Local replicas first: replicas for which the child policy reports {@link HostDistance#LOCAL
  *       LOCAL} distance are returned first, in the order provided by cluster metadata (preserving
  *       primary replica ordering from the token ring).
  *   <li>Remote replicas second: remaining replicas (typically in remote datacenters) are appended,
  *       but only if they are up and not ignored by the child policy.
- *   <li>Replica-only routing when possible: under normal conditions, LWT query plans target only
- *       replicas for the partition in order to reduce coordinator forwarding overhead and improve
- *       performance. When replica information is unavailable, the driver falls back to the child
- *       policy as described in the fallback behavior below, which may include non-replica hosts.
+ *   <li>Non-replica nodes: remaining nodes from the child policy's query plan are appended after
+ *       all replicas, ensuring the query plan always includes all available nodes for failover.
  * </ul>
  *
  * <p><strong>Rack awareness</strong> is intentionally <em>not</em> 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.
    *
-   * <p>Query plan follows a three-pass strategy:
+   * <p>Query plan follows a four-pass strategy:
    *
    * <ol>
    *   <li><strong>Local replicas:</strong> Returns UP replicas marked as LOCAL by the child policy,
    *       in the order provided by cluster metadata (preserving primary replica order).
    *   <li><strong>Remote replicas:</strong> Returns UP replicas marked as REMOTE by the child
    *       policy.
-   *   <li><strong>Child policy fallback:</strong> 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.
+   *   <li><strong>Non-replica nodes:</strong> 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.
+   *   <li><strong>Child policy fallback:</strong> 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.
    * </ol>
    */
   private class PreserveReplicaOrderIterator extends AbstractIterator<Host> {
     private final Iterator<Host> replicasIterator;
+    private final List<Host> replicas;
     private final String keyspace;
     private final Statement statement;
     private List<Host> nonLocalReplicas;
     private Iterator<Host> nonLocalReplicasIterator;
-    private boolean hasReturnedReplicas;
+    private Set<Host> returnedHosts;
     private Iterator<Host> childIterator;
 
-    public PreserveReplicaOrderIterator(
-        String keyspace, Statement statement, Iterator<Host> replicasIterator) {
+    public PreserveReplicaOrderIterator(String keyspace, Statement statement, List<Host> 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<Host> newQueryPlanRegular(
 
   private Iterator<Host> newQueryPlanPreserveReplicaOrder(
       String keyspace, Statement statement, List<Host> replicas) {
-    return new PreserveReplicaOrderIterator(keyspace, statement, replicas.iterator());
+    if (replicas.isEmpty()) {
+      return childPolicy.newQueryPlan(keyspace, statement);
+    }
+    return new PreserveReplicaOrderIterator(keyspace, statement, replicas);
   }
 
   @Override

driver-core/src/test/java/com/datastax/driver/core/policies/TokenAwarePolicyTest.java

@@ -163,15 +163,17 @@ 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);
 
     // when
     Iterator<Host> 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,15 +186,17 @@ 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);
 
     // when
     Iterator<Host> 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 +245,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<Host> 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 +282,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<Host> 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 +324,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.<Host>newArrayList().iterator());
+
     TokenAwarePolicy policy = new TokenAwarePolicy(childPolicy, ordering);
     policy.init(cluster, null);
 
@@ -390,14 +406,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<Host> 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")