support for strong reads

Author: rahul2393

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/EwmaLatencyTracker.java

@@ -29,9 +29,8 @@
 /**
  * Implementation of {@link LatencyTracker} using Exponentially Weighted Moving Average (EWMA).
  *
- * <p>By default, this tracker uses a time-decayed EWMA:
- * $S_{i+1} = \alpha(\Delta t) * new\_latency + (1 - \alpha(\Delta t)) * S_i$, where $\alpha(\Delta
- * t) = 1 - e^{-\Delta t / \tau}$.
+ * <p>By default, this tracker uses a time-decayed EWMA: $S_{i+1} = \alpha(\Delta t) * new\_latency
+ * + (1 - \alpha(\Delta t)) * S_i$, where $\alpha(\Delta t) = 1 - e^{-\Delta t / \tau}$.
  *
  * <p>A fixed-alpha constructor is retained for focused tests.
  */

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java

@@ -44,6 +44,7 @@
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Predicate;
+import java.util.function.ToDoubleFunction;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.stream.IntStream;
@@ -57,6 +58,7 @@ public final class KeyRangeCache {
 
   private static final int MAX_LOCAL_REPLICA_DISTANCE = 5;
   private static final int DEFAULT_MIN_ENTRIES_FOR_RANDOM_PICK = 1000;
+  private static final double LOCAL_LEADER_SELECTION_COST_MULTIPLIER = 0.5D;
 
   /** Determines how to handle ranges that span multiple splits. */
   public enum RangeMode {
@@ -828,16 +830,20 @@ private TabletSnapshot selectTablet(
         List<SkippedTabletDetail> skippedTabletDetails,
         Map<String, ChannelEndpoint> resolvedEndpoints,
         SelectionStats selectionStats) {
-      if (!preferLeader) {
-        return selectLatencyAwareTablet(
+      if (!preferLeader || hintBuilder.getOperationUid() > 0L) {
+        TabletSnapshot preferredLeader =
+            preferLeader ? localLeaderForScoreBias(snapshot, hasDirectedReadOptions) : null;
+        return selectScoreAwareTablet(
             snapshot,
             directedReadOptions,
             hintBuilder,
             excludedEndpoints,
             skippedTabletUids,
             skippedTabletDetails,
             resolvedEndpoints,
-            selectionStats);
+            selectionStats,
+            preferredLeader,
+            preferLeader ? "leader_preferred_latency_score" : "latency_score");
       }
 
       boolean checkedLeader = false;
@@ -884,18 +890,32 @@ private TabletSnapshot selectTablet(
       return null;
     }
 
-    private TabletSnapshot selectLatencyAwareTablet(
+    @javax.annotation.Nullable
+    private TabletSnapshot localLeaderForScoreBias(
+        GroupSnapshot snapshot, boolean hasDirectedReadOptions) {
+      if (!hasDirectedReadOptions
+          && snapshot.hasLeader()
+          && snapshot.leader().distance <= MAX_LOCAL_REPLICA_DISTANCE) {
+        return snapshot.leader();
+      }
+      return null;
+    }
+
+    private TabletSnapshot selectScoreAwareTablet(
         GroupSnapshot snapshot,
         DirectedReadOptions directedReadOptions,
         RoutingHint.Builder hintBuilder,
         Predicate<String> excludedEndpoints,
         Set<Long> skippedTabletUids,
         List<SkippedTabletDetail> skippedTabletDetails,
         Map<String, ChannelEndpoint> resolvedEndpoints,
-        SelectionStats selectionStats) {
+        SelectionStats selectionStats,
+        @javax.annotation.Nullable TabletSnapshot preferredLeader,
+        String selectionReason) {
       long operationUid = hintBuilder.getOperationUid();
       List<TabletSnapshot> eligibleTablets = new ArrayList<>();
       List<ChannelEndpoint> eligibleEndpoints = new ArrayList<>();
+      Map<String, TabletSnapshot> endpointByAddress = new HashMap<>();
       int scoredCandidates = 0;
 
       for (TabletSnapshot tablet : snapshot.tablets) {
@@ -921,6 +941,7 @@ private TabletSnapshot selectLatencyAwareTablet(
         }
         eligibleTablets.add(tablet);
         eligibleEndpoints.add(endpoint);
+        endpointByAddress.put(endpoint.getAddress(), tablet);
         if (EndpointLatencyRegistry.hasScore(operationUid, tablet.serverAddress)) {
           scoredCandidates++;
         }
@@ -936,6 +957,7 @@ private TabletSnapshot selectLatencyAwareTablet(
                 snapshot,
                 eligibleTablets,
                 operationUid,
+                selectionCostLookup(operationUid, preferredLeader),
                 "single_candidate",
                 selected,
                 scoredCandidates);
@@ -946,16 +968,15 @@ private TabletSnapshot selectLatencyAwareTablet(
             eligibleTablets.stream()
                 .min(
                     Comparator.comparingDouble(
-                        tablet ->
-                            EndpointLatencyRegistry.getSelectionCost(
-                                operationUid, tablet.serverAddress)))
+                        tablet -> selectionCost(operationUid, tablet, preferredLeader)))
                 .orElse(eligibleTablets.get(0));
         selectionStats.selectionDetail =
             buildSelectionDetail(
                 snapshot,
                 eligibleTablets,
                 operationUid,
-                "latency_score",
+                selectionCostLookup(operationUid, preferredLeader),
+                selectionReason,
                 selected,
                 scoredCandidates);
         return selected;
@@ -965,15 +986,17 @@ private TabletSnapshot selectLatencyAwareTablet(
           replicaSelector.select(
               eligibleEndpoints,
               endpoint ->
-                  EndpointLatencyRegistry.getSelectionCost(operationUid, endpoint.getAddress()));
+                  selectionCost(
+                      operationUid, endpointByAddress.get(endpoint.getAddress()), preferredLeader));
       if (selectedEndpoint == null) {
         TabletSnapshot selected = eligibleTablets.get(0);
         selectionStats.selectionDetail =
             buildSelectionDetail(
                 snapshot,
                 eligibleTablets,
                 operationUid,
-                "latency_score",
+                selectionCostLookup(operationUid, preferredLeader),
+                selectionReason,
                 selected,
                 scoredCandidates);
         return selected;
@@ -986,7 +1009,8 @@ private TabletSnapshot selectLatencyAwareTablet(
                   snapshot,
                   eligibleTablets,
                   operationUid,
-                  "latency_score",
+                  selectionCostLookup(operationUid, preferredLeader),
+                  selectionReason,
                   selected,
                   scoredCandidates);
           return selected;
@@ -995,10 +1019,35 @@ private TabletSnapshot selectLatencyAwareTablet(
       TabletSnapshot selected = eligibleTablets.get(0);
       selectionStats.selectionDetail =
           buildSelectionDetail(
-              snapshot, eligibleTablets, operationUid, "latency_score", selected, scoredCandidates);
+              snapshot,
+              eligibleTablets,
+              operationUid,
+              selectionCostLookup(operationUid, preferredLeader),
+              selectionReason,
+              selected,
+              scoredCandidates);
       return selected;
     }
 
+    private ToDoubleFunction<TabletSnapshot> selectionCostLookup(
+        long operationUid, @javax.annotation.Nullable TabletSnapshot preferredLeader) {
+      return tablet -> selectionCost(operationUid, tablet, preferredLeader);
+    }
+
+    private double selectionCost(
+        long operationUid,
+        @javax.annotation.Nullable TabletSnapshot tablet,
+        @javax.annotation.Nullable TabletSnapshot preferredLeader) {
+      if (tablet == null) {
+        return Double.MAX_VALUE;
+      }
+      double cost = EndpointLatencyRegistry.getSelectionCost(operationUid, tablet.serverAddress);
+      if (preferredLeader != null && tablet == preferredLeader) {
+        return cost * LOCAL_LEADER_SELECTION_COST_MULTIPLIER;
+      }
+      return cost;
+    }
+
     @javax.annotation.Nullable
     private TabletSnapshot selectRandomExcludedOrCoolingDownTablet(
         GroupSnapshot snapshot,
@@ -1166,30 +1215,26 @@ private SelectionDetail buildSelectionDetail(
         GroupSnapshot snapshot,
         List<TabletSnapshot> eligibleTablets,
         long operationUid,
+        ToDoubleFunction<TabletSnapshot> selectionCostLookup,
         String selectionReason,
         TabletSnapshot selected,
         int scoredCandidates) {
       double bestScore = Double.MAX_VALUE;
       for (TabletSnapshot tablet : eligibleTablets) {
-        bestScore =
-            Math.min(
-                bestScore,
-                EndpointLatencyRegistry.getSelectionCost(operationUid, tablet.serverAddress));
+        bestScore = Math.min(bestScore, selectionCostLookup.applyAsDouble(tablet));
       }
 
       StringBuilder alternatives = new StringBuilder();
       int appended = 0;
-      double selectedScore =
-          EndpointLatencyRegistry.getSelectionCost(operationUid, selected.serverAddress);
+      double selectedScore = selectionCostLookup.applyAsDouble(selected);
       for (TabletSnapshot tablet : eligibleTablets) {
         if (tablet == selected || appended >= 4) {
           continue;
         }
         if (alternatives.length() > 0) {
           alternatives.append(", ");
         }
-        double candidateScore =
-            EndpointLatencyRegistry.getSelectionCost(operationUid, tablet.serverAddress);
+        double candidateScore = selectionCostLookup.applyAsDouble(tablet);
         alternatives
             .append(endpointLabel(snapshot, tablet))
             .append("=")

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java

@@ -656,7 +656,7 @@ public void preferLeaderFalseUsesLowestLatencyReplicaWhenScoresAvailable() {
   }
 
   @Test
-  public void preferLeaderTrueIgnoresLatencyScoresForLeaderSelection() {
+  public void preferLeaderTrueUsesLatencyScoresWhenOperationUidAvailable() {
     FakeEndpointCache endpointCache = new FakeEndpointCache();
     KeyRangeCache cache = new KeyRangeCache(endpointCache);
     cache.useDeterministicRandom();
@@ -679,6 +679,33 @@ public void preferLeaderTrueIgnoresLatencyScoresForLeaderSelection() {
             DirectedReadOptions.getDefaultInstance(),
             hint);
 
+    assertNotNull(server);
+    assertEquals("server2", server.getAddress());
+  }
+
+  @Test
+  public void preferLeaderTrueWithoutOperationUidKeepsLeaderSelection() {
+    FakeEndpointCache endpointCache = new FakeEndpointCache();
+    KeyRangeCache cache = new KeyRangeCache(endpointCache);
+    cache.useDeterministicRandom();
+    cache.addRanges(threeReplicaUpdate());
+
+    endpointCache.get("server1");
+    endpointCache.get("server2");
+    endpointCache.get("server3");
+
+    cache.recordReplicaLatency(TEST_OPERATION_UID, "server1", Duration.ofNanos(300_000L));
+    cache.recordReplicaLatency(TEST_OPERATION_UID, "server2", Duration.ofNanos(100_000L));
+    cache.recordReplicaLatency(TEST_OPERATION_UID, "server3", Duration.ofNanos(200_000L));
+
+    RoutingHint.Builder hint = RoutingHint.newBuilder().setKey(bytes("a"));
+    ChannelEndpoint server =
+        cache.fillRoutingHint(
+            true,
+            KeyRangeCache.RangeMode.COVERING_SPLIT,
+            DirectedReadOptions.getDefaultInstance(),
+            hint);
+
     assertNotNull(server);
     assertEquals("server1", server.getAddress());
   }