Add rack awareness to default load balancing policy

Gor027 · web-flow · commit 437f94bdad65 · 2023-06-27T17:48:14.000+02:00
The default load balancing policy will prefer replicas that are in the
local datacenter and in the local rack.
The rack awareness feature is optional, to enable it a local rack should
be supplied through configuration by defining the option
`DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK`
`basic.load-balancing-policy.local-rack`, otherwise it will be
disabled and unlike the local datacenter it will not be implicitly
fetched from the provided contact points.
diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java
@@ -93,6 +93,12 @@ public enum DefaultDriverOption implements DriverOption {
    * <p>Value-type: {@link String}
    */
   LOAD_BALANCING_LOCAL_DATACENTER("basic.load-balancing-policy.local-datacenter"),
+  /**
+   * The rack that is considered "local".
+   *
+   * <p>Value-type: {@link String}
+   */
+  LOAD_BALANCING_LOCAL_RACK("basic.load-balancing-policy.local-rack"),
   /**
    * A custom filter to include/exclude nodes.
    *
diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java
@@ -137,6 +137,9 @@ public String toString() {
   public static final TypedDriverOption<String> LOAD_BALANCING_LOCAL_DATACENTER =
       new TypedDriverOption<>(
           DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER, GenericType.STRING);
+
+  public static final TypedDriverOption<String> LOAD_BALANCING_LOCAL_RACK =
+      new TypedDriverOption<>(DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK, GenericType.STRING);
   /**
    * A custom filter to include/exclude nodes.
    *
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java
@@ -40,6 +40,7 @@
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.loadbalancing.helper.DefaultNodeDistanceEvaluatorHelper;
 import com.datastax.oss.driver.internal.core.loadbalancing.helper.OptionalLocalDcHelper;
+import com.datastax.oss.driver.internal.core.loadbalancing.helper.OptionalLocalRackHelper;
 import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.DcAgnosticNodeSet;
 import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.MultiDcNodeSet;
 import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.NodeSet;
@@ -121,6 +122,7 @@ public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
   private volatile DistanceReporter distanceReporter;
   private volatile NodeDistanceEvaluator nodeDistanceEvaluator;
   private volatile String localDc;
+  private volatile String localRack;
   private volatile NodeSet liveNodes;
 
   public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String profileName) {
@@ -154,6 +156,11 @@ protected String getLocalDatacenter() {
     return localDc;
   }
 
+  @Nullable
+  protected String getLocalRack() {
+    return localRack;
+  }
+
   /** @return The nodes currently considered as live. */
   protected NodeSet getLiveNodes() {
     return liveNodes;
@@ -163,6 +170,8 @@ protected NodeSet getLiveNodes() {
   public void init(@NonNull Map<UUID, Node> nodes, @NonNull DistanceReporter distanceReporter) {
     this.distanceReporter = distanceReporter;
     localDc = discoverLocalDc(nodes).orElse(null);
+    // If local datacenter is not provided then the rack awareness is disabled
+    localRack = localDc != null ? discoverLocalRack(nodes).orElse(null) : null;
     nodeDistanceEvaluator = createNodeDistanceEvaluator(localDc, nodes);
     liveNodes =
         localDc == null
@@ -207,6 +216,11 @@ protected Optional<String> discoverLocalDc(@NonNull Map<UUID, Node> nodes) {
     return new OptionalLocalDcHelper(context, profile, logPrefix).discoverLocalDc(nodes);
   }
 
+  @NonNull
+  protected Optional<String> discoverLocalRack(@NonNull Map<UUID, Node> nodes) {
+    return new OptionalLocalRackHelper(profile, logPrefix).discoverLocalRack(nodes);
+  }
+
   /**
    * Creates a new node distance evaluator to use with this policy.
    *
@@ -356,6 +370,11 @@ protected void shuffleHead(Object[] currentNodes, int headLength) {
     ArrayUtils.shuffleHead(currentNodes, headLength);
   }
 
+  /** Exposed as a protected method so that it can be accessed by tests */
+  protected void shuffleInRange(Object[] currentNodes, int startIndex, int endIndex) {
+    ArrayUtils.shuffleInRange(currentNodes, startIndex, endIndex);
+  }
+
   @Override
   public void onAdd(@NonNull Node node) {
     NodeDistance distance = computeNodeDistance(node);
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java
@@ -35,6 +35,7 @@
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.util.BitSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
@@ -132,21 +133,36 @@ public Queue<Node> newQueryPlan(@Nullable Request request, @Nullable Session ses
 
     Set<Node> allReplicas = getReplicas(request, session);
     int replicaCount = 0; // in currentNodes
+    int localRackReplicaCount = 0; // in currentNodes
+    String localRack = getLocalRack();
 
     if (!allReplicas.isEmpty()) {
 
       // Move replicas to the beginning of the plan
+      // Replicas in local rack should precede other replicas
       for (int i = 0; i < currentNodes.length; i++) {
         Node node = (Node) currentNodes[i];
         if (allReplicas.contains(node)) {
-          ArrayUtils.bubbleUp(currentNodes, i, replicaCount);
+          if (Objects.equals(node.getRack(), localRack)
+              && Objects.equals(node.getDatacenter(), getLocalDatacenter())) {
+            ArrayUtils.bubbleUp(currentNodes, i, localRackReplicaCount);
+            localRackReplicaCount++;
+          } else {
+            ArrayUtils.bubbleUp(currentNodes, i, replicaCount);
+          }
           replicaCount++;
         }
       }
 
       if (replicaCount > 1) {
-
-        shuffleHead(currentNodes, replicaCount);
+        if (localRack != null && localRackReplicaCount > 0) {
+          // Shuffle only replicas that are in the local rack
+          shuffleHead(currentNodes, localRackReplicaCount);
+          // Shuffles only replicas that are not in local rack
+          shuffleInRange(currentNodes, localRackReplicaCount, replicaCount - 1);
+        } else {
+          shuffleHead(currentNodes, replicaCount);
+        }
 
         if (replicaCount > 2) {
 
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/LocalRackHelper.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/LocalRackHelper.java
@@ -0,0 +1,12 @@
+package com.datastax.oss.driver.internal.core.loadbalancing.helper;
+
+import com.datastax.oss.driver.api.core.metadata.Node;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+
+public interface LocalRackHelper {
+  @NonNull
+  Optional<String> discoverLocalRack(@NonNull Map<UUID, Node> nodes);
+}
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalRackHelper.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalRackHelper.java
@@ -0,0 +1,44 @@
+package com.datastax.oss.driver.internal.core.loadbalancing.helper;
+
+import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
+import com.datastax.oss.driver.api.core.metadata.Node;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.util.Map;
+import java.util.Optional;
+import java.util.UUID;
+import net.jcip.annotations.ThreadSafe;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * An implementation of {@link LocalRackHelper} that fetches the local rack from the driver
+ * configuration. If no user-supplied rack can be retrieved, it returns {@link Optional#empty
+ * empty}.
+ */
+@ThreadSafe
+public class OptionalLocalRackHelper implements LocalRackHelper {
+  private static final Logger LOG = LoggerFactory.getLogger(OptionalLocalRackHelper.class);
+
+  @NonNull protected final DriverExecutionProfile profile;
+  @NonNull protected final String logPrefix;
+
+  public OptionalLocalRackHelper(
+      @NonNull DriverExecutionProfile profile, @NonNull String logPrefix) {
+    this.profile = profile;
+    this.logPrefix = logPrefix;
+  }
+
+  @NonNull
+  @Override
+  public Optional<String> discoverLocalRack(@NonNull Map<UUID, Node> nodes) {
+    if (profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK)) {
+      String localRack = profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK);
+      LOG.debug("[{}] Local rack set from configuration: {}", logPrefix, localRack);
+      return Optional.of(localRack);
+    } else {
+      LOG.debug("[{}] Local rack not set, rack awareness will be disabled", logPrefix);
+      return Optional.empty();
+    }
+  }
+}
diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/util/ArrayUtils.java b/core/src/main/java/com/datastax/oss/driver/internal/core/util/ArrayUtils.java
@@ -89,6 +89,34 @@ public static <ElementT> void shuffleHead(
     }
   }
 
+  /**
+   * Shuffles elements of the array in a range.
+   *
+   * @param elements the array to shuffle.
+   * @param startIndex the start index of the range; must be {@code >= 0}.
+   * @param endIndex the end index of the range; must be {@code <= elements.length}.
+   * @see <a
+   *     href="https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm">Modern
+   *     Fisher-Yates shuffle</a>
+   */
+  public static <ElementT> void shuffleInRange(
+      @NonNull ElementT[] elements, int startIndex, int endIndex) {
+    assert startIndex >= 0 && startIndex < elements.length;
+    assert endIndex >= 0 && endIndex < elements.length;
+
+    if (startIndex > endIndex) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Start index %d should be less than or equal to endIndex %d", startIndex, endIndex));
+    }
+
+    final ThreadLocalRandom random = ThreadLocalRandom.current();
+    for (int i = endIndex; i > startIndex; i--) {
+      int j = random.nextInt(startIndex, i + 1);
+      swap(elements, i, j);
+    }
+  }
+
   /** Rotates the elements in the specified range by the specified amount (round-robin). */
   public static <ElementT> void rotate(
       @NonNull ElementT[] elements, int startIndex, int length, int amount) {
diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DcInferringLoadBalancingPolicyQueryPlanTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DcInferringLoadBalancingPolicyQueryPlanTest.java
@@ -32,6 +32,9 @@ protected DcInferringLoadBalancingPolicy createAndInitPolicy() {
               @Override
               protected void shuffleHead(Object[] array, int n) {}
 
+              @Override
+              protected void shuffleInRange(Object[] currentNodes, int startIndex, int endIndex) {}
+
               @Override
               protected long nanoTime() {
                 return nanoTime;
diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicyInitTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicyInitTest.java
@@ -51,6 +51,23 @@ public void should_use_local_dc_if_provided_via_config() {
     assertThat(policy.getLocalDatacenter()).isEqualTo("dc1");
   }
 
+  @Test
+  public void should_use_local_rack_if_provided_via_config() {
+    // Given
+    when(defaultProfile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK)).thenReturn(true);
+    when(defaultProfile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_RACK))
+        .thenReturn("rack1");
+    when(metadataManager.getContactPoints()).thenReturn(ImmutableSet.of(node1));
+
+    DefaultLoadBalancingPolicy policy = createPolicy();
+
+    // When
+    policy.init(ImmutableMap.of(UUID.randomUUID(), node1), distanceReporter);
+
+    // Then
+    assertThat(policy.getLocalRack()).isEqualTo("rack1");
+  }
+
   @Test
   public void should_use_local_dc_if_provided_via_context() {
     // Given
diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicyQueryPlanTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicyQueryPlanTest.java
@@ -331,6 +331,54 @@ public void should_reorder_first_two_replicas_when_first_has_more_in_flight_than
     then(dsePolicy).should(never()).diceRoll1d4();
   }
 
+  @Test
+  public void should_not_shuffle_local_rack_replica_with_all_replicas_healthy() {
+    // Given
+    given(request.getRoutingKeyspace()).willReturn(KEYSPACE);
+    given(request.getRoutingKey()).willReturn(ROUTING_KEY);
+    given(tokenMap.getReplicas(KEYSPACE, null, ROUTING_KEY))
+        .willReturn(ImmutableSet.of(node1, node3, node5));
+    String localRack = "rack1";
+    given(dsePolicy.getLocalRack()).willReturn(localRack);
+    given(node1.getRack()).willReturn(localRack);
+    given(pool1.getInFlight()).willReturn(0);
+    given(pool3.getInFlight()).willReturn(0);
+    given(pool5.getInFlight()).willReturn(0);
+
+    // When
+    Queue<Node> plan1 = dsePolicy.newQueryPlan(request, session);
+    Queue<Node> plan2 = dsePolicy.newQueryPlan(request, session);
+
+    // Then
+    // nodes 1, 3 and 5 always first, round-robin on the rest
+    assertThat(plan1).containsExactly(node1, node3, node5, node2, node4);
+    assertThat(plan2).containsExactly(node1, node3, node5, node4, node2);
+  }
+
+  @Test
+  public void should_prefer_local_rack_replica_with_less_inflight_requests() {
+    // Given
+    given(request.getRoutingKeyspace()).willReturn(KEYSPACE);
+    given(request.getRoutingKey()).willReturn(ROUTING_KEY);
+    given(tokenMap.getReplicas(KEYSPACE, null, ROUTING_KEY))
+        .willReturn(ImmutableSet.of(node1, node3, node5));
+    String localRack = "rack1";
+    given(dsePolicy.getLocalRack()).willReturn(localRack);
+    given(node3.getRack()).willReturn(localRack);
+    given(node5.getRack()).willReturn(localRack);
+    given(pool1.getInFlight()).willReturn(0);
+    given(pool3.getInFlight()).willReturn(20);
+    given(pool5.getInFlight()).willReturn(10);
+
+    // When
+    Queue<Node> plan1 = dsePolicy.newQueryPlan(request, session);
+    Queue<Node> plan2 = dsePolicy.newQueryPlan(request, session);
+
+    // Then
+    assertThat(plan1).containsExactly(node5, node3, node1, node2, node4);
+    assertThat(plan2).containsExactly(node5, node3, node1, node4, node2);
+  }
+
   @Override
   protected DefaultLoadBalancingPolicy createAndInitPolicy() {
     DefaultLoadBalancingPolicy policy =
@@ -339,6 +387,9 @@ protected DefaultLoadBalancingPolicy createAndInitPolicy() {
               @Override
               protected void shuffleHead(Object[] array, int n) {}
 
+              @Override
+              protected void shuffleInRange(Object[] array, int startIndex, int endIndex) {}
+
               @Override
               protected long nanoTime() {
                 return nanoTime;
diff --git a/manual/core/load_balancing/README.md b/manual/core/load_balancing/README.md
@@ -139,7 +139,24 @@ that case, the driver will connect to 127.0.0.1:9042, and use that node's datace
 for a better out-of-the-box experience for users who have just downloaded the driver; beyond that
 initial development phase, you should provide explicit contact points and a local datacenter.
 
-##### Finding the local datacenter
+##### Rack awareness
+
+The `DefaultLoadBalancingPolicy` and implicitly the `DcInferringLoadBalancingPolicy` prioritize replicas that are in the
+local datacenter, however, sometimes there is a need to prioritize replicas that are in the local rack and to not send
+queries to other replicas in the local datacenter. This will allow to avoid high network traffic between racks/availability zones
+and thus will reduce data transfer costs. The rack-awareness feature is optional and to enable it the local rack should
+be supplied through the configuration:
+
+```
+datastax-java-driver.basic.load-balancing-policy {
+  local-rack = rack1
+}
+```
+
+The feature is disabled by default and unlike the local datacenter it will not be implicitly fetched from the provided
+contact points.
+
+##### Finding the local datacenter & local rack
 
 To check which datacenters are defined in a given cluster, you can run [`nodetool status`]. It will 
 print information about each node in the cluster, grouped by datacenters. Here is an example: 
@@ -165,17 +182,17 @@ UN  <IP5>      1.5 TB     256     ?     <ID5>    rack2
 UN  <IP6>      1.5 TB     256     ?     <ID6>    rack3
 ```
 
-To find out which datacenter should be considered local, you need to first determine which nodes the 
-driver is going to be co-located with, then choose their datacenter as local. In case of doubt, you
+To find out which datacenter and rack(availability zone) should be considered local, you need to first determine which nodes the 
+driver is going to be co-located with, then choose their datacenter and rack as local. In case of doubt, you
 can also use [cqlsh]; if cqlsh is co-located too in the same datacenter, simply run the command 
 below:
 
 ```
-cqlsh> select data_center from system.local;
+cqlsh> select data_center,rack from system.local;
 
-data_center
--------------
-DC1 
+ data_center | rack
+-------------+-------
+ datacenter1 | rack1
 ```
 
 #### Cross-datacenter failover
