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,11 @@ public static Builder builder() {
   private static final String UNSET = "";
 
   private final ConcurrentMap<String, CopyOnWriteArrayList<Host>> perDcLiveHosts =
-      new ConcurrentHashMap<String, CopyOnWriteArrayList<Host>>();
-  private final CopyOnWriteArrayList<Host> liveHostsLocalRackLocalDC =
-      new CopyOnWriteArrayList<Host>();
+      new ConcurrentHashMap<>();
+  private final CopyOnWriteArrayList<Host> liveHostsAllLocalDC = new CopyOnWriteArrayList<>();
+  private final CopyOnWriteArrayList<Host> liveHostsLocalRackLocalDC = new CopyOnWriteArrayList<>();
   private final CopyOnWriteArrayList<Host> liveHostsRemoteRacksLocalDC =
-      new CopyOnWriteArrayList<Host>();
+      new CopyOnWriteArrayList<>();
   @VisibleForTesting final AtomicInteger index = new AtomicInteger();
 
   @VisibleForTesting volatile String localDc;
@@ -147,6 +152,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 {
@@ -240,10 +246,17 @@ public HostDistance distance(Host host) {
   @Override
   public Iterator<Host> newQueryPlan(String loggedKeyspace, final Statement statement) {
 
-    CopyOnWriteArrayList<Host> localLiveHosts = perDcLiveHosts.get(localDc);
-    // Clone for thread safety
-    final List<Host> copyLiveHostsLocalRackLocalDC = cloneList(liveHostsLocalRackLocalDC);
-    final List<Host> copyLiveHostsRemoteRacksLocalDC = cloneList(liveHostsRemoteRacksLocalDC);
+    // For LWT queries, skip rack prioritization and use all local DC hosts equally
+    final boolean isLWT = statement != null && statement.isLWT();
+
+    // For LWT queries, include all local DC hosts in the first part of the plan, not just those in
+    // the local rack
+    final List<Host> copyLiveHostsLocalRackLocalDC =
+        isLWT ? cloneList(liveHostsAllLocalDC) : cloneList(liveHostsLocalRackLocalDC);
+    // For LWT queries, skip the second part of the plan that includes hosts in remote racks of the
+    // local DC
+    final List<Host> copyLiveHostsRemoteRacksLocalDC =
+        isLWT ? Collections.emptyList() : cloneList(liveHostsRemoteRacksLocalDC);
     final int startIdx = index.getAndIncrement();
 
     return new AbstractIterator<Host>() {
@@ -288,7 +301,7 @@ protected Host computeNext() {
           }
 
           ConsistencyLevel cl =
-              statement.getConsistencyLevel() == null
+              statement == null || statement.getConsistencyLevel() == null
                   ? configuration.getQueryOptions().getConsistencyLevel()
                   : statement.getConsistencyLevel();
 
@@ -348,6 +361,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 +379,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[][] {

manual/load_balancing/README.md

@@ -138,6 +138,13 @@ local datacenter and rack. In general, providing the datacenter and rack name ex
 Hosts belonging to the local datacenter are at distance `LOCAL`, and appear first in query plans (in a round-robin
 fashion) with hosts in the local rack having precedence over nodes in remote racks in the local datacenter.
 
+**LWT (Lightweight Transaction) Behavior:** For LWT queries (`Statement.isLWT()` returns true), the policy does not 
+prioritize the local rack. Instead, it round-robins evenly across all hosts in the local datacenter first, then falls 
+back to remote datacenters (if enabled). This design avoids creating rack-level hotspots during Paxos consensus phases. 
+LWT queries involve multiple rounds of coordination between replicas, and concentrating these operations on a single 
+rack can create contention that degrades performance. By distributing LWT load across the entire local datacenter, 
+the driver achieves better throughput while maintaining low latency through datacenter locality.
+
 For example, if there are any UP hosts in the local rack the policy will query those nodes in round-robin fashion:
 * query 1: host1 *(local DC, local rack)*, host2 *(local DC, local rack)*, host3 *(local DC, local rack)*
 * query 2: host2 *(local DC, local rack)*, host3 *(local DC, local rack)*, host1 *(local DC, local rack)*