diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java index f4fdb41fe3c5..35d42ac029e2 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java @@ -32,7 +32,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Objects; -import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedBlockingQueue; @@ -68,6 +67,7 @@ public final class ChannelFinder { new java.util.concurrent.CountDownLatch(0); @Nullable private final EndpointLifecycleManager lifecycleManager; @Nullable private final String finderKey; + private volatile boolean stale = false; public ChannelFinder(ChannelEndpointCache endpointCache) { this(endpointCache, null, null); @@ -97,6 +97,22 @@ String finderKey() { return finderKey; } + /** + * Marks this finder as stale after it has been evicted from the per-database finder cache. A + * stale finder stops processing cache updates and stops publishing active addresses, so an + * in-flight call that still references it cannot re-register endpoints with the lifecycle manager + * after {@link EndpointLifecycleManager#unregisterFinder} has run. + * + *

Synchronizing on {@code updateLock} guarantees that a concurrent publish either completes + * before this method returns — and is then superseded by the finder-generation bump in the + * subsequent {@code unregisterFinder} call — or observes the stale flag and skips publishing. + */ + void markStale() { + synchronized (updateLock) { + stale = true; + } + } + private static ExecutorService createCacheUpdatePool() { ThreadPoolExecutor executor = new ThreadPoolExecutor( @@ -136,16 +152,14 @@ private boolean shouldProcessUpdate(CacheUpdate update) { } public void update(CacheUpdate update) { - Set currentAddresses; synchronized (updateLock) { applyUpdateLocked(update); - currentAddresses = snapshotActiveAddressesLocked(); + publishLifecycleUpdateLocked(); } - publishLifecycleUpdate(currentAddresses); } public void updateAsync(CacheUpdate update) { - if (!shouldProcessUpdate(update)) { + if (stale || !shouldProcessUpdate(update)) { return; } pendingUpdates.add(new PendingCacheUpdate(update)); @@ -187,14 +201,12 @@ private void drainBatch(List batch) { } private void applyBatch(List batch) { - Set currentAddresses; synchronized (updateLock) { for (PendingCacheUpdate pendingUpdate : batch) { applyUpdateLocked(pendingUpdate.update); } - currentAddresses = snapshotActiveAddressesLocked(); + publishLifecycleUpdateLocked(); } - publishLifecycleUpdate(currentAddresses); } private void applyUpdateLocked(CacheUpdate update) { @@ -213,19 +225,16 @@ private void applyUpdateLocked(CacheUpdate update) { rangeCache.addRanges(update); } - @Nullable - private Set snapshotActiveAddressesLocked() { - if (lifecycleManager == null || finderKey == null) { - return null; - } - return rangeCache.getActiveAddresses(); - } - - private void publishLifecycleUpdate(@Nullable Set currentAddresses) { - if (currentAddresses == null) { + /** + * Publishes the active addresses to the lifecycle manager. Must be called while holding {@code + * updateLock} so the publish is ordered with {@link #markStale()}; the call itself only enqueues + * work for the lifecycle manager's reconciliation worker. + */ + private void publishLifecycleUpdateLocked() { + if (stale || lifecycleManager == null || finderKey == null) { return; } - lifecycleManager.updateActiveAddressesAsync(finderKey, currentAddresses); + lifecycleManager.updateActiveAddressesAsync(finderKey, rangeCache.getActiveAddresses()); } /** diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java index c0dfdb8210fe..cfaa72eacbec 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java @@ -25,6 +25,7 @@ import com.google.common.base.Ticker; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; +import com.google.common.cache.RemovalNotification; import com.google.protobuf.ByteString; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; @@ -46,12 +47,8 @@ import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import java.io.IOException; -import java.lang.ref.ReferenceQueue; -import java.lang.ref.SoftReference; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Predicate; import java.util.logging.Level; import java.util.logging.Logger; @@ -72,7 +69,9 @@ final class KeyAwareChannel extends ManagedChannel { private static final long MAX_TRACKED_TRANSACTION_AFFINITIES = 100_000L; private static final long TRANSACTION_AFFINITY_TTL_MINUTES = 10L; - private static final int CHANNEL_FINDER_CLEANUP_INTERVAL = 1024; + + @VisibleForTesting static final long MAX_CHANNEL_FINDERS = 100L; + @VisibleForTesting static final long CHANNEL_FINDER_TTL_MINUTES = 30L; private static final String STREAMING_READ_METHOD = "google.spanner.v1.Spanner/StreamingRead"; private static final String STREAMING_SQL_METHOD = "google.spanner.v1.Spanner/ExecuteStreamingSql"; @@ -87,9 +86,10 @@ final class KeyAwareChannel extends ManagedChannel { @Nullable private final EndpointLifecycleManager lifecycleManager; private final String authority; private final String defaultEndpointAddress; - private final ReferenceQueue channelFinderReferenceQueue = new ReferenceQueue<>(); - private final Map channelFinders = new ConcurrentHashMap<>(); - private final AtomicInteger channelFinderCleanupCounter = new AtomicInteger(); + // Per-database ChannelFinders, each owning bounded range and recipe caches. Bounded and aged out + // so that finders for databases that are no longer accessed do not accumulate for the lifetime + // of the channel; evicted finders are marked stale and unregistered from the lifecycle manager. + private final Cache channelFinders; // Maps read-write transaction IDs to their last routed endpoint. // Bound and age out entries in case application code abandons a transaction // without sending Commit/Rollback or otherwise clearing affinity. @@ -114,7 +114,7 @@ private KeyAwareChannel( @Nullable ChannelEndpointCacheFactory endpointCacheFactory, @Nullable GrpcGcpEndpointChannelConfigurator endpointChannelConfigurator, EndpointOverloadCooldownTracker endpointOverloadCooldowns, - Ticker transactionAffinityTicker) + Ticker cacheTicker) throws IOException { if (endpointCacheFactory == null) { this.endpointCache = @@ -131,7 +131,8 @@ private KeyAwareChannel( this.lifecycleManager = (endpointCacheFactory == null) ? new EndpointLifecycleManager(endpointCache) : null; this.endpointOverloadCooldowns = endpointOverloadCooldowns; - this.transactionAffinities = newTransactionAffinities(transactionAffinityTicker); + this.transactionAffinities = newTransactionAffinities(cacheTicker); + this.channelFinders = newChannelFinders(cacheTicker); } static KeyAwareChannel create( @@ -164,14 +165,10 @@ static KeyAwareChannel create( InstantiatingGrpcChannelProvider channelProvider, @Nullable ChannelEndpointCacheFactory endpointCacheFactory, EndpointOverloadCooldownTracker endpointOverloadCooldowns, - Ticker transactionAffinityTicker) + Ticker cacheTicker) throws IOException { return new KeyAwareChannel( - channelProvider, - endpointCacheFactory, - null, - endpointOverloadCooldowns, - transactionAffinityTicker); + channelProvider, endpointCacheFactory, null, endpointOverloadCooldowns, cacheTicker); } private static Cache newTransactionAffinities(Ticker ticker) { @@ -182,16 +179,22 @@ private static Cache newTransactionAffinities(Ticker ticker) .build(); } - private static final class ChannelFinderReference extends SoftReference { - final String databaseId; - - ChannelFinderReference( - String databaseId, - ChannelFinder referent, - ReferenceQueue referenceQueue) { - super(referent, referenceQueue); - this.databaseId = databaseId; - } + private Cache newChannelFinders(Ticker ticker) { + return CacheBuilder.newBuilder() + .maximumSize(MAX_CHANNEL_FINDERS) + .expireAfterAccess(CHANNEL_FINDER_TTL_MINUTES, TimeUnit.MINUTES) + .ticker(ticker) + .removalListener( + (RemovalNotification notification) -> { + ChannelFinder finder = notification.getValue(); + if (finder != null) { + finder.markStale(); + } + if (lifecycleManager != null) { + lifecycleManager.unregisterFinder(notification.getKey()); + } + }) + .build(); } private String extractDatabaseIdFromSession(String session) { @@ -205,48 +208,20 @@ private String extractDatabaseIdFromSession(String session) { return session.substring(0, sessionsIndex); } - private void cleanupStaleChannelFinders() { - ChannelFinderReference reference; - while ((reference = (ChannelFinderReference) channelFinderReferenceQueue.poll()) != null) { - if (channelFinders.remove(reference.databaseId, reference) && lifecycleManager != null) { - lifecycleManager.unregisterFinder(reference.databaseId); - } - } - } - - private void maybeCleanupStaleChannelFinders() { - if ((channelFinderCleanupCounter.incrementAndGet() & (CHANNEL_FINDER_CLEANUP_INTERVAL - 1)) - == 0) { - cleanupStaleChannelFinders(); - } - } - private ChannelFinder getOrCreateChannelFinder(String databaseId) { - maybeCleanupStaleChannelFinders(); - ChannelFinderReference ref = channelFinders.get(databaseId); - ChannelFinder finder = (ref != null) ? ref.get() : null; - if (finder == null) { - synchronized (channelFinders) { - ref = channelFinders.get(databaseId); - finder = (ref != null) ? ref.get() : null; - if (finder == null) { - finder = new ChannelFinder(endpointCache, lifecycleManager, databaseId); - channelFinders.put( - databaseId, - new ChannelFinderReference(databaseId, finder, channelFinderReferenceQueue)); - } - } + try { + return channelFinders.get( + databaseId, () -> new ChannelFinder(endpointCache, lifecycleManager, databaseId)); + } catch (ExecutionException executionException) { + // The loader only constructs a ChannelFinder and cannot throw a checked exception. + throw new IllegalStateException(executionException.getCause()); } - return finder; } @com.google.common.annotations.VisibleForTesting void awaitPendingCacheUpdates() throws InterruptedException { - for (ChannelFinderReference ref : channelFinders.values()) { - ChannelFinder finder = ref.get(); - if (finder != null) { - finder.awaitPendingUpdates(); - } + for (ChannelFinder finder : channelFinders.asMap().values()) { + finder.awaitPendingUpdates(); } } @@ -262,7 +237,7 @@ private void onRequestRouted(@Nullable ChannelEndpoint selectedEndpoint) { @Override public ManagedChannel shutdown() { - cleanupStaleChannelFinders(); + channelFinders.invalidateAll(); if (lifecycleManager != null) { lifecycleManager.shutdown(); } @@ -272,7 +247,7 @@ public ManagedChannel shutdown() { @Override public ManagedChannel shutdownNow() { - cleanupStaleChannelFinders(); + channelFinders.invalidateAll(); if (lifecycleManager != null) { lifecycleManager.shutdown(); } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java index 98c26381285d..ed67fabaf04f 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java @@ -60,6 +60,8 @@ public final class KeyRangeCache { private static final int DEFAULT_MIN_ENTRIES_FOR_RANDOM_PICK = 1000; private static final double LOCAL_LEADER_SELECTION_COST_MULTIPLIER = 0.5D; + @VisibleForTesting static final int DEFAULT_MAX_RANGES = 10_000; + /** Determines how to handle ranges that span multiple splits. */ public enum RangeMode { /** Consider it a cache miss if the whole range is not in a single split. */ @@ -160,6 +162,7 @@ static String formatTargetEndpointLabel(String address, boolean isLeader) { private volatile boolean deterministicRandom = false; private volatile int minCacheEntriesForRandomPick = DEFAULT_MIN_ENTRIES_FOR_RANDOM_PICK; + private volatile int maxRanges = DEFAULT_MAX_RANGES; public KeyRangeCache(ChannelEndpointCache endpointCache) { this(endpointCache, null, null); @@ -190,6 +193,11 @@ void setMinCacheEntriesForRandomPick(int value) { minCacheEntriesForRandomPick = value; } + @VisibleForTesting + void setMaxRanges(int value) { + maxRanges = value; + } + @VisibleForTesting void recordReplicaLatency(long operationUid, String address, Duration latency) { EndpointLatencyRegistry.recordLatency(databaseScope, operationUid, false, address, latency); @@ -235,6 +243,18 @@ public void addRanges(CacheUpdate cacheUpdate) { writeLock.unlock(); } } + maybeShrink(); + } + + /** + * Evicts least-recently-accessed ranges when the cache has grown beyond {@link #maxRanges}. + * Shrinks below the limit so the O(size) eviction cost is amortized over many insertions. + */ + private void maybeShrink() { + int max = maxRanges; + if (size() > max) { + shrinkTo(max - max / 10); + } } /** diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelFinderTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelFinderTest.java index ee388677f416..2d34aefcd562 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelFinderTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelFinderTest.java @@ -137,6 +137,34 @@ public void updateDoesNotBlockOnLifecycleManagerAddressReconciliation() throws E } } + @Test + public void markStaleDropsAsyncUpdates() throws Exception { + ChannelFinder finder = new ChannelFinder(new FakeEndpointCache()); + + finder.markStale(); + finder.updateAsync(singleRangeUpdate(0)); + finder.awaitPendingUpdates(); + + assertThat(rangeCache(finder).size()).isEqualTo(0); + } + + @Test + public void markStaleStopsPublishingActiveAddressesToLifecycleManager() throws Exception { + RecordingLifecycleManager lifecycleManager = + new RecordingLifecycleManager(new FakeEndpointCache()); + ChannelFinder finder = new ChannelFinder(new FakeEndpointCache(), lifecycleManager, "db-1"); + try { + finder.update(singleRangeUpdate(0)); + assertThat(lifecycleManager.publishCount.get()).isEqualTo(1); + + finder.markStale(); + finder.update(singleRangeUpdate(1)); + assertThat(lifecycleManager.publishCount.get()).isEqualTo(1); + } finally { + lifecycleManager.shutdown(); + } + } + private static CacheUpdate singleRangeUpdate(int index) { String startKey = String.format("k%05d", index); String limitKey = String.format("k%05d", index + 1); @@ -304,6 +332,19 @@ public String authority() { } } + private static final class RecordingLifecycleManager extends EndpointLifecycleManager { + private final AtomicInteger publishCount = new AtomicInteger(); + + private RecordingLifecycleManager(ChannelEndpointCache endpointCache) { + super(endpointCache); + } + + @Override + void updateActiveAddressesAsync(String finderKey, java.util.Set activeAddresses) { + publishCount.incrementAndGet(); + } + } + private static final class BlockingLifecycleManager extends EndpointLifecycleManager { private final CountDownLatch updateStarted = new CountDownLatch(1); private final CountDownLatch releaseUpdate = new CountDownLatch(1); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java index af9366020da9..deb72e59aa49 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java @@ -1155,6 +1155,41 @@ public void abandonedReadWriteTransactionAffinityExpiresAfterInactivity() throws assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(1); } + @Test + public void channelFinderExpiresAfterInactivity() throws Exception { + FakeTicker ticker = new FakeTicker(); + TestHarness harness = createHarness(ticker); + seedCache(harness, createTwoRangeCacheUpdate()); + + // The learned routing cache routes the query to the tablet's server. + ClientCall routedCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + routedCall.start(new CapturingListener(), new Metadata()); + routedCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build()) + .build()); + assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(1); + + ticker.advance(KeyAwareChannel.CHANNEL_FINDER_TTL_MINUTES + 1, TimeUnit.MINUTES); + + // The idle ChannelFinder has been evicted along with its routing caches, so the same query + // falls back to the default channel until the cache is re-learned. + int defaultCallsBefore = harness.defaultManagedChannel.callCount(); + ClientCall expiredCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + expiredCall.start(new CapturingListener(), new Metadata()); + expiredCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build()) + .build()); + + assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(1); + assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(defaultCallsBefore + 1); + } + private static CacheUpdate createTwoRangeCacheUpdate() { return CacheUpdate.newBuilder() .setDatabaseId(7L) diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java index b7645f044a13..1cdeaab9b917 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java @@ -279,6 +279,65 @@ public void shrinkToEvictsRanges() { checkContents(cache, 0, numRanges); } + @Test + public void addRangesEvictsWhenOverMaxRanges() { + FakeEndpointCache endpointCache = new FakeEndpointCache(); + KeyRangeCache cache = new KeyRangeCache(endpointCache); + final int maxRanges = 50; + cache.setMaxRanges(maxRanges); + + final int numRanges = 100; + for (int i = 0; i < numRanges; i++) { + CacheUpdate update = + CacheUpdate.newBuilder() + .addRange( + Range.newBuilder() + .setStartKey(bytes(String.format("%04d", i))) + .setLimitKey(bytes(String.format("%04d", i + 1))) + .setGroupUid(i) + .setSplitId(i) + .setGeneration(bytes("1"))) + .addGroup( + Group.newBuilder() + .setGroupUid(i) + .setGeneration(bytes("1")) + .addTablets( + Tablet.newBuilder() + .setTabletUid(i) + .setServerAddress("server" + i) + .setIncarnation(bytes("1")))) + .build(); + cache.addRanges(update); + // Pre-create endpoint so READY state check passes in shouldSkip. + endpointCache.get("server" + i); + assertTrue(cache.size() <= maxRanges); + } + + // The shrink hysteresis keeps the cache within 10% below the limit. + assertTrue(cache.size() >= maxRanges - maxRanges / 10); + + // Every resident range still routes to its own server, and the most recently added range + // always survives eviction. + int hitCount = 0; + for (int i = 0; i < numRanges; i++) { + RoutingHint.Builder hint = RoutingHint.newBuilder().setKey(bytes(String.format("%04d", i))); + ChannelEndpoint server = + cache.fillRoutingHint( + false, + KeyRangeCache.RangeMode.COVERING_SPLIT, + DirectedReadOptions.getDefaultInstance(), + hint); + if (server != null) { + hitCount++; + assertEquals("server" + i, server.getAddress()); + } + if (i == numRanges - 1) { + assertNotNull(server); + } + } + assertEquals(cache.size(), hitCount); + } + @Test public void readyEndpointIsUsableForLocationAware() { FakeEndpointCache endpointCache = new FakeEndpointCache();