more fixes

Author: rahul2393

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

@@ -24,15 +24,12 @@
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.DirectedReadOptions;
 import com.google.spanner.v1.ExecuteSqlRequest;
-import com.google.spanner.v1.Group;
 import com.google.spanner.v1.Mutation;
 import com.google.spanner.v1.ReadRequest;
 import com.google.spanner.v1.RoutingHint;
-import com.google.spanner.v1.Tablet;
 import com.google.spanner.v1.TransactionOptions;
 import com.google.spanner.v1.TransactionSelector;
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -55,6 +52,7 @@
 @InternalApi
 public final class ChannelFinder {
   private static final Predicate<String> NO_EXCLUDED_ENDPOINTS = address -> false;
+  private static final int CACHE_UPDATE_DRAIN_BATCH_SIZE = 64;
   private static final int MAX_CACHE_UPDATE_THREADS =
       Math.max(2, Runtime.getRuntime().availableProcessors());
   private static final ExecutorService CACHE_UPDATE_POOL = createCacheUpdatePool();
@@ -132,39 +130,12 @@ private boolean shouldProcessUpdate(CacheUpdate update) {
   }
 
   public void update(CacheUpdate update) {
+    Set<String> currentAddresses;
     synchronized (updateLock) {
-      long currentId = databaseId.get();
-      long updateDatabaseId = update.getDatabaseId();
-      if (updateDatabaseId != 0 && currentId != updateDatabaseId) {
-        if (currentId != 0) {
-          recipeCache.clear();
-          rangeCache.clear();
-        }
-        databaseId.set(updateDatabaseId);
-      }
-      if (update.hasKeyRecipes()) {
-        recipeCache.addRecipes(update.getKeyRecipes());
-      }
-      rangeCache.addRanges(update);
-
-      // Notify the lifecycle manager about server addresses so it can create endpoints
-      // in the background and start probing, and evict stale endpoints atomically.
-      if (lifecycleManager != null && finderKey != null) {
-        Set<String> currentAddresses = new HashSet<>();
-        for (Group group : update.getGroupList()) {
-          for (Tablet tablet : group.getTabletsList()) {
-            String addr = tablet.getServerAddress();
-            if (!addr.isEmpty()) {
-              currentAddresses.add(addr);
-            }
-          }
-        }
-        // Also include addresses from existing cached tablets not in this update.
-        currentAddresses.addAll(rangeCache.getActiveAddresses());
-        // Atomically ensure endpoints exist and evict stale ones.
-        lifecycleManager.updateActiveAddresses(finderKey, currentAddresses);
-      }
+      applyUpdateLocked(update);
+      currentAddresses = snapshotActiveAddressesLocked();
     }
+    publishLifecycleUpdate(currentAddresses);
   }
 
   public void updateAsync(CacheUpdate update) {
@@ -187,10 +158,12 @@ public void updateAsync(CacheUpdate update) {
   }
 
   private void drainPendingUpdate() {
+    List<PendingCacheUpdate> batch = new ArrayList<>(CACHE_UPDATE_DRAIN_BATCH_SIZE);
     while (true) {
-      PendingCacheUpdate toApply;
-      while ((toApply = pendingUpdates.poll()) != null) {
-        update(toApply.update);
+      drainBatch(batch);
+      if (!batch.isEmpty()) {
+        applyBatch(batch);
+        batch.clear();
       }
       drainScheduled.set(false);
       if (pendingUpdates.isEmpty() || !drainScheduled.compareAndSet(false, true)) {
@@ -199,6 +172,56 @@ private void drainPendingUpdate() {
     }
   }
 
+  private void drainBatch(List<PendingCacheUpdate> batch) {
+    PendingCacheUpdate toApply;
+    while (batch.size() < CACHE_UPDATE_DRAIN_BATCH_SIZE
+        && (toApply = pendingUpdates.poll()) != null) {
+      batch.add(toApply);
+    }
+  }
+
+  private void applyBatch(List<PendingCacheUpdate> batch) {
+    Set<String> currentAddresses;
+    synchronized (updateLock) {
+      for (PendingCacheUpdate pendingUpdate : batch) {
+        applyUpdateLocked(pendingUpdate.update);
+      }
+      currentAddresses = snapshotActiveAddressesLocked();
+    }
+    publishLifecycleUpdate(currentAddresses);
+  }
+
+  private void applyUpdateLocked(CacheUpdate update) {
+    long currentId = databaseId.get();
+    long updateDatabaseId = update.getDatabaseId();
+    if (updateDatabaseId != 0 && currentId != updateDatabaseId) {
+      if (currentId != 0) {
+        recipeCache.clear();
+        rangeCache.clear();
+      }
+      databaseId.set(updateDatabaseId);
+    }
+    if (update.hasKeyRecipes()) {
+      recipeCache.addRecipes(update.getKeyRecipes());
+    }
+    rangeCache.addRanges(update);
+  }
+
+  @Nullable
+  private Set<String> snapshotActiveAddressesLocked() {
+    if (lifecycleManager == null || finderKey == null) {
+      return null;
+    }
+    return rangeCache.getActiveAddresses();
+  }
+
+  private void publishLifecycleUpdate(@Nullable Set<String> currentAddresses) {
+    if (currentAddresses == null) {
+      return;
+    }
+    lifecycleManager.updateActiveAddressesAsync(finderKey, currentAddresses);
+  }
+
   /**
    * Test-only hook used by {@link KeyAwareChannel#awaitPendingCacheUpdates()} to wait until the
    * async cache update worker has finished applying the latest pending update.

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

@@ -29,6 +29,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -103,6 +105,11 @@ static final class EndpointState {
   private final ChannelEndpointCache endpointCache;
   private final Map<String, EndpointState> endpoints = new ConcurrentHashMap<>();
   private final Set<String> transientFailureEvictedAddresses = ConcurrentHashMap.newKeySet();
+  private final Map<String, Long> finderGenerations = new ConcurrentHashMap<>();
+  private final Map<String, PendingActiveAddressUpdate> pendingActiveAddressUpdates =
+      new ConcurrentHashMap<>();
+  private final Set<String> queuedFinderKeys = ConcurrentHashMap.newKeySet();
+  private final ConcurrentLinkedQueue<String> queuedFinders = new ConcurrentLinkedQueue<>();
 
   /**
    * Active addresses reported by each ChannelFinder, keyed by database id.
@@ -118,15 +125,27 @@ static final class EndpointState {
 
   private final Object activeAddressLock = new Object();
 
+  private final ExecutorService activeAddressUpdateExecutor;
   private final ScheduledExecutorService scheduler;
   private final AtomicBoolean isShutdown = new AtomicBoolean(false);
+  private final AtomicBoolean activeAddressDrainScheduled = new AtomicBoolean(false);
   private final long probeIntervalSeconds;
   private final Duration idleEvictionDuration;
   private final Clock clock;
   private final String defaultEndpointAddress;
 
   private ScheduledFuture<?> evictionFuture;
 
+  private static final class PendingActiveAddressUpdate {
+    private final Set<String> activeAddresses;
+    private final long generation;
+
+    private PendingActiveAddressUpdate(Set<String> activeAddresses, long generation) {
+      this.activeAddresses = activeAddresses;
+      this.generation = generation;
+    }
+  }
+
   EndpointLifecycleManager(ChannelEndpointCache endpointCache) {
     this(
         endpointCache,
@@ -146,6 +165,13 @@ static final class EndpointState {
     this.idleEvictionDuration = idleEvictionDuration;
     this.clock = clock;
     this.defaultEndpointAddress = endpointCache.defaultChannel().getAddress();
+    this.activeAddressUpdateExecutor =
+        Executors.newSingleThreadExecutor(
+            r -> {
+              Thread t = new Thread(r, "spanner-active-address-update");
+              t.setDaemon(true);
+              return t;
+            });
     this.scheduler =
         Executors.newScheduledThreadPool(
             2,
@@ -213,6 +239,59 @@ private void clearTransientFailureEvictionMarker(String address) {
     }
   }
 
+  /**
+   * Enqueues active-address reconciliation on a dedicated worker so cache-map updates do not block
+   * on endpoint lifecycle bookkeeping.
+   */
+  void updateActiveAddressesAsync(String finderKey, Set<String> activeAddresses) {
+    if (isShutdown.get() || finderKey == null || finderKey.isEmpty()) {
+      return;
+    }
+    synchronized (activeAddressLock) {
+      long generation = finderGenerations.getOrDefault(finderKey, 0L);
+      pendingActiveAddressUpdates.put(
+          finderKey, new PendingActiveAddressUpdate(new HashSet<>(activeAddresses), generation));
+      if (queuedFinderKeys.add(finderKey)) {
+        queuedFinders.add(finderKey);
+      }
+    }
+    scheduleActiveAddressDrain();
+  }
+
+  private void scheduleActiveAddressDrain() {
+    if (!activeAddressDrainScheduled.compareAndSet(false, true)) {
+      return;
+    }
+    activeAddressUpdateExecutor.execute(this::drainPendingActiveAddressUpdates);
+  }
+
+  private void drainPendingActiveAddressUpdates() {
+    while (true) {
+      String finderKey = queuedFinders.poll();
+      if (finderKey == null) {
+        activeAddressDrainScheduled.set(false);
+        if (queuedFinders.isEmpty() || !activeAddressDrainScheduled.compareAndSet(false, true)) {
+          return;
+        }
+        continue;
+      }
+
+      queuedFinderKeys.remove(finderKey);
+      PendingActiveAddressUpdate pendingUpdate = pendingActiveAddressUpdates.remove(finderKey);
+      if (pendingUpdate == null) {
+        continue;
+      }
+
+      synchronized (activeAddressLock) {
+        long currentGeneration = finderGenerations.getOrDefault(finderKey, 0L);
+        if (currentGeneration != pendingUpdate.generation) {
+          continue;
+        }
+      }
+      updateActiveAddresses(finderKey, pendingUpdate.activeAddresses);
+    }
+  }
+
   /**
    * Records that real (non-probe) traffic was routed to an endpoint. This refreshes the idle
    * eviction timer for this endpoint.
@@ -295,6 +374,9 @@ void unregisterFinder(String finderKey) {
       return;
     }
     synchronized (activeAddressLock) {
+      finderGenerations.merge(finderKey, 1L, Long::sum);
+      pendingActiveAddressUpdates.remove(finderKey);
+      queuedFinderKeys.remove(finderKey);
       if (activeAddressesPerFinder.remove(finderKey) == null) {
         return;
       }
@@ -588,6 +670,7 @@ void shutdown() {
     }
 
     logger.log(Level.FINE, "Shutting down endpoint lifecycle manager");
+    activeAddressUpdateExecutor.shutdownNow();
 
     if (evictionFuture != null) {
       evictionFuture.cancel(false);
@@ -602,6 +685,9 @@ void shutdown() {
     }
     endpoints.clear();
     transientFailureEvictedAddresses.clear();
+    pendingActiveAddressUpdates.clear();
+    queuedFinderKeys.clear();
+    queuedFinders.clear();
 
     scheduler.shutdown();
     try {

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

@@ -32,6 +32,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import org.junit.Test;
@@ -112,6 +114,28 @@ public void updateAsyncProcessesDatabaseTransitionWithoutRangesOrGroups() throws
     assertThat(rangeCache(finder).size()).isEqualTo(0);
   }
 
+  @Test
+  public void updateDoesNotBlockOnLifecycleManagerAddressReconciliation() throws Exception {
+    BlockingLifecycleManager lifecycleManager =
+        new BlockingLifecycleManager(new FakeEndpointCache());
+    ChannelFinder finder = new ChannelFinder(new FakeEndpointCache(), lifecycleManager, "db-1");
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      Future<?> future = executor.submit(() -> finder.update(singleRangeUpdate(0)));
+
+      future.get(1, TimeUnit.SECONDS);
+      assertThat(lifecycleManager.updateStarted.await(5, TimeUnit.SECONDS)).isTrue();
+      assertThat(rangeCache(finder).size()).isEqualTo(1);
+
+      lifecycleManager.releaseUpdate.countDown();
+      assertThat(lifecycleManager.updateFinished.await(5, TimeUnit.SECONDS)).isTrue();
+    } finally {
+      lifecycleManager.releaseUpdate.countDown();
+      lifecycleManager.shutdown();
+      executor.shutdownNow();
+    }
+  }
+
   private static CacheUpdate singleRangeUpdate(int index) {
     String startKey = String.format("k%05d", index);
     String limitKey = String.format("k%05d", index + 1);
@@ -262,4 +286,26 @@ public String authority() {
       return "fake";
     }
   }
+
+  private static final class BlockingLifecycleManager extends EndpointLifecycleManager {
+    private final CountDownLatch updateStarted = new CountDownLatch(1);
+    private final CountDownLatch releaseUpdate = new CountDownLatch(1);
+    private final CountDownLatch updateFinished = new CountDownLatch(1);
+
+    private BlockingLifecycleManager(ChannelEndpointCache endpointCache) {
+      super(endpointCache);
+    }
+
+    @Override
+    void updateActiveAddresses(String finderKey, java.util.Set<String> activeAddresses) {
+      updateStarted.countDown();
+      try {
+        releaseUpdate.await(5, TimeUnit.SECONDS);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      } finally {
+        updateFinished.countDown();
+      }
+    }
+  }
 }