feat: Implement LWT routing optimization in RackAwareRoundRobinPolicy

Author: nikagra

driver-core/src/main/java/com/datastax/driver/core/policies/RackAwareRoundRobinPolicy.java

@@ -56,6 +56,11 @@
  * but those are always tried after the local nodes. In other words, this policy guarantees that no
  * host in a remote data center will be queried unless no host in the local data center can be
  * reached.
+ *
+ * <p>For LWT (Lightweight Transaction) queries (where {@link Statement#isLWT()} returns {@code
+ * true}), the policy skips local rack prioritization and treats all hosts in the local datacenter
+ * equally, distributing queries in round-robin fashion across the entire local DC. Remote
+ * datacenters are still only used as fallback after all local DC hosts have been tried.
  */
 public class RackAwareRoundRobinPolicy implements LoadBalancingPolicy {
 
@@ -73,11 +78,12 @@ public static Builder builder() {
   private static final String UNSET = "";
 
   private final ConcurrentMap<String, CopyOnWriteArrayList<Host>> perDcLiveHosts =
-      new ConcurrentHashMap<String, CopyOnWriteArrayList<Host>>();
+      new ConcurrentHashMap<>();
+  private final CopyOnWriteArrayList<Host> liveHostsAllLocalDC = new CopyOnWriteArrayList<>();
   private final CopyOnWriteArrayList<Host> liveHostsLocalRackLocalDC =
-      new CopyOnWriteArrayList<Host>();
+      new CopyOnWriteArrayList<>();
   private final CopyOnWriteArrayList<Host> liveHostsRemoteRacksLocalDC =
-      new CopyOnWriteArrayList<Host>();
+      new CopyOnWriteArrayList<>();
   @VisibleForTesting final AtomicInteger index = new AtomicInteger();
 
   @VisibleForTesting volatile String localDc;
@@ -147,6 +153,7 @@ public void init(Cluster cluster, Collection<Host> hosts) {
       else prev.addIfAbsent(host);
 
       if (dc.equals(localDc)) {
+        liveHostsAllLocalDC.add(host);
         if (rack.equals(localRack)) {
           liveHostsLocalRackLocalDC.add(host);
         } else {
@@ -241,14 +248,20 @@ public HostDistance distance(Host host) {
   public Iterator<Host> newQueryPlan(String loggedKeyspace, final Statement statement) {
 
     CopyOnWriteArrayList<Host> localLiveHosts = perDcLiveHosts.get(localDc);
+
+    // For LWT queries, skip rack prioritization and use all local DC hosts equally
+    final boolean isLWT = statement != null && statement.isLWT();
+
     // Clone for thread safety
+    final List<Host> copyLiveHostsAllLocalDC = cloneList(liveHostsAllLocalDC);
     final List<Host> copyLiveHostsLocalRackLocalDC = cloneList(liveHostsLocalRackLocalDC);
     final List<Host> copyLiveHostsRemoteRacksLocalDC = cloneList(liveHostsRemoteRacksLocalDC);
     final int startIdx = index.getAndIncrement();
 
     return new AbstractIterator<Host>() {
 
       private int idx = startIdx;
+      private int remainingLiveHostsAllLocalDC = copyLiveHostsAllLocalDC.size();
       private int remainingLiveHostsLocalRackLocalDC = copyLiveHostsLocalRackLocalDC.size();
       private int remainingLiveHostsRemoteRacksLocalDC = copyLiveHostsRemoteRacksLocalDC.size();
 
@@ -260,22 +273,35 @@ public Iterator<Host> newQueryPlan(String loggedKeyspace, final Statement statem
       @Override
       protected Host computeNext() {
         while (true) {
-          if (remainingLiveHostsLocalRackLocalDC > 0) {
-            remainingLiveHostsLocalRackLocalDC--;
-            int c = idx++ % copyLiveHostsLocalRackLocalDC.size();
-            if (c < 0) {
-              c += copyLiveHostsLocalRackLocalDC.size();
+          // For LWT queries, use all local DC hosts without rack prioritization
+          if (isLWT) {
+            if (remainingLiveHostsAllLocalDC > 0) {
+              remainingLiveHostsAllLocalDC--;
+              int c = idx++ % copyLiveHostsAllLocalDC.size();
+              if (c < 0) {
+                c += copyLiveHostsAllLocalDC.size();
+              }
+              return copyLiveHostsAllLocalDC.get(c);
+            }
+          } else {
+            // For non-LWT queries, use rack-aware routing
+            if (remainingLiveHostsLocalRackLocalDC > 0) {
+              remainingLiveHostsLocalRackLocalDC--;
+              int c = idx++ % copyLiveHostsLocalRackLocalDC.size();
+              if (c < 0) {
+                c += copyLiveHostsLocalRackLocalDC.size();
+              }
+              return copyLiveHostsLocalRackLocalDC.get(c);
             }
-            return copyLiveHostsLocalRackLocalDC.get(c);
-          }
 
-          if (remainingLiveHostsRemoteRacksLocalDC > 0) {
-            remainingLiveHostsRemoteRacksLocalDC--;
-            int c = idx++ % copyLiveHostsRemoteRacksLocalDC.size();
-            if (c < 0) {
-              c += copyLiveHostsRemoteRacksLocalDC.size();
+            if (remainingLiveHostsRemoteRacksLocalDC > 0) {
+              remainingLiveHostsRemoteRacksLocalDC--;
+              int c = idx++ % copyLiveHostsRemoteRacksLocalDC.size();
+              if (c < 0) {
+                c += copyLiveHostsRemoteRacksLocalDC.size();
+              }
+              return copyLiveHostsRemoteRacksLocalDC.get(c);
             }
-            return copyLiveHostsRemoteRacksLocalDC.get(c);
           }
 
           if (currentRemoteDcHosts != null && currentRemoteDcRemaining > 0) {
@@ -288,7 +314,7 @@ protected Host computeNext() {
           }
 
           ConsistencyLevel cl =
-              statement.getConsistencyLevel() == null
+              statement == null || statement.getConsistencyLevel() == null
                   ? configuration.getQueryOptions().getConsistencyLevel()
                   : statement.getConsistencyLevel();
 
@@ -348,6 +374,7 @@ public void onUp(Host host) {
     dcHosts.addIfAbsent(host);
 
     if (dc.equals(localDc)) {
+      liveHostsAllLocalDC.addIfAbsent(host);
       if (rack.equals(localRack)) {
         liveHostsLocalRackLocalDC.add(host);
       } else {
@@ -365,6 +392,7 @@ public void onDown(Host host) {
     if (dcHosts != null) dcHosts.remove(host);
 
     if (dc.equals(localDc)) {
+      liveHostsAllLocalDC.remove(host);
       if (rack.equals(localRack)) {
         liveHostsLocalRackLocalDC.remove(host);
       } else {

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

@@ -86,11 +86,14 @@ public void setUpUnitTests() {
     cluster = mock(Cluster.class);
     Configuration configuration = mock(Configuration.class);
     ProtocolOptions protocolOptions = mock(ProtocolOptions.class);
+    QueryOptions queryOptions = mock(QueryOptions.class);
     Metadata metadata = mock(Metadata.class);
     childPolicy = mock(LoadBalancingPolicy.class);
     when(cluster.getConfiguration()).thenReturn(configuration);
     when(configuration.getCodecRegistry()).thenReturn(codecRegistry);
     when(configuration.getProtocolOptions()).thenReturn(protocolOptions);
+    when(configuration.getQueryOptions()).thenReturn(queryOptions);
+    when(queryOptions.getConsistencyLevel()).thenReturn(ConsistencyLevel.ONE);
     when(protocolOptions.getProtocolVersion()).thenReturn(ProtocolVersion.DEFAULT);
     when(cluster.getMetadata()).thenReturn(metadata);
     when(host1.isUp()).thenReturn(true);
@@ -1107,6 +1110,110 @@ public void should_follow_configuration_on_query_planning(
         .containsExactly(queryPlanForNonLocalConsistencyLevel2.toArray(new Host[0]));
   }
 
+  /**
+   * Ensures that {@link RackAwareRoundRobinPolicy} skips rack prioritization for LWT queries,
+   * treating all local DC hosts equally while still prioritizing local DC over remote DC.
+   *
+   * @test_category load_balancing:rack_aware,lwt
+   */
+  @Test(groups = "unit")
+  public void should_skip_rack_prioritization_for_lwt_queries() {
+    // given: a policy with 4 local DC hosts (2 in local rack, 2 in remote rack) and 2 remote DC
+    // hosts
+    // Initialize hosts in a mixed order: remoteRack, localRack, remoteRack, localRack
+    // This ensures that when LWT skips rack prioritization, we get a different order
+    // than the rack-aware order
+    RackAwareRoundRobinPolicy policy =
+        new RackAwareRoundRobinPolicy("localDC", "localRack", 1, false, false, false);
+    policy.init(cluster, ImmutableList.of(host3, host1, host4, host2, host5, host6));
+
+    // Create a mock LWT statement
+    Statement lwtStatement = mock(Statement.class);
+    when(lwtStatement.isLWT()).thenReturn(true);
+    when(lwtStatement.getConsistencyLevel()).thenReturn(ConsistencyLevel.ONE);
+
+    // when: generating query plans for LWT queries
+    policy.index.set(0);
+    List<Host> queryPlan1 = Lists.newArrayList(policy.newQueryPlan("keyspace", lwtStatement));
+    List<Host> queryPlan2 = Lists.newArrayList(policy.newQueryPlan("keyspace", lwtStatement));
+
+    // then: all 4 local DC hosts should appear before any remote DC host (no rack prioritization)
+    Assertions.assertThat(queryPlan1.subList(0, 4)).containsOnly(host1, host2, host3, host4);
+    Assertions.assertThat(queryPlan2.subList(0, 4)).containsOnly(host1, host2, host3, host4);
+
+    // then: remote DC hosts should appear after all local DC hosts
+    Assertions.assertThat(queryPlan1.subList(4, 5)).containsOnly(host5);
+    Assertions.assertThat(queryPlan2.subList(4, 5)).containsOnly(host5);
+
+    // then: for LWT queries, order should follow insertion order (host3, host1, host4, host2)
+    // not rack-aware order (host1, host2, host3, host4)
+    Assertions.assertThat(queryPlan1).startsWith(host3);
+    Assertions.assertThat(queryPlan2).startsWith(host1);
+  }
+
+  /**
+   * Ensures that {@link RackAwareRoundRobinPolicy} preserves rack-aware routing for non-LWT
+   * queries.
+   *
+   * @test_category load_balancing:rack_aware
+   */
+  @Test(groups = "unit")
+  public void should_preserve_rack_aware_routing_for_non_lwt_queries() {
+    // given: a policy with 4 local DC hosts (2 in local rack, 2 in remote rack) and 2 remote DC
+    // hosts
+    // Initialize hosts in a mixed order to ensure rack-aware routing reorganizes them
+    RackAwareRoundRobinPolicy policy =
+        new RackAwareRoundRobinPolicy("localDC", "localRack", 1, false, false, false);
+    policy.init(cluster, ImmutableList.of(host3, host1, host4, host2, host5, host6));
+
+    // Create a normal (non-LWT) statement
+    Statement normalStatement = mock(Statement.class);
+    when(normalStatement.isLWT()).thenReturn(false);
+    when(normalStatement.getConsistencyLevel()).thenReturn(ConsistencyLevel.ONE);
+
+    // when: generating query plans for non-LWT queries
+    policy.index.set(0);
+    List<Host> queryPlan1 = Lists.newArrayList(policy.newQueryPlan("keyspace", normalStatement));
+    List<Host> queryPlan2 = Lists.newArrayList(policy.newQueryPlan("keyspace", normalStatement));
+
+    // then: local rack hosts (host1, host2) should appear first regardless of init order
+    Assertions.assertThat(queryPlan1.subList(0, 2)).containsOnly(host1, host2);
+    Assertions.assertThat(queryPlan2.subList(0, 2)).containsOnly(host1, host2);
+
+    // then: remote rack local DC hosts (host3, host4) should appear next
+    Assertions.assertThat(queryPlan1.subList(2, 4)).containsOnly(host3, host4);
+    Assertions.assertThat(queryPlan2.subList(2, 4)).containsOnly(host3, host4);
+
+    // then: remote DC hosts should appear last
+    Assertions.assertThat(queryPlan1.subList(4, 5)).containsOnly(host5);
+    Assertions.assertThat(queryPlan2.subList(4, 5)).containsOnly(host5);
+
+    // then: query plans should follow round-robin pattern within rack boundaries
+    Assertions.assertThat(queryPlan1).startsWith(host1);
+    Assertions.assertThat(queryPlan2).startsWith(host2);
+  }
+
+  /**
+   * Ensures that {@link RackAwareRoundRobinPolicy} handles null statement correctly.
+   *
+   * @test_category load_balancing:rack_aware
+   */
+  @Test(groups = "unit")
+  public void should_handle_null_statement() {
+    // given: a policy with hosts in local and remote DC
+    RackAwareRoundRobinPolicy policy =
+        new RackAwareRoundRobinPolicy("localDC", "localRack", 1, false, false, false);
+    policy.init(cluster, ImmutableList.of(host1, host2, host3, host4, host5, host6));
+
+    // when: generating query plan with null statement
+    policy.index.set(0);
+    List<Host> queryPlan = Lists.newArrayList(policy.newQueryPlan("keyspace", null));
+
+    // then: should use rack-aware routing (default behavior for non-LWT)
+    // Local rack hosts should appear first
+    Assertions.assertThat(queryPlan.subList(0, 2)).containsOnly(host1, host2);
+  }
+
   @DataProvider(name = "distanceTestCases")
   public Object[][] distanceTestCases() {
     return new Object[][] {