For get operations, if there are very few local replicas, then shuffle the pool with remote replicas so that the few remaining local replicas don't get all the traffic. (#2653)

Also Define a config with the threshold for minimum number of local replicas to do this.

Co-authored-by: github-actions <[email protected]>

Author: ankagrawal

Diff for: ambry-api/src/main/java/com/github/ambry/config/RouterConfig.java

@@ -38,6 +38,8 @@ public class RouterConfig {
   // This is a theoretical maximum value. Configured value may be much smaller since we might need to respond back to
   // client with either success or failure much sooner.
   public static final int MAX_OVERALL_TIMEOUT_VALUE_FOR_A_REQUEST_IN_MS = 60 * 60 * 1000;
+  // By default a get request should prioritize remote replicas only if there are no local replicas available.
+  public static final int DEFAULT_ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL = 0;
 
   // config keys
   public static final String ROUTER_SCALING_UNIT_COUNT = "router.scaling.unit.count";
@@ -147,6 +149,10 @@ public class RouterConfig {
   public static final String ROUTER_GET_OPERATION_DEPRIORITIZE_BOOTSTRAP_REPLICAS =
       "router.get.operation.deprioritize.bootstrap.replicas";
 
+  // minimum number of local replicas that should be live for a get request so that local replicas are prioritized.
+  public static final String ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL =
+      "router.get.operation.min.local.replica.count.to.prioritize.local";
+
   /**
    * Number of independent scaling units for the router.
    */
@@ -754,6 +760,10 @@ public class RouterConfig {
   @Config(ROUTER_GET_OPERATION_DEPRIORITIZE_BOOTSTRAP_REPLICAS)
   public final boolean routerGetOperationDeprioritizeBootstrapReplicas;
 
+  @Config(ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL)
+  @Default("0")
+  public final int routerGetOperationMinLocalReplicaCountToPrioritizeLocal;
+
   /**
    * Create a RouterConfig instance.
    * @param verifiableProperties the properties map to refer to.
@@ -921,6 +931,8 @@ public RouterConfig(VerifiableProperties verifiableProperties) {
     routerReservedMetadataEnabled = verifiableProperties.getBoolean(RESERVED_METADATA_ENABLED, false);
     routerGetOperationDeprioritizeBootstrapReplicas =
         verifiableProperties.getBoolean(ROUTER_GET_OPERATION_DEPRIORITIZE_BOOTSTRAP_REPLICAS, false);
+    routerGetOperationMinLocalReplicaCountToPrioritizeLocal =
+        verifiableProperties.getInt(ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL, DEFAULT_ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL);
   }
 
   /**

Diff for: ambry-router/src/main/java/com/github/ambry/router/NonBlockingRouterMetrics.java

@@ -292,6 +292,10 @@ public class NonBlockingRouterMetrics {
   public final Counter requestExpiryOnNetworkTimeoutCount;
   // Number of requests timed out due to not being able to sent to network
   public final Counter requestExpiryOnFinalTimeoutCount;
+  // Number of Get requests to server for which remote replicas were prioritized due to few local replicas.
+  public final Counter remoteReplicaPrioritizedForGetDueToFewLocalReplicas;
+  // Number of Get requests to server for which local replicas were shuffled remote replicas due to few local replicas.
+  public final Counter shuffledWithRemoteReplicasForGetDueToFewLocalReplicas;
 
   // Map that stores dataNode-level metrics.
   private final Map<DataNodeId, NodeLevelMetrics> dataNodeToMetrics;
@@ -673,6 +677,12 @@ public NonBlockingRouterMetrics(ClusterMap clusterMap, RouterConfig routerConfig
         metricRegistry.counter(MetricRegistry.name(NonBlockingRouter.class, "RequestExpiryOnNetworkTimeoutCount"));
     requestExpiryOnFinalTimeoutCount =
         metricRegistry.counter(MetricRegistry.name(NonBlockingRouter.class, "RequestExpiryOnFinalTimeoutCount"));
+    remoteReplicaPrioritizedForGetDueToFewLocalReplicas =
+        metricRegistry.counter(MetricRegistry.name(SimpleOperationTracker.class,
+            "RemoteReplicaPrioritizedForGetDueToFewLocalReplicas"));
+    shuffledWithRemoteReplicasForGetDueToFewLocalReplicas =
+        metricRegistry.counter(MetricRegistry.name(SimpleOperationTracker.class,
+            "ShuffledWithRemoteReplicasForGetDueToFewLocalReplicas"));
   }
 
   /**

Diff for: ambry-router/src/main/java/com/github/ambry/router/SimpleOperationTracker.java

@@ -299,6 +299,7 @@ class SimpleOperationTracker implements OperationTracker {
     List<ReplicaId> examinedReplicas = new ArrayList<>();
     originatingDcName = originatingDcName == null ? reassignedOriginDc : originatingDcName;
     int numLocalAndLiveReplicas = 0;
+    int numRemoteOriginatingDcAndLiveReplicas = 0;
     for (ReplicaId replicaId : replicas) {
       examinedReplicas.add(replicaId);
       String replicaDcName = replicaId.getDataNodeId().getDatacenterName();
@@ -310,6 +311,7 @@ class SimpleOperationTracker implements OperationTracker {
           numLocalAndLiveReplicas++;
           addToBeginningOfPool(replicaId);
         } else if (crossColoEnabled && isOriginatingDcReplica) {
+          numRemoteOriginatingDcAndLiveReplicas++;
           addToEndOfPool(replicaId);
         } else if (crossColoEnabled) {
           backupReplicas.addFirst(replicaId);
@@ -349,6 +351,7 @@ class SimpleOperationTracker implements OperationTracker {
     }
 
     maybeDeprioritizeLocalBootstrapReplicas(numLocalAndLiveReplicas);
+    maybeShuffleWithRemoteReplicas(numLocalAndLiveReplicas, numRemoteOriginatingDcAndLiveReplicas);
     totalReplicaCount = replicaPool.size();
 
     // MockPartitionId.getReplicaIds() is returning a shared reference which may cause race condition.
@@ -614,6 +617,39 @@ void maybeDeprioritizeLocalBootstrapReplicas(int numLocalAndLiveReplicas) {
     }
   }
 
+  /**
+   * For get operations, if there are very few local replicas, then shuffle the pool with remote replicas so that the
+   * few (or one) remaining local replicas don't get all the traffic. The threshold for minimum number of local replicas
+   * is defined in {@link RouterConfig#routerGetOperationMinLocalReplicaCountToPrioritizeLocal}.
+   * @param numLocalAndLiveReplicas the number of local and live replicas.
+   * @param numRemoteOriginatingDcAndLiveReplicas the number of remote originating DC and live replicas.
+   */
+  void maybeShuffleWithRemoteReplicas(int numLocalAndLiveReplicas, int numRemoteOriginatingDcAndLiveReplicas) {
+    if (isGetOperation() && numLocalAndLiveReplicas > 0
+        && numLocalAndLiveReplicas <= routerConfig.routerGetOperationMinLocalReplicaCountToPrioritizeLocal) {
+      logger.debug("Shuffling replicas for {} because there are only {} local and live replicas for {}",
+          routerOperation, numLocalAndLiveReplicas, partitionId);
+      routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.inc();
+      List<ReplicaId> replicasToReshuffle = new ArrayList<>();
+      if (numRemoteOriginatingDcAndLiveReplicas > 0) {
+        // If the local DC is not the originating DC, we shuffle only with originating DC replicas.
+        replicasToReshuffle.addAll(
+            replicaPool.subList(0, numLocalAndLiveReplicas + numRemoteOriginatingDcAndLiveReplicas));
+      } else {
+        replicasToReshuffle.addAll(replicaPool);
+      }
+      Collections.shuffle(replicasToReshuffle);
+      ListIterator<ReplicaId> iter = replicaPool.listIterator();
+      for (ReplicaId replicaId : replicasToReshuffle) {
+        iter.next();
+        iter.set(replicaId);
+      }
+      if (!replicaPool.get(0).getDataNodeId().getDatacenterName().equals(datacenterName)) {
+        routerMetrics.remoteReplicaPrioritizedForGetDueToFewLocalReplicas.inc();
+      }
+    }
+  }
+
   public boolean hasFailed() {
     if (routerOperation == RouterOperation.PutOperation && routerConfig.routerPutUseDynamicSuccessTarget) {
       return totalReplicaCount - failedCount < Math.max(totalReplicaCount - 1,

Diff for: ambry-router/src/test/java/com/github/ambry/router/GetBlobOperationTest.java

@@ -436,6 +436,7 @@ public void testSimpleBlobGetSuccess() throws Exception {
       if (!quotaChargeCallback.getQuotaConfig().bandwidthThrottlingFeatureEnabled) {
         Assert.assertEquals(i + 1, quotaChargeCallback.numCheckAndChargeCalls);
       }
+      assertEquals(0, routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.getCount());
     }
   }
 
@@ -455,6 +456,7 @@ public void testSimpleBlobRawMode() throws Exception {
       GetBlobOperation op = createOperationAndComplete(null);
       Assert.assertEquals(IllegalStateException.class, op.getOperationException().getClass());
     }
+    assertEquals(0, routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.getCount());
   }
 
   /**
@@ -502,6 +504,7 @@ public void testCompositeBlobRawMode() throws Exception {
       } else {
         // Only supported for encrypted blobs now
       }
+      assertEquals(0, routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.getCount());
     }
   }
 
@@ -578,6 +581,7 @@ public void testCompositeBlobChunkSizeMultipleGetSuccess() throws Exception {
       }
       getAndAssertSuccess(false, false, lifeVersion);
     }
+    assertEquals(0, routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.getCount());
   }
 
   /**
@@ -897,6 +901,38 @@ public void testOrigDcUnavailability() throws Exception {
     getErrorCodeChecker.testAndAssert(RouterErrorCode.AmbryUnavailable);
   }
 
+  /**
+   * Test that if local DC has less than {@code RouterConfig#routerGetOperationMinLocalReplicaCountToPrioritizeLocal}
+   * replicas available, then we will prioritize remote DC.
+   * @throws Exception
+   */
+  @Test
+  public void testGetGoesToRemoteIfLocalDcHasFewAvailableReplicas() throws Exception {
+    blobSize = maxChunkSize;
+    doPut();
+    Properties props = getDefaultNonBlockingRouterProperties(true);
+    props.setProperty(RouterConfig.ROUTER_GET_ELIGIBLE_REPLICAS_BY_STATE_ENABLED, "true");
+    props.setProperty(RouterConfig.ROUTER_GET_OPERATION_MIN_LOCAL_REPLICA_COUNT_TO_PRIORITIZE_LOCAL, "1");
+    RouterConfig oldRouterConfig = routerConfig;
+    routerConfig = new RouterConfig(new VerifiableProperties(props));
+
+    try {
+      // For originating DC, let's bring down two replicas
+      MockPartitionId partitionId = (MockPartitionId) blobId.getPartition();
+      List<ReplicaId> localReplicas = partitionId.replicaIds.stream()
+          .filter(replicaId -> replicaId.getDataNodeId().getDatacenterName().equals(localDcName))
+          .collect(Collectors.toList());
+      partitionId.setReplicaState(localReplicas.get(0), ReplicaState.LEADER);
+      partitionId.setReplicaState(localReplicas.get(1), ReplicaState.OFFLINE);
+      partitionId.setReplicaState(localReplicas.get(2), ReplicaState.OFFLINE);
+
+      getAndAssertSuccess();
+      assertEquals(1, routerMetrics.shuffledWithRemoteReplicasForGetDueToFewLocalReplicas.getCount());
+    } finally {
+      routerConfig = oldRouterConfig;
+    }
+  }
+
   /**
    * Test the case where deletes and unavailability happens at the same time, delete should take priority.
    * @throws Exception