add more test for R/W shape

Author: rahul2393

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

@@ -82,7 +82,8 @@ public ChannelFinder(
       ChannelEndpointCache endpointCache,
       @Nullable EndpointLifecycleManager lifecycleManager,
       @Nullable String finderKey) {
-    this.rangeCache = new KeyRangeCache(Objects.requireNonNull(endpointCache), lifecycleManager);
+    this.rangeCache =
+        new KeyRangeCache(Objects.requireNonNull(endpointCache), lifecycleManager, finderKey);
     this.lifecycleManager = lifecycleManager;
     this.finderKey = finderKey;
   }
@@ -91,6 +92,11 @@ void useDeterministicRandom() {
     rangeCache.useDeterministicRandom();
   }
 
+  @Nullable
+  String finderKey() {
+    return finderKey;
+  }
+
   private static ExecutorService createCacheUpdatePool() {
     ThreadPoolExecutor executor =
         new ThreadPoolExecutor(

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

@@ -17,37 +17,51 @@
 package com.google.cloud.spanner.spi.v1;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Ticker;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /** Shared process-local latency scores for routed Spanner endpoints. */
 final class EndpointLatencyRegistry {
+  private static final String GLOBAL_SCOPE = "__global__";
 
   static final Duration DEFAULT_ERROR_PENALTY = Duration.ofSeconds(10);
   static final Duration DEFAULT_RTT = Duration.ofMillis(10);
   static final double DEFAULT_PENALTY_VALUE = 1_000_000.0;
+  @VisibleForTesting static final Duration TRACKER_EXPIRE_AFTER_ACCESS = Duration.ofMinutes(10);
+  @VisibleForTesting static final long MAX_TRACKERS = 100_000L;
 
-  private static final ConcurrentHashMap<TrackerKey, LatencyTracker> TRACKERS =
-      new ConcurrentHashMap<>();
+  private static volatile Cache<TrackerKey, LatencyTracker> TRACKERS =
+      newTrackerCache(Ticker.systemTicker());
   private static final ConcurrentHashMap<String, AtomicInteger> INFLIGHT_REQUESTS =
       new ConcurrentHashMap<>();
 
   private EndpointLatencyRegistry() {}
 
-  static boolean hasScore(long operationUid, String endpointLabelOrAddress) {
-    TrackerKey trackerKey = trackerKey(operationUid, endpointLabelOrAddress);
-    return trackerKey != null && TRACKERS.containsKey(trackerKey);
+  static boolean hasScore(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress) {
+    TrackerKey trackerKey = trackerKey(databaseScope, operationUid, endpointLabelOrAddress);
+    return trackerKey != null && TRACKERS.getIfPresent(trackerKey) != null;
   }
 
-  static double getSelectionCost(long operationUid, String endpointLabelOrAddress) {
-    TrackerKey trackerKey = trackerKey(operationUid, endpointLabelOrAddress);
+  static double getSelectionCost(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress) {
+    TrackerKey trackerKey = trackerKey(databaseScope, operationUid, endpointLabelOrAddress);
     if (trackerKey == null) {
       return Double.MAX_VALUE;
     }
     double activeRequests = getInflight(endpointLabelOrAddress);
-    LatencyTracker tracker = TRACKERS.get(trackerKey);
+    LatencyTracker tracker = TRACKERS.getIfPresent(trackerKey);
     if (tracker != null) {
       return tracker.getScore() * (activeRequests + 1.0);
     }
@@ -57,24 +71,35 @@ static double getSelectionCost(long operationUid, String endpointLabelOrAddress)
     return defaultRttMicros() * (activeRequests + 1.0);
   }
 
-  static void recordLatency(long operationUid, String endpointLabelOrAddress, Duration latency) {
-    TrackerKey trackerKey = trackerKey(operationUid, endpointLabelOrAddress);
+  static void recordLatency(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress,
+      Duration latency) {
+    TrackerKey trackerKey = trackerKey(databaseScope, operationUid, endpointLabelOrAddress);
     if (trackerKey == null || latency == null) {
       return;
     }
-    TRACKERS.computeIfAbsent(trackerKey, ignored -> new EwmaLatencyTracker()).update(latency);
+    getOrCreateTracker(trackerKey).update(latency);
   }
 
-  static void recordError(long operationUid, String endpointLabelOrAddress) {
-    recordError(operationUid, endpointLabelOrAddress, DEFAULT_ERROR_PENALTY);
+  static void recordError(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress) {
+    recordError(databaseScope, operationUid, endpointLabelOrAddress, DEFAULT_ERROR_PENALTY);
   }
 
-  static void recordError(long operationUid, String endpointLabelOrAddress, Duration penalty) {
-    TrackerKey trackerKey = trackerKey(operationUid, endpointLabelOrAddress);
+  static void recordError(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress,
+      Duration penalty) {
+    TrackerKey trackerKey = trackerKey(databaseScope, operationUid, endpointLabelOrAddress);
     if (trackerKey == null || penalty == null) {
       return;
     }
-    TRACKERS.computeIfAbsent(trackerKey, ignored -> new EwmaLatencyTracker()).recordError(penalty);
+    getOrCreateTracker(trackerKey).recordError(penalty);
   }
 
   static void beginRequest(String endpointLabelOrAddress) {
@@ -94,10 +119,7 @@ static void finishRequest(String endpointLabelOrAddress) {
     if (counter == null) {
       return;
     }
-    int updated = counter.decrementAndGet();
-    if (updated <= 0) {
-      INFLIGHT_REQUESTS.remove(address, counter);
-    }
+    counter.updateAndGet(current -> current > 0 ? current - 1 : 0);
   }
 
   static int getInflight(String endpointLabelOrAddress) {
@@ -111,10 +133,15 @@ static int getInflight(String endpointLabelOrAddress) {
 
   @VisibleForTesting
   static void clear() {
-    TRACKERS.clear();
+    TRACKERS.invalidateAll();
     INFLIGHT_REQUESTS.clear();
   }
 
+  @VisibleForTesting
+  static void useTrackerTicker(Ticker ticker) {
+    TRACKERS = newTrackerCache(ticker);
+  }
+
   @VisibleForTesting
   static String normalizeAddress(String endpointLabelOrAddress) {
     if (endpointLabelOrAddress == null || endpointLabelOrAddress.isEmpty()) {
@@ -124,24 +151,49 @@ static String normalizeAddress(String endpointLabelOrAddress) {
   }
 
   @VisibleForTesting
-  static TrackerKey trackerKey(long operationUid, String endpointLabelOrAddress) {
+  static TrackerKey trackerKey(
+      @javax.annotation.Nullable String databaseScope,
+      long operationUid,
+      String endpointLabelOrAddress) {
     String address = normalizeAddress(endpointLabelOrAddress);
     if (operationUid <= 0 || address == null) {
       return null;
     }
-    return new TrackerKey(operationUid, address);
+    return new TrackerKey(normalizeScope(databaseScope), operationUid, address);
   }
 
   private static long defaultRttMicros() {
     return DEFAULT_RTT.toNanos() / 1_000L;
   }
 
+  private static String normalizeScope(@javax.annotation.Nullable String databaseScope) {
+    return (databaseScope == null || databaseScope.isEmpty()) ? GLOBAL_SCOPE : databaseScope;
+  }
+
+  private static LatencyTracker getOrCreateTracker(TrackerKey trackerKey) {
+    try {
+      return TRACKERS.get(trackerKey, EwmaLatencyTracker::new);
+    } catch (ExecutionException e) {
+      throw new IllegalStateException("Failed to create latency tracker", e);
+    }
+  }
+
+  private static Cache<TrackerKey, LatencyTracker> newTrackerCache(Ticker ticker) {
+    return CacheBuilder.newBuilder()
+        .maximumSize(MAX_TRACKERS)
+        .expireAfterAccess(TRACKER_EXPIRE_AFTER_ACCESS.toNanos(), TimeUnit.NANOSECONDS)
+        .ticker(ticker)
+        .build();
+  }
+
   @VisibleForTesting
   static final class TrackerKey {
+    private final String databaseScope;
     private final long operationUid;
     private final String address;
 
-    private TrackerKey(long operationUid, String address) {
+    private TrackerKey(String databaseScope, long operationUid, String address) {
+      this.databaseScope = databaseScope;
       this.operationUid = operationUid;
       this.address = address;
     }
@@ -155,17 +207,19 @@ public boolean equals(Object other) {
         return false;
       }
       TrackerKey that = (TrackerKey) other;
-      return operationUid == that.operationUid && Objects.equals(address, that.address);
+      return operationUid == that.operationUid
+          && Objects.equals(databaseScope, that.databaseScope)
+          && Objects.equals(address, that.address);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(operationUid, address);
+      return Objects.hash(databaseScope, operationUid, address);
     }
 
     @Override
     public String toString() {
-      return operationUid + "@" + address;
+      return databaseScope + ":" + operationUid + "@" + address;
     }
   }
 }

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

@@ -232,7 +232,10 @@ private void recordFirstResponseLatency(
     }
     long latencyNanos = Math.max(0L, System.nanoTime() - startedAtNanos);
     EndpointLatencyRegistry.recordLatency(
-        routingTarget.operationUid, routingTarget.targetEndpoint, Duration.ofNanos(latencyNanos));
+        routingTarget.databaseScope,
+        routingTarget.operationUid,
+        routingTarget.targetEndpoint,
+        Duration.ofNanos(latencyNanos));
   }
 
   private Map<String, Float> parseServerTimingHeader(String serverTiming) {