[C-S] Make it use DCAwareRoundRobinPolicy unless rack is provided

Currently c-s uses RackAwareRoundRobinPolicy in any case which leads it targeting particular rack only, even if rack is not provided. This PR makes it use DCAwareRoundRobinPolicy unless rack is provided. Ref: scylladb/cassandra-stress#20 Ref: scylladb/cassandra-stress#21 Closes: #404
scylladb · Sep 12, 2024 · e505a6d · e505a6d
1 parent 0b4accd
commit e505a6d
Showing 1 changed file with 15 additions and 12 deletions.
diff --git a/tools/stress/src/org/apache/cassandra/stress/util/JavaDriverClient.java b/tools/stress/src/org/apache/cassandra/stress/util/JavaDriverClient.java
@@ -25,6 +25,7 @@
 
 import com.datastax.driver.core.*;
 import com.datastax.driver.core.policies.RackAwareRoundRobinPolicy;
+import com.datastax.driver.core.policies.DCAwareRoundRobinPolicy;
 import com.datastax.driver.core.policies.LoadBalancingPolicy;
 import com.datastax.driver.core.policies.TokenAwarePolicy;
 import com.datastax.driver.core.policies.TokenAwarePolicy.ReplicaOrdering;
@@ -94,21 +95,23 @@ public JavaDriverClient(StressSettings settings, List<String> hosts, int port, E
 
     private LoadBalancingPolicy loadBalancingPolicy(StressSettings settings)
     {
-        RackAwareRoundRobinPolicy.Builder policyBuilder = RackAwareRoundRobinPolicy.builder();
-
-        if (settings.node.datacenter != null)
-            policyBuilder.withLocalDc(settings.node.datacenter);
-        if (settings.node.rack != null)
-            policyBuilder.withLocalRack(settings.node.rack);
-
         LoadBalancingPolicy ret = null;
-        if (settings.node.datacenter != null)
-            ret = policyBuilder.build();
 
+        if (settings.node.rack != null) {
+            RackAwareRoundRobinPolicy.Builder policyBuilder = RackAwareRoundRobinPolicy.builder();
+            if (settings.node.datacenter != null)
+                policyBuilder.withLocalDc(settings.node.datacenter);
+            policyBuilder = policyBuilder.withLocalRack(settings.node.rack);
+            ret = policyBuilder.build();
+        } else {
+            DCAwareRoundRobinPolicy.Builder policyBuilder = DCAwareRoundRobinPolicy.builder();
+            if (settings.node.datacenter != null)
+                policyBuilder.withLocalDc(settings.node.datacenter);
+            ret = policyBuilder.build();
+        }
         if (settings.node.isWhiteList)
-            ret = new WhiteListPolicy(ret == null ? policyBuilder.build() : ret, settings.node.resolveAll(settings.port.nativePort));
-
-        return new TokenAwarePolicy(ret == null ? policyBuilder.build() : ret, ReplicaOrdering.NEUTRAL);
+            ret = new WhiteListPolicy(ret, settings.node.resolveAll(settings.port.nativePort));
+        return new TokenAwarePolicy(ret, ReplicaOrdering.NEUTRAL);
     }
 
     public PreparedStatement prepare(String query)