Skip to content

fix(spanner): bound the location-aware routing caches (P0-5) #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fornwall wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(spanner): bound the location-aware routing caches (P0-5) #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2394790
fix(spanner): bound the location-aware routing caches
fornwall Jul 19, 2026
d7e99b0
test(spanner): cover routing-cache bounding and finder eviction
fornwall Jul 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 28 additions & 19 deletions ...ner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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.
*
* <p>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(
Expand Down Expand Up @@ -136,16 +152,14 @@ private boolean shouldProcessUpdate(CacheUpdate update) {
}

public void update(CacheUpdate update) {
Set<String> 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));
Expand Down Expand Up @@ -187,14 +201,12 @@ private void drainBatch(List<PendingCacheUpdate> batch) {
}

private void applyBatch(List<PendingCacheUpdate> batch) {
Set<String> currentAddresses;
synchronized (updateLock) {
for (PendingCacheUpdate pendingUpdate : batch) {
applyUpdateLocked(pendingUpdate.update);
}
currentAddresses = snapshotActiveAddressesLocked();
publishLifecycleUpdateLocked();
}
publishLifecycleUpdate(currentAddresses);
}

private void applyUpdateLocked(CacheUpdate update) {
Expand All @@ -213,19 +225,16 @@ private void applyUpdateLocked(CacheUpdate update) {
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) {
/**
* 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());
}

/**
Expand Down
105 changes: 40 additions & 65 deletions ...r/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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";
Expand All @@ -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<ChannelFinder> channelFinderReferenceQueue = new ReferenceQueue<>();
private final Map<String, ChannelFinderReference> 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<String, ChannelFinder> 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.
Expand All @@ -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 =
Expand All @@ -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(
Expand Down Expand Up @@ -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<ByteString, String> newTransactionAffinities(Ticker ticker) {
Expand All @@ -182,16 +179,22 @@ private static Cache<ByteString, String> newTransactionAffinities(Ticker ticker)
.build();
}

private static final class ChannelFinderReference extends SoftReference<ChannelFinder> {
final String databaseId;

ChannelFinderReference(
String databaseId,
ChannelFinder referent,
ReferenceQueue<? super ChannelFinder> referenceQueue) {
super(referent, referenceQueue);
this.databaseId = databaseId;
}
private Cache<String, ChannelFinder> newChannelFinders(Ticker ticker) {
return CacheBuilder.newBuilder()
.maximumSize(MAX_CHANNEL_FINDERS)
.expireAfterAccess(CHANNEL_FINDER_TTL_MINUTES, TimeUnit.MINUTES)
.ticker(ticker)
.removalListener(
(RemovalNotification<String, ChannelFinder> notification) -> {
ChannelFinder finder = notification.getValue();
if (finder != null) {
finder.markStale();
}
if (lifecycleManager != null) {
lifecycleManager.unregisterFinder(notification.getKey());
}
})
.build();
}

private String extractDatabaseIdFromSession(String session) {
Expand All @@ -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();
}
}

Expand All @@ -262,7 +237,7 @@ private void onRequestRouted(@Nullable ChannelEndpoint selectedEndpoint) {

@Override
public ManagedChannel shutdown() {
cleanupStaleChannelFinders();
channelFinders.invalidateAll();
if (lifecycleManager != null) {
lifecycleManager.shutdown();
}
Expand All @@ -272,7 +247,7 @@ public ManagedChannel shutdown() {

@Override
public ManagedChannel shutdownNow() {
cleanupStaleChannelFinders();
channelFinders.invalidateAll();
if (lifecycleManager != null) {
lifecycleManager.shutdown();
}
Expand Down
20 changes: 20 additions & 0 deletions ...ner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
}

/**
Expand Down
41 changes: 41 additions & 0 deletions ...google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelFinderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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<String> activeAddresses) {
publishCount.incrementAndGet();
}
}

private static final class BlockingLifecycleManager extends EndpointLifecycleManager {
private final CountDownLatch updateStarted = new CountDownLatch(1);
private final CountDownLatch releaseUpdate = new CountDownLatch(1);
Expand Down
Loading