JAVA-3142: Ability to specify ordering of remote local dc's via new configuration for graceful automatic failovers

nitinitt · nitinitt · commit a3bd191df3e8 · 2024-05-07T16:51:57.000-07:00
patch by Nitin Chhabra; reviewed by Alexandre Dutra, Andy Tolbert, and Bret McGuire for JAVA-3142
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
@@ -982,7 +982,13 @@ public enum DefaultDriverOption implements DriverOption {
    * <p>Value-type: {@link java.time.Duration}
    */
   SSL_KEYSTORE_RELOAD_INTERVAL("advanced.ssl-engine-factory.keystore-reload-interval"),
-  ;
+  /**
+   * Ordered preference list of remote dcs optionally supplied for automatic failover.
+   *
+   * <p>Value type: {@link java.util.List List}&#60;{@link String}&#62;
+   */
+  LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS(
+      "advanced.load-balancing-policy.dc-failover.preferred-remote-dcs");
 
   private final String path;
 
diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java
@@ -381,6 +381,8 @@ protected static void fillWithDriverDefaults(OptionsMap map) {
     map.put(TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC, 0);
     map.put(TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS, false);
     map.put(TypedDriverOption.METRICS_GENERATE_AGGREGABLE_HISTOGRAMS, true);
+    map.put(
+        TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS, ImmutableList.of(""));
   }
 
   @Immutable
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
@@ -892,6 +892,16 @@ public String toString() {
               DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS,
               GenericType.BOOLEAN);
 
+  /**
+   * Ordered preference list of remote dcs optionally supplied for automatic failover and included
+   * in query plan. This feature is enabled only when max-nodes-per-remote-dc is greater than 0.
+   */
+  public static final TypedDriverOption<List<String>>
+      LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS =
+          new TypedDriverOption<>(
+              DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS,
+              GenericType.listOf(String.class));
+
   private static Iterable<TypedDriverOption<?>> introspectBuiltInValues() {
     try {
       ImmutableList.Builder<TypedDriverOption<?>> result = ImmutableList.builder();
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
@@ -45,10 +45,14 @@
 import com.datastax.oss.driver.internal.core.util.collection.QueryPlan;
 import com.datastax.oss.driver.internal.core.util.collection.SimpleQueryPlan;
 import com.datastax.oss.driver.shaded.guava.common.base.Predicates;
+import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
+import com.datastax.oss.driver.shaded.guava.common.collect.Sets;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.nio.ByteBuffer;
 import java.util.Collections;
+import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
@@ -117,6 +121,7 @@ public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
   private volatile NodeDistanceEvaluator nodeDistanceEvaluator;
   private volatile String localDc;
   private volatile NodeSet liveNodes;
+  private final LinkedHashSet<String> preferredRemoteDcs;
 
   public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String profileName) {
     this.context = (InternalDriverContext) context;
@@ -131,6 +136,11 @@ public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String
         this.context
             .getConsistencyLevelRegistry()
             .nameToLevel(profile.getString(DefaultDriverOption.REQUEST_CONSISTENCY));
+
+    preferredRemoteDcs =
+        new LinkedHashSet<>(
+            profile.getStringList(
+                DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS));
   }
 
   /**
@@ -320,27 +330,59 @@ protected Queue<Node> maybeAddDcFailover(@Nullable Request request, @NonNull Que
         return local;
       }
     }
-    QueryPlan remote =
-        new LazyQueryPlan() {
-
-          @Override
-          protected Object[] computeNodes() {
-            Object[] remoteNodes =
-                liveNodes.dcs().stream()
-                    .filter(Predicates.not(Predicates.equalTo(localDc)))
-                    .flatMap(dc -> liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc))
-                    .toArray();
-
-            int remoteNodesLength = remoteNodes.length;
-            if (remoteNodesLength == 0) {
-              return EMPTY_NODES;
-            }
-            shuffleHead(remoteNodes, remoteNodesLength);
-            return remoteNodes;
-          }
-        };
-
-    return new CompositeQueryPlan(local, remote);
+    if (preferredRemoteDcs.isEmpty()) {
+      return new CompositeQueryPlan(local, buildRemoteQueryPlanAll());
+    }
+    return new CompositeQueryPlan(local, buildRemoteQueryPlanPreferred());
+  }
+
+  private QueryPlan buildRemoteQueryPlanAll() {
+
+    return new LazyQueryPlan() {
+      @Override
+      protected Object[] computeNodes() {
+
+        Object[] remoteNodes =
+            liveNodes.dcs().stream()
+                .filter(Predicates.not(Predicates.equalTo(localDc)))
+                .flatMap(dc -> liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc))
+                .toArray();
+        if (remoteNodes.length == 0) {
+          return EMPTY_NODES;
+        }
+        shuffleHead(remoteNodes, remoteNodes.length);
+        return remoteNodes;
+      }
+    };
+  }
+
+  private QueryPlan buildRemoteQueryPlanPreferred() {
+
+    Set<String> dcs = liveNodes.dcs();
+    List<String> orderedDcs = Lists.newArrayListWithCapacity(dcs.size());
+    orderedDcs.addAll(preferredRemoteDcs);
+    orderedDcs.addAll(Sets.difference(dcs, preferredRemoteDcs));
+
+    QueryPlan[] queryPlans =
+        orderedDcs.stream()
+            .filter(Predicates.not(Predicates.equalTo(localDc)))
+            .map(
+                (dc) -> {
+                  return new LazyQueryPlan() {
+                    @Override
+                    protected Object[] computeNodes() {
+                      Object[] rv = liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc).toArray();
+                      if (rv.length == 0) {
+                        return EMPTY_NODES;
+                      }
+                      shuffleHead(rv, rv.length);
+                      return rv;
+                    }
+                  };
+                })
+            .toArray(QueryPlan[]::new);
+
+    return new CompositeQueryPlan(queryPlans);
   }
 
   /** Exposed as a protected method so that it can be accessed by tests */
diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf
@@ -574,6 +574,11 @@ datastax-java-driver {
       # Modifiable at runtime: no
       # Overridable in a profile: yes
       allow-for-local-consistency-levels = false
+      # Ordered preference list of remote dc's (in order) optionally supplied for automatic failover. While building a query plan, the driver uses the DC's supplied in order together with max-nodes-per-remote-dc
+      # Required: no
+      # Modifiable at runtime: no
+      # Overridable in a profile: no
+      preferred-remote-dcs = [""]
     }
   }
 
diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicyPreferredRemoteDcsTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicyPreferredRemoteDcsTest.java
@@ -0,0 +1,184 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.datastax.oss.driver.internal.core.loadbalancing;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+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 com.datastax.oss.driver.internal.core.metadata.DefaultNode;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
+import java.util.Map;
+import java.util.UUID;
+import org.junit.Test;
+import org.mockito.Mock;
+
+public class BasicLoadBalancingPolicyPreferredRemoteDcsTest
+    extends BasicLoadBalancingPolicyDcFailoverTest {
+  @Mock protected DefaultNode node10;
+  @Mock protected DefaultNode node11;
+  @Mock protected DefaultNode node12;
+  @Mock protected DefaultNode node13;
+  @Mock protected DefaultNode node14;
+
+  @Override
+  @Test
+  public void should_prioritize_single_replica() {
+    when(request.getRoutingKeyspace()).thenReturn(KEYSPACE);
+    when(request.getRoutingKey()).thenReturn(ROUTING_KEY);
+    when(tokenMap.getReplicas(KEYSPACE, ROUTING_KEY)).thenReturn(ImmutableSet.of(node3));
+
+    // node3 always first, round-robin on the rest
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node3, node1, node2, node4, node5, node9, node10, node6, node7, node12, node13);
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node3, node2, node4, node5, node1, node9, node10, node6, node7, node12, node13);
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node3, node4, node5, node1, node2, node9, node10, node6, node7, node12, node13);
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node3, node5, node1, node2, node4, node9, node10, node6, node7, node12, node13);
+
+    // Should not shuffle replicas since there is only one
+    verify(policy, never()).shuffleHead(any(), eq(1));
+    // But should shuffle remote nodes
+    verify(policy, times(12)).shuffleHead(any(), eq(2));
+  }
+
+  @Override
+  @Test
+  public void should_prioritize_and_shuffle_replicas() {
+    when(request.getRoutingKeyspace()).thenReturn(KEYSPACE);
+    when(request.getRoutingKey()).thenReturn(ROUTING_KEY);
+    when(tokenMap.getReplicas(KEYSPACE, ROUTING_KEY))
+        .thenReturn(ImmutableSet.of(node1, node2, node3, node6, node9));
+
+    // node 6 and 9 being in a remote DC, they don't get a boost for being a replica
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node1, node2, node3, node4, node5, node9, node10, node6, node7, node12, node13);
+    assertThat(policy.newQueryPlan(request, session))
+        .containsExactly(
+            node1, node2, node3, node5, node4, node9, node10, node6, node7, node12, node13);
+
+    // should shuffle replicas
+    verify(policy, times(2)).shuffleHead(any(), eq(3));
+    // should shuffle remote nodes
+    verify(policy, times(6)).shuffleHead(any(), eq(2));
+    // No power of two choices with only two replicas
+    verify(session, never()).getPools();
+  }
+
+  @Override
+  protected void assertRoundRobinQueryPlans() {
+    for (int i = 0; i < 3; i++) {
+      assertThat(policy.newQueryPlan(request, session))
+          .containsExactly(
+              node1, node2, node3, node4, node5, node9, node10, node6, node7, node12, node13);
+      assertThat(policy.newQueryPlan(request, session))
+          .containsExactly(
+              node2, node3, node4, node5, node1, node9, node10, node6, node7, node12, node13);
+      assertThat(policy.newQueryPlan(request, session))
+          .containsExactly(
+              node3, node4, node5, node1, node2, node9, node10, node6, node7, node12, node13);
+      assertThat(policy.newQueryPlan(request, session))
+          .containsExactly(
+              node4, node5, node1, node2, node3, node9, node10, node6, node7, node12, node13);
+      assertThat(policy.newQueryPlan(request, session))
+          .containsExactly(
+              node5, node1, node2, node3, node4, node9, node10, node6, node7, node12, node13);
+    }
+
+    verify(policy, atLeast(15)).shuffleHead(any(), eq(2));
+  }
+
+  @Override
+  protected BasicLoadBalancingPolicy createAndInitPolicy() {
+    when(node4.getDatacenter()).thenReturn("dc1");
+    when(node5.getDatacenter()).thenReturn("dc1");
+    when(node6.getDatacenter()).thenReturn("dc2");
+    when(node7.getDatacenter()).thenReturn("dc2");
+    when(node8.getDatacenter()).thenReturn("dc2");
+    when(node9.getDatacenter()).thenReturn("dc3");
+    when(node10.getDatacenter()).thenReturn("dc3");
+    when(node11.getDatacenter()).thenReturn("dc3");
+    when(node12.getDatacenter()).thenReturn("dc4");
+    when(node13.getDatacenter()).thenReturn("dc4");
+    when(node14.getDatacenter()).thenReturn("dc4");
+
+    // Accept 2 nodes per remote DC
+    when(defaultProfile.getInt(
+            DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC))
+        .thenReturn(2);
+    when(defaultProfile.getBoolean(
+            DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS))
+        .thenReturn(false);
+
+    when(defaultProfile.getStringList(
+            DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS))
+        .thenReturn(ImmutableList.of("dc3", "dc2"));
+
+    // Use a subclass to disable shuffling, we just spy to make sure that the shuffling method was
+    // called (makes tests easier)
+    BasicLoadBalancingPolicy policy =
+        spy(
+            new BasicLoadBalancingPolicy(context, DriverExecutionProfile.DEFAULT_NAME) {
+              @Override
+              protected void shuffleHead(Object[] currentNodes, int headLength) {
+                // nothing (keep in same order)
+              }
+            });
+    Map<UUID, Node> nodes =
+        ImmutableMap.<UUID, Node>builder()
+            .put(UUID.randomUUID(), node1)
+            .put(UUID.randomUUID(), node2)
+            .put(UUID.randomUUID(), node3)
+            .put(UUID.randomUUID(), node4)
+            .put(UUID.randomUUID(), node5)
+            .put(UUID.randomUUID(), node6)
+            .put(UUID.randomUUID(), node7)
+            .put(UUID.randomUUID(), node8)
+            .put(UUID.randomUUID(), node9)
+            .put(UUID.randomUUID(), node10)
+            .put(UUID.randomUUID(), node11)
+            .put(UUID.randomUUID(), node12)
+            .put(UUID.randomUUID(), node13)
+            .put(UUID.randomUUID(), node14)
+            .build();
+    policy.init(nodes, distanceReporter);
+    assertThat(policy.getLiveNodes().dc("dc1")).containsExactly(node1, node2, node3, node4, node5);
+    assertThat(policy.getLiveNodes().dc("dc2")).containsExactly(node6, node7); // only 2 allowed
+    assertThat(policy.getLiveNodes().dc("dc3")).containsExactly(node9, node10); // only 2 allowed
+    assertThat(policy.getLiveNodes().dc("dc4")).containsExactly(node12, node13); // only 2 allowed
+    return policy;
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -381,6 +381,8 @@ protected static void fillWithDriverDefaults(OptionsMap map) {`
`381`	`381`	`map.put(TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC, 0);`
`382`	`382`	`map.put(TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS, false);`
`383`	`383`	`map.put(TypedDriverOption.METRICS_GENERATE_AGGREGABLE_HISTOGRAMS, true);`
	`384`	`+ map.put(`
	`385`	`+ TypedDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS, ImmutableList.of(""));`
`384`	`386`	`}`
`385`	`387`
`386`	`388`	`@Immutable`
Original file line number	Diff line number	Diff line change
`@@ -574,6 +574,11 @@ datastax-java-driver {`
`574`	`574`	`# Modifiable at runtime: no`
`575`	`575`	`# Overridable in a profile: yes`
`576`	`576`	`allow-for-local-consistency-levels = false`
	`577`	`+ # Ordered preference list of remote dc's (in order) optionally supplied for automatic failover. While building a query plan, the driver uses the DC's supplied in order together with max-nodes-per-remote-dc`
	`578`	`+ # Required: no`
	`579`	`+ # Modifiable at runtime: no`
	`580`	`+ # Overridable in a profile: no`
	`581`	`+ preferred-remote-dcs = [""]`
`577`	`582`	`}`
`578`	`583`	`}`
`579`	`584`