From d3c0d537ece70c5bf12b09e7258c454558baefab Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 08:46:18 +0530 Subject: [PATCH 1/7] feat(spanner): prime scaled-up dynamic channel pool channels with SELECT 1 The grpc-gcp dynamic channel pool (DCP) now supports a GcpChannelPrimer that runs before a scaled-up channel is published to the pool, but java-spanner never registered one, so channels added under load were published cold: the first requests on them paid for the TCP handshake, TLS session, and server-side connection setup at the moment the pool was already saturated. The Go Spanner client primes scaled-up channels by executing SELECT 1 with the multiplexed session through the new channel. This change brings java-spanner to parity. --- .../MultiplexedSessionDatabaseClient.java | 97 +- .../google/cloud/spanner/SessionClient.java | 44 +- .../com/google/cloud/spanner/SpannerImpl.java | 8 +- .../google/cloud/spanner/SpannerOptions.java | 47 +- .../spi/v1/DynamicChannelPoolPrimer.java | 620 +++++++++++ .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 134 ++- .../cloud/spanner/spi/v1/SpannerRpc.java | 27 +- .../DynamicChannelPoolPrimeSessionTest.java | 387 +++++++ .../MultiplexedSessionDatabaseClientTest.java | 205 +++- .../google/cloud/spanner/SpannerImplTest.java | 78 ++ .../cloud/spanner/SpannerOptionsTest.java | 51 + .../spanner/spi/v1/ChannelPrimerTestRpc.java | 99 ++ .../spi/v1/DynamicChannelPoolPrimerTest.java | 998 ++++++++++++++++++ .../spanner/spi/v1/GapicSpannerRpcTest.java | 855 ++++++++++++++- 14 files changed, 3598 insertions(+), 52 deletions(-) create mode 100644 java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java create mode 100644 java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java create mode 100644 java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java create mode 100644 java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index ece6d862f87b..8b1e768d680e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -27,7 +27,9 @@ import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; +import com.google.cloud.spanner.SessionClient.SessionOption; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.spanner.v1.BatchWriteResponse; @@ -220,6 +222,18 @@ private SharedChannelUsage(int numChannels) { private final AtomicLong numSessionsReleased = new AtomicLong(); + /** Source of the {@link #channelPrimeOwnerTicket owner tickets} of all clients of the process. */ + private static final AtomicLong CHANNEL_PRIME_OWNER_TICKETS = new AtomicLong(); + + /** + * The ticket under which this client owns its database for priming channels that the dynamic + * channel pool adds during scale-up. The ticket is registered with the rpc before the first + * multiplexed session is created, carried by every CreateSession of this client, and unregistered + * when this client is closed, so a session that arrives after the close is never used for + * priming. + */ + private final long channelPrimeOwnerTicket = CHANNEL_PRIME_OWNER_TICKETS.incrementAndGet(); + MultiplexedSessionDatabaseClient(SessionClient sessionClient) { this(sessionClient, Clock.systemUTC()); } @@ -253,15 +267,36 @@ private SharedChannelUsage(int numChannels) { final SettableApiFuture initialSessionReferenceFuture = SettableApiFuture.create(); this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); + // Register as the owner of the database before the first CreateSession, so that session and + // every refreshed session can be attributed to this client for channel priming. + spanner + .getRpc() + .registerChannelPrimeOwner( + sessionClient.getDatabaseId().getName(), channelPrimeOwnerTicket); + + try { + Duration waitDuration = + sessionClient.getSpanner().getOptions().getSessionPoolOptions().getWaitForMinSessions(); + int initialAttempts = + waitDuration == null || waitDuration.isZero() ? MAX_INITIAL_CREATE_SESSION_ATTEMPTS : 1; + asyncCreateMultiplexedSession(initialSessionReferenceFuture, initialAttempts); + maybeWaitForSessionCreation( + sessionClient.getSpanner().getOptions().getSessionPoolOptions(), + initialSessionReferenceFuture); + } catch (Throwable t) { + // The constructor did not complete, so the caller never gets a reference to this client and + // will never close it. Undo everything that was registered above: mark the client closed so + // that a CreateSession that is still in flight neither starts the maintainer nor is handed + // to a waiter, stop the maintainer, drop the channel prime owner ticket so that a late + // session is not recorded for priming, and release the shared channel usage. + close(); + throw t; + } + } - Duration waitDuration = - sessionClient.getSpanner().getOptions().getSessionPoolOptions().getWaitForMinSessions(); - int initialAttempts = - waitDuration == null || waitDuration.isZero() ? MAX_INITIAL_CREATE_SESSION_ATTEMPTS : 1; - asyncCreateMultiplexedSession(initialSessionReferenceFuture, initialAttempts); - maybeWaitForSessionCreation( - sessionClient.getSpanner().getOptions().getSessionPoolOptions(), - initialSessionReferenceFuture); + /** The options of every CreateSession of this client: they carry the owner ticket. */ + private Map createSessionOptions() { + return SessionClient.optionMap(SessionOption.channelPrimeOwner(channelPrimeOwnerTicket)); } private void asyncCreateMultiplexedSession( @@ -270,10 +305,19 @@ private void asyncCreateMultiplexedSession( new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { + synchronized (MultiplexedSessionDatabaseClient.this) { + if (isClosed) { + // The client was closed while the session was being created. Ignore the session: + // it must neither be handed to waiters of a closed client nor keep a maintainer + // running for a client that no longer exists. + sessionReferenceFuture.setException(newClosedException()); + return; + } + // only start the maintainer if we actually managed to create a session in the first + // place. Starting it under the lock guarantees that a concurrent close() stops it. + maintainer.start(); + } sessionReferenceFuture.set(session.getSessionReference()); - // only start the maintainer if we actually managed to create a session in the first - // place. - maintainer.start(); if (sessionClient .getSpanner() .getOptions() @@ -296,16 +340,19 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount (ResourceNotFoundException) spannerException); } // Set the exception to trigger an error for all waiters. - // Then retry the session creation if the error is (potentially) transient. + // Then retry the session creation if the error is (potentially) transient and the + // client has not been closed in the meantime. sessionReferenceFuture.setException(t); if (remainingAttempts > 1 + && !isClientClosed() && RETRYABLE_ERROR_CODES.contains(spannerException.getErrorCode())) { final SettableApiFuture future = SettableApiFuture.create(); MultiplexedSessionDatabaseClient.this.multiplexedSessionReference.set(future); asyncCreateMultiplexedSession(future, remainingAttempts - 1); } } - }); + }, + createSessionOptions()); } private void maybeWaitForSessionCreation( @@ -363,6 +410,15 @@ boolean isValid() { return resourceNotFoundException.get() == null; } + private synchronized boolean isClientClosed() { + return isClosed; + } + + private static SpannerException newClosedException() { + return SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "This client has been closed"); + } + AtomicLong getNumSessionsAcquired() { return this.numSessionsAcquired; } @@ -381,6 +437,13 @@ void close() { } } if (releaseChannelUsage) { + // The multiplexed session is no longer maintained, so it must no longer be used for priming + // dynamic channel pool channels, and a CreateSession of this client that is still in flight + // must not register its session either. + spanner + .getRpc() + .unregisterChannelPrimeOwner( + sessionClient.getDatabaseId().getName(), channelPrimeOwnerTicket); synchronized (CHANNEL_USAGE) { SharedChannelUsage sharedChannelUsage = CHANNEL_USAGE.get(this.spanner); if (sharedChannelUsage != null) { @@ -659,6 +722,11 @@ void maintain() { new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { + if (isClientClosed()) { + // The client was closed while the session was being refreshed. The refreshed + // session belongs to a client that no longer exists and is ignored. + return; + } multiplexedSessionReference.set( ApiFutures.immediateFuture(session.getSessionReference())); expirationDate.set( @@ -673,7 +741,8 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount // we continue to use the session that has passed its expiration date for now, and // that a new attempt at creating a new session will be done in 10 minutes from now. } - }); + }, + createSessionOptions()); } } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java index 234675a221bc..e055faa6b636 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java @@ -34,6 +34,7 @@ import java.util.Map; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; /** Client for creating single sessions and batches of sessions. */ @@ -94,6 +95,10 @@ static SessionOption channelAffinityRef(ChannelAffinityRef channelAffinityRef) { return new SessionOption(SpannerRpc.Option.CHANNEL_ID_AFFINITY, channelAffinityRef); } + static SessionOption channelPrimeOwner(long ownerTicket) { + return new SessionOption(SpannerRpc.Option.CHANNEL_PRIME_OWNER, ownerTicket); + } + SpannerRpc.Option rpcOption() { return rpcOption; } @@ -264,8 +269,19 @@ SessionImpl createSession() { * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available. */ void createMultiplexedSession(SessionConsumer consumer) { + createMultiplexedSession(consumer, /* options= */ null); + } + + /** + * Creates a multiplexed session with the given {@link SpannerRpc.Option options} for the {@code + * CreateSession} call and returns it to the given {@link SessionConsumer}. + * + * @see #createMultiplexedSession(SessionConsumer) + */ + void createMultiplexedSession( + SessionConsumer consumer, @Nullable Map options) { try { - SessionImpl sessionImpl = createMultiplexedSession(); + SessionImpl sessionImpl = createMultiplexedSession(options); consumer.onSessionReady(sessionImpl); } catch (Throwable t) { consumer.onSessionCreateFailure(t, 1); @@ -277,6 +293,16 @@ void createMultiplexedSession(SessionConsumer consumer) { * GRPC channel. In case of an error during the gRPC calls, an exception will be thrown. */ SessionImpl createMultiplexedSession() { + return createMultiplexedSession((Map) null); + } + + /** + * Creates a multiplexed session with the given {@link SpannerRpc.Option options} for the {@code + * CreateSession} call and returns it. + * + * @see #createMultiplexedSession() + */ + SessionImpl createMultiplexedSession(@Nullable Map options) { ISpan span = spanner .getTracer() @@ -289,7 +315,7 @@ SessionImpl createMultiplexedSession() { db.getName(), spanner.getOptions().getDatabaseRole(), spanner.getOptions().getSessionLabels(), - null, + options, true); SessionImpl sessionImpl = new SessionImpl( @@ -319,10 +345,13 @@ SessionImpl createMultiplexedSession() { * SessionConsumer#onSessionCreateFailure(Throwable, int)} call with the error. * * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available. + * @param options The {@link SpannerRpc.Option options} of the {@code CreateSession} call, or + * {@code null} for none. */ - void asyncCreateMultiplexedSession(SessionConsumer consumer) { + void asyncCreateMultiplexedSession( + SessionConsumer consumer, @Nullable Map options) { try { - executor.submit(new CreateMultiplexedSessionsRunnable(consumer)); + executor.submit(new CreateMultiplexedSessionsRunnable(consumer, options)); } catch (Throwable t) { consumer.onSessionCreateFailure(t, 1); } @@ -330,15 +359,18 @@ void asyncCreateMultiplexedSession(SessionConsumer consumer) { private final class CreateMultiplexedSessionsRunnable implements Runnable { private final SessionConsumer consumer; + @Nullable private final Map options; - private CreateMultiplexedSessionsRunnable(SessionConsumer consumer) { + private CreateMultiplexedSessionsRunnable( + SessionConsumer consumer, @Nullable Map options) { Preconditions.checkNotNull(consumer); this.consumer = consumer; + this.options = options; } @Override public void run() { - createMultiplexedSession(consumer); + createMultiplexedSession(consumer, options); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 206f58726257..a3950ba3f382 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -298,7 +298,9 @@ public DatabaseClient getDatabaseClient(DatabaseId db) { checkClosed(); String clientId = null; if (dbClients.containsKey(db) && !dbClients.get(db).isValid()) { - // Close the invalidated client and remove it. + // Close the invalidated client and remove it. Closing it unregisters it as the owner of + // its database for priming dynamic channel pool channels, because its multiplexed session + // is no longer maintained. dbClients.get(db).closeAsync(new ClosedException()); clientId = dbClients.get(db).clientId; dbClients.remove(db); @@ -360,8 +362,8 @@ void close(long timeout, TimeUnit unit) { } try { closureFutures = new ArrayList<>(); - for (DatabaseClientImpl dbClient : dbClients.values()) { - closureFutures.add(dbClient.closeAsync(closedException)); + for (Map.Entry dbClient : dbClients.entrySet()) { + closureFutures.add(dbClient.getValue().closeAsync(closedException)); } dbClients.clear(); Futures.successfulAsList(closureFutures).get(timeout, unit); diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 3dc14d44c8f6..d2b9d2c5b4b8 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -178,6 +178,19 @@ public class SpannerOptions extends ServiceOptions { */ public static final Duration DEFAULT_DYNAMIC_POOL_CLEANUP_INTERVAL = Duration.ofMinutes(1); + /** + * Default maximum time for one attempt to prime a channel that the dynamic channel pool adds + * during scale-up. When multiplexed sessions are enabled, scaled-up channels are primed by + * executing {@code SELECT 1} with the multiplexed session before they are published to the pool. + */ + public static final Duration DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT = Duration.ofSeconds(10); + + /** + * Default maximum number of attempts to prime a channel that the dynamic channel pool adds during + * scale-up before the channel is discarded. + */ + public static final int DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS = 3; + /** * Creates a {@link GcpChannelPoolOptions} instance with Spanner-specific defaults for dynamic * channel pooling. These defaults are optimized for typical Spanner workloads. @@ -193,8 +206,18 @@ public class SpannerOptions extends ServiceOptions { *
  • Scale down interval: 3 minutes *
  • Affinity key lifetime: 10 minutes *
  • Cleanup interval: 1 minute + *
  • Channel prime timeout: 10 seconds + *
  • Channel prime max attempts: {@value #DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS} * * + *

    When multiplexed sessions are enabled (the default), channels that the pool adds during + * scale-up are primed with {@code SELECT 1} on the multiplexed session before they are published. + * The primer is registered by the Spanner client when dynamic channel pooling is enabled, unless + * these options already contain a primer. Priming uses the most recently created multiplexed + * session of any database client of the {@link Spanner} instance; a session that turns out to be + * invalid is dropped and the next most recent one is used. When multiplexed sessions are + * disabled, scaled-up channels are published without priming. + * * @return a new {@link GcpChannelPoolOptions} instance with Spanner defaults */ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() { @@ -208,13 +231,16 @@ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() { DEFAULT_DYNAMIC_POOL_SCALE_DOWN_INTERVAL) .setAffinityKeyLifetime(DEFAULT_DYNAMIC_POOL_AFFINITY_KEY_LIFETIME) .setCleanupInterval(DEFAULT_DYNAMIC_POOL_CLEANUP_INTERVAL) + .setChannelPrimeTimeout(DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT) + .setChannelPrimeMaxAttempts(DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS) .build(); } /** * Merges user-provided {@link GcpChannelPoolOptions} with Spanner-specific defaults. Any value * that the user has not explicitly set (i.e. left at the builder's default of 0 or null) will be - * filled in from {@link #createDefaultDynamicChannelPoolOptions()}. + * filled in from {@link #createDefaultDynamicChannelPoolOptions()}. A user-provided channel + * primer, prime timeout, and prime attempt count are always preserved. */ static GcpChannelPoolOptions mergeWithDefaultChannelPoolOptions( GcpChannelPoolOptions userOptions) { @@ -2031,6 +2057,14 @@ public Builder disableGrpcGcpExtension() { * Enables dynamic channel pooling. When enabled, the client will automatically scale the number * of channels based on load. This requires the gRPC-GCP extension to be enabled. * + *

    When multiplexed sessions are enabled (the default), channels that the pool adds during + * scale-up are primed before they serve traffic: the client executes {@code SELECT 1} with the + * multiplexed session on the new channel, and the pool only publishes the channel once that + * succeeds. When multiplexed sessions are disabled, scaled-up channels are published without + * priming. See {@link #createDefaultDynamicChannelPoolOptions()} for the prime timeout and + * attempt defaults, and {@link #setGcpChannelPoolOptions(GcpChannelPoolOptions)} to customize + * them. + * *

    Dynamic channel pooling is disabled by default. Use this method to explicitly enable it. * Note that calling {@link #setNumChannels(int)} will disable dynamic channel pooling even if * this method was called. @@ -2068,7 +2102,16 @@ public Builder setGrpcGcpOtelMetricsEnabled(boolean enableGrpcGcpOtelMetrics) { * channel pool behavior when {@link #enableDynamicChannelPool()} is enabled. * *

    If not set, Spanner-specific defaults will be used (see {@link - * #createDefaultDynamicChannelPoolOptions()}). + * #createDefaultDynamicChannelPoolOptions()}). Values that are left unset in the given options + * are filled in from those defaults. + * + *

    When multiplexed sessions are enabled (the default), channels that the pool adds during + * scale-up are primed with {@code SELECT 1} on the multiplexed session before they are + * published; otherwise they are published without priming. A channel primer, prime timeout, or + * prime attempt count that is set in the given options takes precedence over the Spanner primer + * and its defaults. The deadline of the Spanner primer's RPC is derived from the prime timeout + * and normally stays below it; only a prime timeout of two milliseconds or less yields the + * primer's minimum deadline of one millisecond, and the prime timeout then bounds the attempt. * *

    Example usage: * diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java new file mode 100644 index 000000000000..543d6a2e3b22 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -0,0 +1,620 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.spi.v1; + +import com.google.cloud.grpc.GcpChannelPrimer; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.SpannerExceptionFactory; +import com.google.cloud.spanner.SpannerOptions; +import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; +import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ResultSet; +import com.google.spanner.v1.SpannerGrpc; +import io.grpc.CallCredentials; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ClientInterceptors; +import io.grpc.ForwardingClientCall.SimpleForwardingClientCall; +import io.grpc.ManagedChannel; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.stub.ClientCalls; +import java.io.UnsupportedEncodingException; +import java.net.URLEncoder; +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.regex.Pattern; +import javax.annotation.Nullable; + +/** + * Primes channels that the grpc-gcp dynamic channel pool adds during scale-up by executing {@code + * SELECT 1} on the new channel with a multiplexed session, so the transport, TLS session, and + * server-side connection are established before the channel serves live traffic. Mirrors the + * priming that the Go Spanner client performs for its dynamic channel pool. Priming requires + * multiplexed sessions: when they are disabled on the {@link SpannerOptions} there is nothing to + * prime with, and scaled-up channels are published unprimed, which is the behavior of the pool + * without a primer. + * + *

    The channel pool is shared by all database clients of a {@link Spanner} instance, while + * multiplexed sessions are created per database. The primer therefore keeps a small registry of + * prime sessions keyed by database name. {@link GapicSpannerRpc} registers every multiplexed + * session that it successfully creates, which replaces the previous entry of that database, and the + * registry hands out the entry with the highest creation generation, that is the most recently + * created multiplexed session. A multiplexed {@code CreateSession} is always the first RPC of a + * database client before any data RPC is possible, and the periodic refresh of a multiplexed + * session goes through the same path, so the most recently created multiplexed session is normally + * a valid session to prime with. {@code SELECT 1} against any database is sufficient to prime a + * channel. + * + *

    A registered session can nevertheless become invalid, for example when its database is + * dropped. The registry is therefore maintained on rare paths only: the entries of a database are + * removed when its database client is invalidated or closed, and a priming attempt that fails for a + * reason that is specific to the session it used, such as the session not being found, the session + * or its database role being invalid, or permission being denied, evicts exactly the entry that it + * used, so the pool's next attempt falls back to the next most recent session. Nothing is written + * from any data RPC, and the scale-up path performs a single volatile read to select the session. + * + *

    A {@code CreateSession} that is still in flight when its database client is retired must not + * register its session after the client has been retired. Every database client that creates + * multiplexed sessions therefore registers itself as the owner of its database through {@link + * #registerPrimeOwner(String, long)} with a unique owner ticket before its first {@code + * CreateSession}, includes that ticket in every {@code CreateSession} it issues, and unregisters + * itself through {@link #unregisterPrimeOwner(String, long)} when it is invalidated or closed. A + * session is registered only when the ticket of its {@code CreateSession} is the current owner + * ticket of its database, so a late response of a retired client finds no owner, or the ticket of + * the replacement client, and is dropped. The owner map holds one entry per live database client + * and nothing for retired ones, so it is bounded by the number of live clients. This mirrors the Go + * client, which installs a newly created multiplexed session only while its creation is still the + * current one and the session manager is still valid. + */ +final class DynamicChannelPoolPrimer implements GcpChannelPrimer { + /** The statement executed on every scaled-up channel. */ + static final String PRIME_SQL = "SELECT 1"; + + /** + * Upper bound of the deadline of a single priming RPC. Deliberately well below the pool's default + * prime timeout of 10 seconds so an unresponsive backend never leaves the call lingering on the + * server after the pool has given up on the attempt. See {@link #rpcDeadlineFor(Duration)}. + */ + static final Duration MAX_RPC_DEADLINE = Duration.ofSeconds(5); + + /** + * Safety margin that the RPC deadline keeps below the pool's prime timeout, so the RPC fails with + * DEADLINE_EXCEEDED and is cleaned up before the pool times out the attempt. + */ + static final Duration RPC_DEADLINE_MARGIN = Duration.ofSeconds(1); + + /** + * Lower bound of the deadline of a single priming RPC. A prime timeout of at most twice this + * value yields this deadline, which may then reach or exceed the prime timeout itself; the pool's + * prime timeout cancels the attempt, and with it the RPC, in that case. + */ + static final Duration MIN_RPC_DEADLINE = Duration.ofMillis(1); + + private static final Metadata.Key ROUTE_TO_LEADER_KEY = + Metadata.Key.of("x-goog-spanner-route-to-leader", Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key REQUEST_PARAMS_KEY = + Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); + + /** Matches failure messages about the database role of the session, such as an invalid role. */ + private static final Pattern ROLE_PATTERN = Pattern.compile("\\brole\\b"); + + /** + * Channel number of the request id of a priming RPC. The channel that is being primed is not part + * of the pool yet, so it has no pool channel id and the request id uses 0 (unknown). + */ + static final int REQUEST_ID_CHANNEL = 0; + + private final SpannerMetadataProvider metadataProvider; + private final String projectName; + private final RequestIdCreator requestIdCreator; + @Nullable private final CallCredentials defaultCallCredentials; + @Nullable private final CallCredentialsProvider callCredentialsProvider; + private final boolean routeToLeader; + private final boolean multiplexedSessionsEnabled; + private final Duration rpcDeadline; + + /** Source of the creation generations of {@link PrimeSession} entries. */ + private final AtomicLong generations = new AtomicLong(); + + /** + * The registered prime sessions, at most one per database, sorted by descending generation. The + * list is replaced as a whole on every write, so a read is a single volatile read. + */ + private volatile ImmutableList primeSessions = ImmutableList.of(); + + /** + * The owner ticket of the live database client of each database. Guarded by {@link #writeLock}. + * An entry is put by {@link #registerPrimeOwner(String, long)} and removed by {@link + * #unregisterPrimeOwner(String, long)}, so the map holds exactly one entry per live database + * client that creates multiplexed sessions and nothing for retired clients. + */ + private final Map primeOwners = new HashMap<>(); + + /** Guards the writers of {@link #primeSessions}; all of them run on rare paths only. */ + private final Object writeLock = new Object(); + + /** A multiplexed session that is registered for priming. */ + static final class PrimeSession { + private final String databaseName; + private final String sessionName; + private final long ownerTicket; + private final long generation; + + private PrimeSession( + String databaseName, String sessionName, long ownerTicket, long generation) { + this.databaseName = databaseName; + this.sessionName = sessionName; + this.ownerTicket = ownerTicket; + this.generation = generation; + } + + String getDatabaseName() { + return databaseName; + } + + String getSessionName() { + return sessionName; + } + + /** The owner ticket of the database client that created the session. */ + long getOwnerTicket() { + return ownerTicket; + } + + /** The creation generation; a higher generation means a more recently created session. */ + long getGeneration() { + return generation; + } + + @Override + public boolean equals(Object other) { + if (!(other instanceof PrimeSession)) { + return false; + } + PrimeSession that = (PrimeSession) other; + return generation == that.generation + && ownerTicket == that.ownerTicket + && databaseName.equals(that.databaseName) + && sessionName.equals(that.sessionName); + } + + @Override + public int hashCode() { + return Objects.hash(databaseName, sessionName, ownerTicket, generation); + } + + @Override + public String toString() { + return sessionName + "#" + generation; + } + } + + /** + * @param metadataProvider provides the fixed headers and the resource-prefix header of a normal + * Spanner call + * @param projectName the project resource name that serves as the default resource-prefix value + * @param requestIdCreator the request id creator of the owning rpc, so priming RPCs carry a + * request id with the same client id as every other call of the rpc + * @param defaultCallCredentials the credentials that GAX attaches to normal calls, or {@code + * null} when the client runs without credentials + * @param callCredentialsProvider the optional user-supplied provider that takes precedence over + * {@code defaultCallCredentials} for each call, exactly like for normal Spanner calls + * @param routeToLeader whether normal Spanner calls carry the route-to-leader header + * @param multiplexedSessionsEnabled whether the client creates multiplexed sessions at all + * @param rpcDeadline the deadline of a single priming RPC + */ + DynamicChannelPoolPrimer( + SpannerMetadataProvider metadataProvider, + String projectName, + RequestIdCreator requestIdCreator, + @Nullable CallCredentials defaultCallCredentials, + @Nullable CallCredentialsProvider callCredentialsProvider, + boolean routeToLeader, + boolean multiplexedSessionsEnabled, + Duration rpcDeadline) { + this.metadataProvider = Preconditions.checkNotNull(metadataProvider); + this.projectName = Preconditions.checkNotNull(projectName); + this.requestIdCreator = Preconditions.checkNotNull(requestIdCreator); + this.defaultCallCredentials = defaultCallCredentials; + this.callCredentialsProvider = callCredentialsProvider; + this.routeToLeader = routeToLeader; + this.multiplexedSessionsEnabled = multiplexedSessionsEnabled; + Preconditions.checkArgument( + rpcDeadline != null && !rpcDeadline.isZero() && !rpcDeadline.isNegative(), + "rpcDeadline must be positive"); + this.rpcDeadline = rpcDeadline; + } + + /** + * Returns the deadline of a single priming RPC for a pool whose prime timeout is {@code + * primeTimeout}: at most {@link #MAX_RPC_DEADLINE}, at least {@link #MIN_RPC_DEADLINE}, and + * otherwise below the prime timeout. The RPC deadline keeps {@link #RPC_DEADLINE_MARGIN} below + * the prime timeout, but never drops below half of the prime timeout so short prime timeouts + * still leave the RPC a usable deadline. Prime timeouts of at most twice {@link + * #MIN_RPC_DEADLINE} yield the minimum deadline, which may reach or exceed the prime timeout; the + * pool's prime timeout then cancels the attempt and the RPC. + * + * @param primeTimeout the (positive) prime timeout of the pool + */ + static Duration rpcDeadlineFor(Duration primeTimeout) { + Preconditions.checkArgument( + primeTimeout != null && !primeTimeout.isZero() && !primeTimeout.isNegative(), + "primeTimeout must be positive"); + Duration withMargin = primeTimeout.minus(RPC_DEADLINE_MARGIN); + Duration half = primeTimeout.dividedBy(2); + Duration deadline = withMargin.compareTo(half) > 0 ? withMargin : half; + if (deadline.compareTo(MAX_RPC_DEADLINE) > 0) { + return MAX_RPC_DEADLINE; + } + return deadline.compareTo(MIN_RPC_DEADLINE) < 0 ? MIN_RPC_DEADLINE : deadline; + } + + /** Returns the deadline of a single priming RPC. */ + @VisibleForTesting + Duration getRpcDeadline() { + return rpcDeadline; + } + + /** + * Registers the database client with the given owner ticket as the current owner of the given + * database. Must be called before the client issues its first {@code CreateSession}, so every + * session that the client creates can be attributed to it through {@link + * #registerPrimeSession(String, String, long)}. A previous owner of the database, if any, is + * replaced and its registered session, if any, is removed, because the previous client has been + * or is being retired. + * + * @param ownerTicket a ticket that is unique among all database clients of the process + */ + void registerPrimeOwner(String databaseName, long ownerTicket) { + checkDatabaseName(databaseName); + synchronized (writeLock) { + Long previous = primeOwners.put(databaseName, ownerTicket); + if (previous != null && previous != ownerTicket) { + removeEntries(databaseName, previous); + } + } + } + + /** + * Unregisters the database client with the given owner ticket as the owner of the given database, + * and removes the session that it registered, if any. Called when the client is invalidated or + * closed, because its multiplexed session is then no longer maintained. A late {@code + * CreateSession} response of the client is dropped afterwards, because its ticket is no longer + * the current owner ticket of the database. If the database has already been taken over by a + * replacement client, the replacement and its session are left untouched. + */ + void unregisterPrimeOwner(String databaseName, long ownerTicket) { + checkDatabaseName(databaseName); + synchronized (writeLock) { + Long current = primeOwners.get(databaseName); + if (current != null && current == ownerTicket) { + primeOwners.remove(databaseName); + } + removeEntries(databaseName, ownerTicket); + } + } + + /** + * Registers a multiplexed session that was just created successfully for the given database as a + * prime session. The entry replaces any previous entry of the same database, for example when a + * database client refreshes its multiplexed session, and receives the highest generation, so it + * becomes the session that the next priming attempt uses. + * + *

    The registration is dropped unless the given owner ticket, which the {@code CreateSession} + * request carried, is the current owner ticket of the database: the session then belongs to a + * database client that has been invalidated or closed while its {@code CreateSession} was in + * flight, and must not replace the session of the client that succeeded it. + * + * @param ownerTicket the owner ticket that the {@code CreateSession} request carried + * @return whether the session was registered + */ + boolean registerPrimeSession(String databaseName, String sessionName, long ownerTicket) { + checkDatabaseName(databaseName); + Preconditions.checkArgument( + sessionName != null && !sessionName.isEmpty(), "sessionName must not be empty"); + synchronized (writeLock) { + Long current = primeOwners.get(databaseName); + if (current == null || current != ownerTicket) { + return false; + } + PrimeSession entry = + new PrimeSession(databaseName, sessionName, ownerTicket, generations.incrementAndGet()); + ImmutableList.Builder builder = ImmutableList.builder(); + // The new entry has the highest generation, so it goes first to keep the descending order. + builder.add(entry); + for (PrimeSession existing : primeSessions) { + if (!existing.databaseName.equals(databaseName)) { + builder.add(existing); + } + } + primeSessions = builder.build(); + return true; + } + } + + /** Removes the entries of the given database that were registered with the given ticket. */ + private void removeEntries(String databaseName, long ownerTicket) { + ImmutableList.Builder builder = ImmutableList.builder(); + boolean removed = false; + for (PrimeSession existing : primeSessions) { + if (existing.databaseName.equals(databaseName) && existing.ownerTicket == ownerTicket) { + removed = true; + } else { + builder.add(existing); + } + } + if (removed) { + primeSessions = builder.build(); + } + } + + private static void checkDatabaseName(String databaseName) { + Preconditions.checkArgument( + databaseName != null && !databaseName.isEmpty(), "databaseName must not be empty"); + } + + /** + * Removes exactly the given entry, and returns whether it was still registered. An entry that has + * already been replaced by a newer session of the same database, or that has already been + * removed, is left alone, so an eviction never clobbers a newer registration. + */ + @VisibleForTesting + boolean evictPrimeSession(PrimeSession entry) { + synchronized (writeLock) { + if (!primeSessions.contains(entry)) { + return false; + } + ImmutableList.Builder builder = ImmutableList.builder(); + for (PrimeSession existing : primeSessions) { + if (!existing.equals(entry)) { + builder.add(existing); + } + } + primeSessions = builder.build(); + return true; + } + } + + /** + * Returns the entry with the highest generation, that is the most recently created registered + * multiplexed session, or {@code null} if none is registered. + */ + @VisibleForTesting + @Nullable + PrimeSession getPrimeSession() { + ImmutableList current = primeSessions; + return current.isEmpty() ? null : current.get(0); + } + + /** Returns the name of the session that the next priming attempt uses, or {@code null}. */ + @VisibleForTesting + @Nullable + String getPrimeSessionName() { + PrimeSession entry = getPrimeSession(); + return entry == null ? null : entry.sessionName; + } + + /** Returns all registered entries in descending generation order. */ + @VisibleForTesting + ImmutableList getPrimeSessions() { + return primeSessions; + } + + /** Returns a copy of the current owner ticket of every database with a live database client. */ + @VisibleForTesting + Map getPrimeOwners() { + synchronized (writeLock) { + return new HashMap<>(primeOwners); + } + } + + @Override + public ListenableFuture prime(ManagedChannel channel) { + if (!multiplexedSessionsEnabled) { + // Priming executes SELECT 1 with a multiplexed session, because that is the only session + // that lives at the level of the shared channel pool rather than inside a session pool. + // Without multiplexed sessions there is nothing to prime with, so the channel is published + // unprimed, which is the behavior of a pool without a primer. + return Futures.immediateFuture(null); + } + PrimeSession entry = getPrimeSession(); + if (entry == null) { + // The Go client refuses to scale up at all before a multiplexed session exists. The Java + // primer cannot gate the pool's scale-up decision, so the attempt fails fast instead and the + // pool's retry with backoff and its close-on-failure behaviour handle it. The wasted dial + // and the prime-failure metrics of such an attempt are the known divergence from Go. + return Futures.immediateFailedFuture( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "Cannot prime a dynamic channel pool channel before a multiplexed session is" + + " available")); + } + return executePrimeStatement(channel, entry); + } + + private ListenableFuture executePrimeStatement(ManagedChannel channel, PrimeSession entry) { + String sessionName = entry.sessionName; + ExecuteSqlRequest request = + ExecuteSqlRequest.newBuilder().setSession(sessionName).setSql(PRIME_SQL).build(); + // The priming query always uses the unary ExecuteSql method with an explicit short deadline, + // as the Go client does. ExecuteStreamingSql must never be used here: its configured default + // deadline is one hour, which would let a hung priming call outlive the pool's prime timeout + // by far. + CallOptions callOptions = + CallOptions.DEFAULT.withDeadlineAfter(rpcDeadline.toNanos(), TimeUnit.NANOSECONDS); + CallCredentials callCredentials = resolveCallCredentials(); + if (callCredentials != null) { + callOptions = callOptions.withCallCredentials(callCredentials); + } + Channel channelWithHeaders = + ClientInterceptors.intercept( + channel, new AttachHeadersInterceptor(newHeaders(sessionName))); + ClientCall call = + channelWithHeaders.newCall(SpannerGrpc.getExecuteSqlMethod(), callOptions); + // Cancelling the future returned by futureUnaryCall cancels the underlying ClientCall, and + // cancelling the derived futures cancels their input. A channel that is shut down under the + // running call fails the call with UNAVAILABLE or CANCELLED, which fails this future. + ListenableFuture result = + Futures.transform( + ClientCalls.futureUnaryCall(call, request), + resultSet -> null, + MoreExecutors.directExecutor()); + return Futures.catchingAsync( + result, + Throwable.class, + failure -> { + if (isCandidateSpecificFailure(failure)) { + // The failure is specific to the session that was used, for example because its + // database was dropped or the caller lost access to it. Evict exactly the entry that + // was used, so the pool's next attempt falls back to the next most recent session. A + // newer session of the same database that was registered in the meantime is never + // clobbered. + evictPrimeSession(entry); + } + return Futures.immediateFailedFuture(failure); + }, + MoreExecutors.directExecutor()); + } + + /** + * Returns whether a priming failure is terminal for the session that was used, so priming should + * rotate to the next candidate: the session or its database is not found, permission to use it is + * denied, or its session or database role is reported as invalid. Only such failures evict the + * session; transient failures such as UNAVAILABLE or DEADLINE_EXCEEDED say nothing about the + * session and keep it registered. + */ + @VisibleForTesting + static boolean isCandidateSpecificFailure(Throwable failure) { + Status status = Status.fromThrowable(failure); + switch (status.getCode()) { + case NOT_FOUND: + case PERMISSION_DENIED: + return true; + case FAILED_PRECONDITION: + case INVALID_ARGUMENT: + break; + default: + return false; + } + String description = status.getDescription(); + if (description == null) { + return false; + } + String message = description.toLowerCase(Locale.ENGLISH); + if (ROLE_PATTERN.matcher(message).find()) { + return true; + } + return message.contains("session") + && (message.contains("not found") + || message.contains("invalid") + || message.contains("does not exist")); + } + + @Nullable + private CallCredentials resolveCallCredentials() { + if (callCredentialsProvider != null) { + CallCredentials callCredentials = callCredentialsProvider.getCallCredentials(); + if (callCredentials != null) { + return callCredentials; + } + } + return defaultCallCredentials; + } + + /** + * Returns the per-call headers of a priming RPC. The delegate channel that the pool hands to the + * primer is built by the same channel provider as every other channel of the pool, so it already + * carries the fixed headers of the client (x-goog-api-client, the user agent, and any custom + * headers) through the GAX header interceptor. Only the headers that a normal call adds per call + * are attached here: the resource-prefix header, x-goog-request-params, the route-to-leader + * header when leader-aware routing is enabled, and x-goog-spanner-request-id. + */ + @VisibleForTesting + Metadata newHeaders(String sessionName) { + Metadata headers = new Metadata(); + // The session name starts with the database name, which the metadata provider extracts as the + // value of the resource-prefix header, exactly as for normal calls that pass the session name + // as the resource. + for (Map.Entry> header : + metadataProvider.newExtraHeaders(sessionName, projectName).entrySet()) { + Metadata.Key key = Metadata.Key.of(header.getKey(), Metadata.ASCII_STRING_MARSHALLER); + for (String value : header.getValue()) { + headers.put(key, value); + } + } + headers.put(REQUEST_PARAMS_KEY, "session=" + urlEncode(sessionName)); + if (routeToLeader) { + headers.put(ROUTE_TO_LEADER_KEY, "true"); + } + // The pool invokes prime() once per attempt, so every attempt gets a fresh request id with + // attempt number 1. The header is written directly and the request id is deliberately not set + // as a call option: the RequestIdInterceptor on the delegate channel only acts on the call + // option, so this never produces a duplicate header. + XGoogSpannerRequestId requestId = requestIdCreator.nextRequestId(REQUEST_ID_CHANNEL); + requestId.incrementAttempt(); + headers.put(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY, requestId.getHeaderValue()); + return headers; + } + + private static String urlEncode(String value) { + try { + return URLEncoder.encode(value, "UTF-8"); + } catch (UnsupportedEncodingException e) { + throw new IllegalStateException("UTF-8 is not supported", e); + } + } + + /** Merges the per-call headers of a priming RPC into the request headers of every call. */ + private static final class AttachHeadersInterceptor implements ClientInterceptor { + private final Metadata headers; + + private AttachHeadersInterceptor(Metadata headers) { + this.headers = headers; + } + + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + return new SimpleForwardingClientCall(next.newCall(method, callOptions)) { + @Override + public void start(Listener responseListener, Metadata requestHeaders) { + requestHeaders.merge(headers); + super.start(responseListener, requestHeaders); + } + }; + } + } +} diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index 8e2bb1c5ebdf..48278c552515 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -61,6 +61,7 @@ import com.google.auth.Credentials; import com.google.cloud.RetryHelper; import com.google.cloud.RetryHelper.RetryHelperException; +import com.google.cloud.grpc.GcpChannelPrimer; import com.google.cloud.grpc.GcpManagedChannel; import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; import com.google.cloud.grpc.GcpManagedChannelBuilder; @@ -309,6 +310,7 @@ public class GapicSpannerRpc implements SpannerRpc { private final boolean isDynamicChannelPoolEnabled; @Nullable private final KeyAwareChannel keyAwareChannel; @Nullable private final GcpManagedChannel grpcGcpChannel; + @Nullable private final DynamicChannelPoolPrimer channelPrimer; private final GrpcCallContext baseGrpcCallContext; @@ -373,6 +375,7 @@ public GapicSpannerRpc(final SpannerOptions options) { if (initializeStubs) { CredentialsProvider credentialsProvider = GrpcTransportOptions.setUpCredentialsProvider(options); + this.channelPrimer = createChannelPrimer(options, credentialsProvider); InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder = createBaseChannelProviderBuilder( @@ -388,9 +391,10 @@ public GapicSpannerRpc(final SpannerOptions options) { defaultChannelProviderBuilder, options, headerProviderWithUserAgent, - credentialsProvider); + credentialsProvider, + channelPrimer); } else { - maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options); + maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options, channelPrimer); } boolean enableLocationApi = options.isEnableLocationApi(); @@ -571,6 +575,7 @@ public UnaryCallable createUnaryCalla } else { this.keyAwareChannel = null; this.grpcGcpChannel = null; + this.channelPrimer = null; this.databaseAdminStub = null; this.instanceAdminStub = null; this.spannerStub = null; @@ -646,7 +651,8 @@ private void setupGcpFallback( InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder, final SpannerOptions options, final HeaderProvider headerProviderWithUserAgent, - final CredentialsProvider credentialsProvider) { + final CredentialsProvider credentialsProvider, + @Nullable final DynamicChannelPoolPrimer channelPrimer) { InstantiatingGrpcChannelProvider.Builder cloudPathProviderBuilder = createBaseChannelProviderBuilder( options, headerProviderWithUserAgent, /* isEnableDirectAccess= */ false); @@ -698,7 +704,8 @@ public ClientCall interceptCall( ManagedChannelBuilder fallbackBuilder = cloudPathBuilder; if (options.isGrpcGcpExtensionEnabled()) { String jsonApiConfig = parseGrpcGcpApiConfig(); - GcpManagedChannelOptions gcpOptions = grpcGcpOptionsWithMetricsAndDcp(options); + GcpManagedChannelOptions gcpOptions = + grpcGcpOptionsWithMetricsAndDcp(options, channelPrimer); if (gcpOptions == null) { gcpOptions = GcpManagedChannelOptions.newBuilder().build(); } @@ -731,7 +738,7 @@ private InstantiatingGrpcChannelProvider.Builder createChannelProviderBuilder( InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder = createBaseChannelProviderBuilder( options, headerProviderWithUserAgent, isEnableDirectAccess); - maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options); + maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options, channelPrimer); return defaultChannelProviderBuilder; } @@ -790,8 +797,65 @@ private InstantiatingGrpcChannelProvider.Builder createBaseChannelProviderBuilde return defaultChannelProviderBuilder; } + /** + * Creates the primer for channels that the dynamic channel pool adds during scale-up, or {@code + * null} if the pool is not dynamic. The primer carries the same credentials and per-call headers + * as a normal Spanner call: the user-supplied {@link CallCredentialsProvider} takes precedence, + * otherwise the scoped credentials that GAX attaches to every call are used, and the request ids + * of priming RPCs come from the request id creator of this rpc. The deadline of the priming RPC + * is derived from the pool's prime timeout and normally stays below it; see {@link + * DynamicChannelPoolPrimer#rpcDeadlineFor(Duration)} for the minimum deadline that a prime + * timeout of two milliseconds or less yields. The sessions that the primer uses are registered by + * {@link #createSession(String, String, Map, Map, boolean)} whenever a multiplexed session is + * created successfully for the current owner of its database, which the database clients register + * through {@link #registerChannelPrimeOwner(String, long)} and remove through {@link + * #unregisterChannelPrimeOwner(String, long)}, so a session whose creation was still in flight + * when its database client was retired is never registered. + */ + @Nullable + private DynamicChannelPoolPrimer createChannelPrimer( + SpannerOptions options, CredentialsProvider credentialsProvider) { + if (!options.isGrpcGcpExtensionEnabled() || !options.isDynamicChannelPoolEnabled()) { + return null; + } + CallCredentials defaultCallCredentials = null; + try { + Credentials credentials = credentialsProvider.getCredentials(); + if (credentials != null) { + defaultCallCredentials = MoreCallCredentials.from(credentials); + } + } catch (IOException e) { + throw newSpannerException(e); + } + return new DynamicChannelPoolPrimer( + metadataProvider, + projectName, + requestIdCreator, + defaultCallCredentials, + callCredentialsProvider, + leaderAwareRoutingEnabled, + options.getSessionPoolOptions().getUseMultiplexedSession(), + // The options already contain the merged prime timeout, including a user-provided one. + DynamicChannelPoolPrimer.rpcDeadlineFor( + options.getGcpChannelPoolOptions().getChannelPrimeTimeout())); + } + + /** Returns the primer of scaled-up dynamic channel pool channels, or {@code null} if none. */ + @VisibleForTesting + @Nullable + DynamicChannelPoolPrimer getChannelPrimer() { + return channelPrimer; + } + // Enhance gRPC-GCP options with metrics and dynamic channel pool configuration. private static GcpManagedChannelOptions grpcGcpOptionsWithMetricsAndDcp(SpannerOptions options) { + return grpcGcpOptionsWithMetricsAndDcp(options, /* channelPrimer= */ null); + } + + // Enhance gRPC-GCP options with metrics and dynamic channel pool configuration, including the + // given primer for scaled-up channels. + private static GcpManagedChannelOptions grpcGcpOptionsWithMetricsAndDcp( + SpannerOptions options, @Nullable GcpChannelPrimer channelPrimer) { GcpManagedChannelOptions grpcGcpOptions = MoreObjects.firstNonNull(options.getGrpcGcpOptions(), new GcpManagedChannelOptions()); GcpManagedChannelOptions.Builder optionsBuilder = @@ -817,7 +881,7 @@ private static GcpManagedChannelOptions grpcGcpOptionsWithMetricsAndDcp(SpannerO // applied regardless of whether dynamic channel pool is enabled. In the non-DCP path, only // propagate the affinity cleanup configuration to avoid implicitly turning on dynamic scaling. if (options.isGrpcGcpExtensionEnabled()) { - optionsBuilder.withChannelPoolOptions(getGrpcGcpChannelPoolOptions(options)); + optionsBuilder.withChannelPoolOptions(getGrpcGcpChannelPoolOptions(options, channelPrimer)); } return optionsBuilder.build(); @@ -825,8 +889,26 @@ private static GcpManagedChannelOptions grpcGcpOptionsWithMetricsAndDcp(SpannerO @VisibleForTesting static GcpChannelPoolOptions getGrpcGcpChannelPoolOptions(SpannerOptions options) { + return getGrpcGcpChannelPoolOptions(options, /* channelPrimer= */ null); + } + + /** + * Returns the grpc-gcp channel pool options. With dynamic channel pooling, the given primer is + * registered for scaled-up channels unless the user already supplied a primer through their own + * {@link GcpChannelPoolOptions}. A user-provided primer, prime timeout, and attempt count are + * never overridden. Without dynamic channel pooling, the pool never scales up and no primer is + * registered. + */ + @VisibleForTesting + static GcpChannelPoolOptions getGrpcGcpChannelPoolOptions( + SpannerOptions options, @Nullable GcpChannelPrimer channelPrimer) { GcpChannelPoolOptions channelPoolOptions = options.getGcpChannelPoolOptions(); if (options.isDynamicChannelPoolEnabled()) { + if (channelPrimer != null && channelPoolOptions.getChannelPrimer() == null) { + return GcpChannelPoolOptions.newBuilder(channelPoolOptions) + .setChannelPrimer(channelPrimer) + .build(); + } return channelPoolOptions; } @@ -875,13 +957,15 @@ private static GrpcGcpEndpointChannelConfigurator createGrpcGcpEndpointChannelCo @SuppressWarnings("rawtypes") private static void maybeEnableGrpcGcpExtension( InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder, - final SpannerOptions options) { + final SpannerOptions options, + @Nullable final DynamicChannelPoolPrimer channelPrimer) { if (!options.isGrpcGcpExtensionEnabled()) { return; } final String jsonApiConfig = parseGrpcGcpApiConfig(); - final GcpManagedChannelOptions grpcGcpOptions = grpcGcpOptionsWithMetricsAndDcp(options); + final GcpManagedChannelOptions grpcGcpOptions = + grpcGcpOptionsWithMetricsAndDcp(options, channelPrimer); ApiFunction baseConfigurator = defaultChannelProviderBuilder.getChannelConfigurator(); @@ -1969,7 +2053,39 @@ public Session createSession( CreateSessionRequest request = requestBuilder.build(); GrpcCallContext context = newCallContext(options, databaseName, request, SpannerGrpc.getCreateSessionMethod(), true); - return get(spannerStub.createSessionCallable().futureCall(request, context)); + Session session = get(spannerStub.createSessionCallable().futureCall(request, context)); + // Every multiplexed session, including the periodic refresh of an existing one, is created + // through this method, so the most recently created multiplexed session is normally a valid + // session for priming channels that the dynamic channel pool adds during scale-up. The + // registration is keyed on both the request flag and the response flag: a backend that does + // not echo the multiplexed flag, for example because it ignores the request flag and creates a + // regular session, must never turn that session into a prime session. Regular sessions are + // not safe for the concurrent use that priming implies. The owner ticket of the database + // client that issued the request is passed on, so the primer can drop the session if the + // client has been retired while the request was in flight; a request without an owner ticket + // never registers a prime session. + Long primeOwnerTicket = Option.CHANNEL_PRIME_OWNER.getLong(options); + if (channelPrimer != null + && isMultiplexed + && session.getMultiplexed() + && primeOwnerTicket != null) { + channelPrimer.registerPrimeSession(databaseName, session.getName(), primeOwnerTicket); + } + return session; + } + + @Override + public void registerChannelPrimeOwner(String databaseName, long ownerTicket) { + if (channelPrimer != null) { + channelPrimer.registerPrimeOwner(databaseName, ownerTicket); + } + } + + @Override + public void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) { + if (channelPrimer != null) { + channelPrimer.unregisterPrimeOwner(databaseName, ownerTicket); + } } @Override diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 282979b522b8..4e07cfa1d835 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -82,7 +82,14 @@ public interface SpannerRpc extends ServiceRpc { /** Options passed in {@link SpannerRpc} methods to control how an RPC is issued. */ enum Option { CHANNEL_HINT("Channel Hint"), - CHANNEL_ID_AFFINITY("Channel ID Affinity"); + CHANNEL_ID_AFFINITY("Channel ID Affinity"), + /** + * The owner ticket of the database client that issues a multiplexed {@code CreateSession}. The + * created session is used for priming channels that the dynamic channel pool adds during + * scale-up only while the ticket is the current one that was registered through {@link + * SpannerRpc#registerChannelPrimeOwner(String, long)} for the database. Internal. + */ + CHANNEL_PRIME_OWNER("Channel Prime Owner"); private final String value; @@ -376,6 +383,24 @@ default Session createSession( void deleteSession(String sessionName, @Nullable Map options) throws SpannerException; + /** + * Registers the database client with the given owner ticket as the current owner of the given + * database for priming channels that the dynamic channel pool adds during scale-up. A multiplexed + * session that is created with {@link Option#CHANNEL_PRIME_OWNER} set to the ticket is used for + * priming only while the ticket is the current owner ticket of its database. Called before the + * client issues its first multiplexed {@code CreateSession}. The default implementation does + * nothing. + */ + default void registerChannelPrimeOwner(String databaseName, long ownerTicket) {} + + /** + * Unregisters the database client with the given owner ticket as the owner of the given database + * and removes the multiplexed session that it created from the sessions that are used for priming + * channels that the dynamic channel pool adds during scale-up. Called when the client is + * invalidated or closed. The default implementation does nothing. + */ + default void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) {} + ApiFuture asyncDeleteSession(String sessionName, @Nullable Map options) throws SpannerException; diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java new file mode 100644 index 000000000000..9b48d55d42be --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java @@ -0,0 +1,387 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.spi.v1.ChannelPrimerTestRpc; +import com.google.common.util.concurrent.Uninterruptibles; +import com.google.spanner.v1.CreateSessionRequest; +import com.google.spanner.v1.Session; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import io.grpc.stub.StreamObserver; +import java.net.InetSocketAddress; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Verifies that a multiplexed CreateSession that is still in flight when its database client is + * retired never becomes the prime session of the dynamic channel pool, and that the retired client + * ignores the late session. + */ +@RunWith(JUnit4.class) +public class DynamicChannelPoolPrimeSessionTest { + private static final DatabaseId DATABASE_ID = + DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"); + private static final Statement SELECT1 = Statement.of("SELECT 1"); + private static final Duration TIMEOUT = Duration.ofSeconds(15); + + /** A mock Spanner whose CreateSession calls wait for a gate before they are served. */ + private static final class GatedMockSpanner extends MockSpannerServiceImpl { + final AtomicInteger arrivedCreateSessions = new AtomicInteger(); + volatile CountDownLatch createSessionGate = new CountDownLatch(0); + + @Override + public void createSession( + CreateSessionRequest request, StreamObserver responseObserver) { + arrivedCreateSessions.incrementAndGet(); + Uninterruptibles.awaitUninterruptibly(createSessionGate); + super.createSession(request, responseObserver); + } + + void blockCreateSessions() { + createSessionGate = new CountDownLatch(1); + } + + void releaseCreateSessions() { + createSessionGate.countDown(); + } + } + + private GatedMockSpanner mockSpanner; + private Server server; + private ChannelPrimerTestRpc rpc; + private SpannerImpl spanner; + + /** A second Spanner that a test creates with its own session pool options, closed after it. */ + private SpannerImpl waitingSpanner; + + private ChannelPrimerTestRpc waitingRpc; + + @Before + public void startServer() throws Exception { + mockSpanner = new GatedMockSpanner(); + mockSpanner.setAbortProbability(0.0D); + mockSpanner.putStatementResult( + StatementResult.query(SELECT1, MockSpannerTestUtil.SELECT1_RESULTSET)); + server = + NettyServerBuilder.forAddress(new InetSocketAddress("localhost", 0)) + .addService(mockSpanner) + .build() + .start(); + SpannerOptions options = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setEnableDirectAccess(false) + .setHost("http://localhost:" + server.getPort()) + .setCredentials(NoCredentials.getInstance()) + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + // Every multiplexed session is due for a refresh right after its creation, so + // a test can trigger the refresh by running the maintainer. The maintainer is + // only run explicitly: its scheduled run is minutes away. + .setMultiplexedSessionMaintenanceDuration(Duration.ofMillis(1)) + .build()) + .build(); + rpc = new ChannelPrimerTestRpc(options); + spanner = new SpannerImpl(rpc, options); + } + + @After + public void stopServer() throws Exception { + mockSpanner.releaseCreateSessions(); + if (waitingSpanner != null && !waitingSpanner.isClosed()) { + waitingSpanner.close(); + } + if (!spanner.isClosed()) { + spanner.close(); + } + server.shutdown(); + server.awaitTermination(); + } + + /** + * Creates a second Spanner instance that waits for the given duration for the multiplexed session + * of every database client, and whose maintainer runs every ten milliseconds once it is started. + */ + private void createWaitingSpanner(Duration waitForMinSessions) { + SpannerOptions options = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setEnableDirectAccess(false) + .setHost("http://localhost:" + server.getPort()) + .setCredentials(NoCredentials.getInstance()) + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + // A started maintainer refreshes the session almost immediately, so a test can + // detect a maintainer that should never have been started. + .setMultiplexedSessionMaintenanceDuration(Duration.ofMillis(1)) + .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMillis(10)) + .setWaitForMinSessionsDuration(waitForMinSessions) + .build()) + .build(); + waitingRpc = new ChannelPrimerTestRpc(options); + waitingSpanner = new SpannerImpl(waitingRpc, options); + } + + private static void awaitCondition(BooleanSupplier condition) throws InterruptedException { + long deadline = System.nanoTime() + TIMEOUT.toNanos(); + while (!condition.getAsBoolean()) { + if (System.nanoTime() > deadline) { + throw new AssertionError("Condition not met within " + TIMEOUT); + } + Thread.sleep(5L); + } + } + + private DatabaseClientImpl getDatabaseClient() { + return (DatabaseClientImpl) spanner.getDatabaseClient(DATABASE_ID); + } + + private static MultiplexedSessionDatabaseClient multiplexedClient(DatabaseClientImpl client) { + return client.multiplexedSessionDatabaseClient; + } + + /** Creates a database client, waits for its multiplexed session, and returns the client. */ + private DatabaseClientImpl createClientWithPrimeSession() throws Exception { + DatabaseClientImpl client = getDatabaseClient(); + awaitCondition(() -> rpc.getPrimeSessionName() != null); + String session = multiplexedClient(client).getCurrentSessionReference().getName(); + assertThat(rpc.getPrimeSessionNames()).containsExactly(session); + return client; + } + + /** Blocks CreateSession on the server and triggers a refresh of the client's session. */ + private void blockRefreshInFlight(DatabaseClientImpl client) throws Exception { + int arrivedBefore = mockSpanner.arrivedCreateSessions.get(); + mockSpanner.blockCreateSessions(); + // The session is due for a refresh once its one millisecond maintenance duration has passed. + Thread.sleep(10L); + multiplexedClient(client).getMaintainer().maintain(); + awaitCondition(() -> mockSpanner.arrivedCreateSessions.get() == arrivedBefore + 1); + } + + /** Releases the blocked CreateSession calls and waits until the given number have returned. */ + private void releaseAndAwaitCompleted(int expectedCompleted) throws Exception { + mockSpanner.releaseCreateSessions(); + awaitCondition(() -> rpc.getCompletedMultiplexedCreateSessions() >= expectedCompleted); + assertEquals(expectedCompleted, rpc.getCompletedMultiplexedCreateSessions()); + } + + private void invalidate(DatabaseClientImpl client) { + invalidate(client, DATABASE_ID); + } + + private void invalidate(DatabaseClientImpl client, DatabaseId databaseId) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.stickyDatabaseNotFoundException(databaseId.getName())); + assertThrows( + DatabaseNotFoundException.class, + () -> { + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1)) { + resultSet.next(); + } + }); + mockSpanner.removeAllExecutionTimes(); + assertFalse(client.isValid()); + } + + @Test + public void refreshInFlightWhenClientIsReplacedNeverBecomesPrimeSession() throws Exception { + DatabaseClientImpl retired = createClientWithPrimeSession(); + String retiredSession = multiplexedClient(retired).getCurrentSessionReference().getName(); + invalidate(retired); + blockRefreshInFlight(retired); + + // Getting the client again retires the invalid client and creates its replacement. The + // replacement's CreateSession is blocked as well and returns together with the refresh. + DatabaseClientImpl replacement = getDatabaseClient(); + assertNotSame(retired, replacement); + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + releaseAndAwaitCompleted(3); + + String replacementSession = + multiplexedClient(replacement).getCurrentSessionReference().getName(); + assertThat(replacementSession).isNotEqualTo(retiredSession); + // The late refresh of the retired client is dropped, and the replacement stays selected. + assertThat(rpc.getPrimeSessionNames()).containsExactly(replacementSession); + assertThat(rpc.getPrimeOwnerDatabases()).containsExactly(DATABASE_ID.getName()); + // The retired client ignored its refreshed session as well. + assertEquals(retiredSession, multiplexedClient(retired).getCurrentSessionReference().getName()); + assertTrue(replacement.isValid()); + } + + @Test + public void refreshInFlightWhenSpannerIsClosedNeverBecomesPrimeSession() throws Exception { + DatabaseClientImpl client = createClientWithPrimeSession(); + String session = multiplexedClient(client).getCurrentSessionReference().getName(); + blockRefreshInFlight(client); + // The refresh returns right after the database has been unregistered, while the rpc is still + // open, which is the latest point at which it could re-register the retired session. + rpc.setAfterUnregisterHook( + () -> { + try { + releaseAndAwaitCompleted(2); + } catch (Exception e) { + throw new AssertionError(e); + } + }); + + spanner.close(); + + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + assertEquals(session, multiplexedClient(client).getCurrentSessionReference().getName()); + } + + @Test + public void initialCreateSessionInFlightWhenSpannerIsClosedNeverBecomesPrimeSession() + throws Exception { + mockSpanner.blockCreateSessions(); + DatabaseClientImpl client = getDatabaseClient(); + awaitCondition(() -> mockSpanner.arrivedCreateSessions.get() == 1); + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + rpc.setAfterUnregisterHook( + () -> { + try { + releaseAndAwaitCompleted(1); + } catch (Exception e) { + throw new AssertionError(e); + } + }); + + spanner.close(); + + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + // The closed client ignored the session that arrived after it was closed. + SpannerException exception = + assertThrows( + SpannerException.class, () -> multiplexedClient(client).getCurrentSessionReference()); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + + @Test + public void churnOfDatabaseClientsLeavesNoPrimeOwnersBehind() throws Exception { + // A long-lived Spanner instance that churns through many database names must not retain + // anything for the retired database clients: the registry holds exactly one owner and one + // session per live database client. + List liveDatabases = new ArrayList<>(); + for (int i = 0; i < 20; i++) { + DatabaseId databaseId = DatabaseId.of("[PROJECT]", "[INSTANCE]", "churn-" + i); + liveDatabases.add(databaseId.getName()); + DatabaseClientImpl client = (DatabaseClientImpl) spanner.getDatabaseClient(databaseId); + String session = multiplexedClient(client).getCurrentSessionReference().getName(); + awaitCondition(() -> rpc.getPrimeSessionNames().contains(session)); + invalidate(client, databaseId); + // Getting the client again retires it and creates its replacement. + DatabaseClientImpl replacement = (DatabaseClientImpl) spanner.getDatabaseClient(databaseId); + assertNotSame(client, replacement); + String replacementSession = + multiplexedClient(replacement).getCurrentSessionReference().getName(); + awaitCondition(() -> rpc.getPrimeSessionNames().contains(replacementSession)); + assertThat(rpc.getPrimeSessionNames()).doesNotContain(session); + assertThat(rpc.getPrimeSessionNames()).hasSize(liveDatabases.size()); + assertThat(rpc.getPrimeOwnerDatabases()).containsExactlyElementsIn(liveDatabases); + } + + spanner.close(); + + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + } + + @Test + public void failedClientConstructionLeavesNoPrimeOwnerBehind() throws Exception { + // A client that waits for its first multiplexed session throws from its constructor when the + // CreateSession fails, and is never cached by the Spanner instance, so nothing will ever close + // it. It must still have released the owner ticket that it registered before the CreateSession. + createWaitingSpanner(Duration.ofSeconds(5)); + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofStickyException( + Status.PERMISSION_DENIED.withDescription("no access").asRuntimeException())); + + for (int i = 0; i < 5; i++) { + DatabaseId databaseId = DatabaseId.of("[PROJECT]", "[INSTANCE]", "failing-" + i); + SpannerException exception = + assertThrows(SpannerException.class, () -> waitingSpanner.getDatabaseClient(databaseId)); + assertEquals(ErrorCode.PERMISSION_DENIED, exception.getErrorCode()); + assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(waitingRpc.getPrimeSessionNames()).isEmpty(); + } + + // The maintainer of a client that never finished its construction is not running: it would + // have refreshed the session within its ten millisecond loop frequency. + mockSpanner.removeAllExecutionTimes(); + int arrived = mockSpanner.arrivedCreateSessions.get(); + Thread.sleep(200L); + assertEquals(arrived, mockSpanner.arrivedCreateSessions.get()); + } + + @Test + public void sessionArrivingAfterFailedClientConstructionNeverBecomesPrimeSession() + throws Exception { + // The CreateSession outlives the wait for the first multiplexed session, so the constructor + // throws while the CreateSession is still in flight. Its late success must not be recorded as + // a prime session, and must not start the maintainer of a client that does not exist. + createWaitingSpanner(Duration.ofMillis(200)); + mockSpanner.blockCreateSessions(); + DatabaseId databaseId = DatabaseId.of("[PROJECT]", "[INSTANCE]", "abandoned"); + + SpannerException exception = + assertThrows(SpannerException.class, () -> waitingSpanner.getDatabaseClient(databaseId)); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode()); + // The owner ticket is released as soon as the constructor fails, before the session arrives. + assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + + int arrivedBeforeRelease = mockSpanner.arrivedCreateSessions.get(); + mockSpanner.releaseCreateSessions(); + awaitCondition(() -> waitingRpc.getCompletedMultiplexedCreateSessions() >= 1); + + assertThat(waitingRpc.getPrimeSessionNames()).isEmpty(); + assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + // No maintainer was started for the abandoned client, so no session is ever refreshed. + Thread.sleep(200L); + assertEquals(arrivedBeforeRelease, mockSpanner.arrivedCreateSessions.get()); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index d398a78643aa..8a99746084de 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -19,39 +19,54 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeFalse; import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; import com.google.cloud.NoCredentials; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionConsumer; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import java.io.PrintWriter; import java.io.StringWriter; import java.lang.reflect.Field; import java.time.Clock; import java.time.Duration; import java.time.Instant; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; import org.junit.After; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; import org.mockito.stubbing.Answer; @RunWith(JUnit4.class) public class MultiplexedSessionDatabaseClientTest { + private static final DatabaseId TEST_DATABASE_ID = + DatabaseId.of("test-project", "test-instance", "test-database"); + @After public void tearDown() throws Exception { clearChannelUsage(); @@ -70,6 +85,8 @@ public void testMaintainer() { SpannerOptions spannerOptions = mock(SpannerOptions.class); SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class); when(sessionClient.getSpanner()).thenReturn(spanner); + when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); + when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class)); when(spanner.getOptions()).thenReturn(spannerOptions); when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()) @@ -102,7 +119,7 @@ public void testMaintainer() { return null; }) .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class)); + .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); // Create a client. This should get session1. MultiplexedSessionDatabaseClient client = @@ -126,6 +143,110 @@ public void testMaintainer() { assertEquals(client.getCurrentSessionReference(), session2.getSessionReference()); } + @Test + public void testClosedClientIgnoresInitialSessionThatArrivesAfterClose() { + Clock clock = mock(Clock.class); + when(clock.instant()).thenReturn(Instant.now()); + SessionClient sessionClient = mock(SessionClient.class); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class); + when(sessionClient.getSpanner()).thenReturn(spanner); + when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); + when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class)); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()) + .thenReturn(Duration.ofDays(7)); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency()) + .thenReturn(Duration.ofMinutes(10)); + when(sessionPoolOptions.getWaitForMinSessions()).thenReturn(Duration.ZERO); + + SessionImpl session = mock(SessionImpl.class); + when(session.getSessionReference()).thenReturn(mock(SessionReference.class)); + // Capture the consumer of the initial session without delivering a session yet. + AtomicReference consumer = new AtomicReference<>(); + doAnswer( + (Answer) + invocationOnMock -> { + consumer.set(invocationOnMock.getArgument(0)); + return null; + }) + .when(sessionClient) + .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + assertNotNull(consumer.get()); + + // The client is closed while its initial session is still being created. + client.close(); + consumer.get().onSessionReady(session); + + // The late session is not handed to waiters of the closed client, and the maintainer is not + // started for it. + SpannerException exception = + assertThrows(SpannerException.class, client::getCurrentSessionReference); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + verify(sessionPoolOptions, never()).getMultiplexedSessionMaintenanceLoopFrequency(); + } + + @Test + public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() { + Instant now = Instant.now(); + Clock clock = mock(Clock.class); + when(clock.instant()).thenReturn(now); + SessionClient sessionClient = mock(SessionClient.class); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class); + when(sessionClient.getSpanner()).thenReturn(spanner); + when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); + when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class)); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()) + .thenReturn(Duration.ofDays(7)); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency()) + .thenReturn(Duration.ofMinutes(10)); + + SessionImpl session1 = mock(SessionImpl.class); + SessionReference sessionReference1 = mock(SessionReference.class); + when(session1.getSessionReference()).thenReturn(sessionReference1); + SessionImpl session2 = mock(SessionImpl.class); + when(session2.getSessionReference()).thenReturn(mock(SessionReference.class)); + + // Deliver the initial session immediately, but capture the consumer of the refresh. + AtomicReference refreshConsumer = new AtomicReference<>(); + doAnswer( + (Answer) + invocationOnMock -> { + SessionConsumer consumer = invocationOnMock.getArgument(0); + consumer.onSessionReady(session1); + return null; + }) + .doAnswer( + (Answer) + invocationOnMock -> { + refreshConsumer.set(invocationOnMock.getArgument(0)); + return null; + }) + .when(sessionClient) + .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + assertEquals(sessionReference1, client.getCurrentSessionReference()); + + // The session is due for a refresh, and the refresh is in flight when the client is closed. + when(clock.instant()).thenReturn(now.plus(Duration.ofDays(8))); + client.getMaintainer().maintain(); + assertNotNull(refreshConsumer.get()); + client.close(); + refreshConsumer.get().onSessionReady(session2); + + // The refreshed session of the closed client is ignored. + assertEquals(sessionReference1, client.getCurrentSessionReference()); + } + @Test public void testDisableMultiplexedSessionEnvVar() throws Exception { assumeTrue(isJava8() && !isWindows()); @@ -270,6 +391,8 @@ public void testGrpcGcpSingleUseDoesNotReserveBitsetChannelHint() throws Excepti ISpan span = mock(ISpan.class); when(sessionClient.getSpanner()).thenReturn(spanner); + when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); + when(spanner.getRpc()).thenReturn(mock(SpannerRpc.class)); when(spanner.getOptions()).thenReturn(spannerOptions); when(spanner.getTracer()).thenReturn(tracer); when(tracer.getCurrentSpan()).thenReturn(span); @@ -348,6 +471,80 @@ public void testCloseKeepsChannelUsageEntryWhileAnotherClientIsUsingSameSpanner( } } + @Test + public void testChannelPrimeOwnerTicketIsRegisteredCarriedAndUnregistered() { + Clock clock = mock(Clock.class); + when(clock.instant()).thenReturn(Instant.now()); + SessionClient sessionClient = mock(SessionClient.class); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerRpc rpc = mock(SpannerRpc.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class); + when(sessionClient.getSpanner()).thenReturn(spanner); + when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); + when(spanner.getRpc()).thenReturn(rpc); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceDuration()) + .thenReturn(Duration.ofDays(7)); + when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency()) + .thenReturn(Duration.ofMinutes(10)); + SessionImpl session = mock(SessionImpl.class); + when(session.getSessionReference()).thenReturn(mock(SessionReference.class)); + // Capture the options of every CreateSession without delivering a session yet. + List> createSessionOptions = new ArrayList<>(); + AtomicReference consumer = new AtomicReference<>(); + doAnswer( + (Answer) + invocationOnMock -> { + consumer.set(invocationOnMock.getArgument(0)); + createSessionOptions.add(invocationOnMock.getArgument(1)); + return null; + }) + .when(sessionClient) + .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + + MultiplexedSessionDatabaseClient first = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + MultiplexedSessionDatabaseClient second = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + + // Every client registers itself as an owner with its own ticket before its first + // CreateSession, and that CreateSession carries the ticket. + ArgumentCaptor tickets = ArgumentCaptor.forClass(Long.class); + verify(rpc, times(2)) + .registerChannelPrimeOwner(eq(TEST_DATABASE_ID.getName()), tickets.capture()); + long firstTicket = tickets.getAllValues().get(0); + long secondTicket = tickets.getAllValues().get(1); + assertNotEquals(firstTicket, secondTicket); + assertEquals(2, createSessionOptions.size()); + assertEquals( + Long.valueOf(firstTicket), + SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(0))); + assertEquals( + Long.valueOf(secondTicket), + SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(1))); + verify(rpc, never()).unregisterChannelPrimeOwner(any(), anyLong()); + + // A refresh of the session carries the same ticket as the initial CreateSession. + consumer.get().onSessionReady(session); + when(clock.instant()).thenReturn(Instant.now().plus(Duration.ofDays(8))); + second.getMaintainer().maintain(); + assertEquals(3, createSessionOptions.size()); + assertEquals( + Long.valueOf(secondTicket), + SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(2))); + + // Closing a client unregisters exactly its own ticket, once. + first.close(); + verify(rpc).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), firstTicket); + verify(rpc, never()).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), secondTicket); + first.close(); + verify(rpc, times(1)).unregisterChannelPrimeOwner(any(), anyLong()); + second.close(); + verify(rpc).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), secondTicket); + } + private SessionClient createSessionClient(SpannerImpl spanner) { return new FailingMultiplexedSessionClient(spanner); } @@ -406,15 +603,13 @@ public void release(ScheduledExecutorService executor) { } private static final class FailingMultiplexedSessionClient extends SessionClient { - private static final DatabaseId TEST_DATABASE_ID = - DatabaseId.of("test-project", "test-instance", "test-database"); - private FailingMultiplexedSessionClient(SpannerImpl spanner) { super(spanner, TEST_DATABASE_ID, new TestExecutorFactory()); } @Override - void asyncCreateMultiplexedSession(SessionConsumer consumer) { + void asyncCreateMultiplexedSession( + SessionConsumer consumer, @Nullable Map options) { consumer.onSessionCreateFailure( SpannerExceptionFactory.newSpannerException(ErrorCode.UNAUTHENTICATED, "test"), 1); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java index dbcfa7f03fe5..cc919756b581 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java @@ -21,6 +21,11 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -39,12 +44,14 @@ import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.Attributes; import java.io.IOException; import java.io.PrintWriter; import java.io.StringWriter; import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.After; import org.junit.Before; import org.junit.BeforeClass; @@ -126,6 +133,77 @@ public void getDbclientAgainGivesSame() { verify(spannerOptions).initializeBuiltInMetrics(db); } + @Test + public void invalidatedDatabaseClientUnregistersItsChannelPrimeOwnerTicket() { + DatabaseId db = DatabaseId.of("projects/p1/instances/i1/databases/d1"); + Mockito.when(spannerOptions.getTransportOptions()) + .thenReturn(GrpcTransportOptions.newBuilder().build()); + Mockito.when(spannerOptions.getSessionPoolOptions()) + .thenReturn(SessionPoolOptions.newBuilder().setMinSessions(0).build()); + AtomicBoolean valid = new AtomicBoolean(true); + SpannerImpl spanner = + new SpannerImpl(rpc, spannerOptions) { + @Override + DatabaseClientImpl createDatabaseClient( + String clientId, + MultiplexedSessionDatabaseClient multiplexedSessionClient, + Attributes databaseAttributes) { + return new DatabaseClientImpl( + clientId, multiplexedSessionClient, tracer, databaseAttributes) { + @Override + boolean isValid() { + return valid.get(); + } + }; + } + }; + try { + DatabaseClient first = spanner.getDatabaseClient(db); + ArgumentCaptor tickets = ArgumentCaptor.forClass(Long.class); + verify(rpc).registerChannelPrimeOwner(eq(db.getName()), tickets.capture()); + long firstTicket = tickets.getValue(); + verify(rpc, never()).unregisterChannelPrimeOwner(anyString(), anyLong()); + + // The invalidated client is replaced. Closing it unregisters its ticket, and the + // replacement registers a new ticket of its own. + valid.set(false); + DatabaseClient second = spanner.getDatabaseClient(db); + + assertThat(second).isNotSameInstanceAs(first); + verify(rpc).unregisterChannelPrimeOwner(db.getName(), firstTicket); + verify(rpc, times(2)).registerChannelPrimeOwner(eq(db.getName()), tickets.capture()); + long secondTicket = tickets.getValue(); + assertThat(secondTicket).isNotEqualTo(firstTicket); + verify(rpc, never()).unregisterChannelPrimeOwner(db.getName(), secondTicket); + } finally { + spanner.close(); + } + } + + @Test + public void closeUnregistersChannelPrimeOwnerTicketOfEveryDatabaseClient() { + DatabaseId db1 = DatabaseId.of("projects/p1/instances/i1/databases/d1"); + DatabaseId db2 = DatabaseId.of("projects/p1/instances/i1/databases/d2"); + Mockito.when(spannerOptions.getTransportOptions()) + .thenReturn(GrpcTransportOptions.newBuilder().build()); + Mockito.when(spannerOptions.getSessionPoolOptions()) + .thenReturn(SessionPoolOptions.newBuilder().setMinSessions(0).build()); + SpannerImpl spanner = new SpannerImpl(rpc, spannerOptions); + spanner.getDatabaseClient(db1); + spanner.getDatabaseClient(db2); + ArgumentCaptor ticket1 = ArgumentCaptor.forClass(Long.class); + ArgumentCaptor ticket2 = ArgumentCaptor.forClass(Long.class); + verify(rpc).registerChannelPrimeOwner(eq(db1.getName()), ticket1.capture()); + verify(rpc).registerChannelPrimeOwner(eq(db2.getName()), ticket2.capture()); + verify(rpc, never()).unregisterChannelPrimeOwner(anyString(), anyLong()); + + spanner.close(); + + verify(rpc).unregisterChannelPrimeOwner(db1.getName(), ticket1.getValue()); + verify(rpc).unregisterChannelPrimeOwner(db2.getName(), ticket2.getValue()); + verify(rpc, times(2)).unregisterChannelPrimeOwner(anyString(), anyLong()); + } + @Test public void queryOptions() { QueryOptions queryOptions = diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java index bdcdf89afc2d..25813260944f 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java @@ -39,6 +39,7 @@ import com.google.cloud.NoCredentials; import com.google.cloud.ServiceOptions; import com.google.cloud.TransportOptions; +import com.google.cloud.grpc.GcpChannelPrimer; import com.google.cloud.grpc.GcpManagedChannelOptions.GcpChannelPoolOptions; import com.google.cloud.spanner.SpannerOptions.Builder.DefaultReadWriteTransactionOptions; import com.google.cloud.spanner.SpannerOptions.FixedCloseableExecutorProvider; @@ -48,6 +49,7 @@ import com.google.cloud.spanner.omni.SpannerOmniCredentials; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.base.Strings; +import com.google.common.util.concurrent.Futures; import com.google.spanner.v1.BatchCreateSessionsRequest; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; @@ -1390,6 +1392,55 @@ public void testDynamicChannelPoolDefaultValues() { assertEquals(2, SpannerOptions.DEFAULT_DYNAMIC_POOL_MIN_CHANNELS); assertEquals(Duration.ofMinutes(10), SpannerOptions.DEFAULT_DYNAMIC_POOL_AFFINITY_KEY_LIFETIME); assertEquals(Duration.ofMinutes(1), SpannerOptions.DEFAULT_DYNAMIC_POOL_CLEANUP_INTERVAL); + assertEquals(Duration.ofSeconds(10), SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT); + assertEquals(3, SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS); + } + + @Test + public void testDynamicChannelPoolPrimeSettingsDefaultsAndUserValues() { + GcpChannelPoolOptions defaults = SpannerOptions.createDefaultDynamicChannelPoolOptions(); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT, + defaults.getChannelPrimeTimeout()); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS, + defaults.getChannelPrimeMaxAttempts()); + // The options themselves never register a primer; the Spanner client does that when the pool + // is created, unless the user supplied one. + assertNull(defaults.getChannelPrimer()); + + // Prime settings that the user left unset are filled in from the defaults. + SpannerOptions merged = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .enableDynamicChannelPool() + .setGcpChannelPoolOptions(GcpChannelPoolOptions.newBuilder().setMaxSize(6).build()) + .build(); + assertEquals(6, merged.getGcpChannelPoolOptions().getMaxSize()); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT, + merged.getGcpChannelPoolOptions().getChannelPrimeTimeout()); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS, + merged.getGcpChannelPoolOptions().getChannelPrimeMaxAttempts()); + assertNull(merged.getGcpChannelPoolOptions().getChannelPrimer()); + + // A user-provided primer, timeout, and attempt count are preserved. + GcpChannelPrimer userPrimer = channel -> Futures.immediateFuture(null); + SpannerOptions custom = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setChannelPrimer(userPrimer) + .setChannelPrimeTimeout(Duration.ofSeconds(2)) + .setChannelPrimeMaxAttempts(5) + .build()) + .build(); + assertSame(userPrimer, custom.getGcpChannelPoolOptions().getChannelPrimer()); + assertEquals(Duration.ofSeconds(2), custom.getGcpChannelPoolOptions().getChannelPrimeTimeout()); + assertEquals(5, custom.getGcpChannelPoolOptions().getChannelPrimeMaxAttempts()); } @Test diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java new file mode 100644 index 000000000000..71bd66b11756 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java @@ -0,0 +1,99 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.spi.v1; + +import com.google.cloud.spanner.SpannerOptions; +import com.google.cloud.spanner.spi.v1.DynamicChannelPoolPrimer.PrimeSession; +import com.google.spanner.v1.Session; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; + +/** + * A {@link GapicSpannerRpc} for tests outside this package that need to observe the prime sessions + * and prime owners of the dynamic channel pool primer and to hook into the retirement of a database + * client. + */ +public class ChannelPrimerTestRpc extends GapicSpannerRpc { + private final AtomicInteger completedMultiplexedCreateSessions = new AtomicInteger(); + private volatile Runnable afterUnregisterHook = () -> {}; + + public ChannelPrimerTestRpc(SpannerOptions options) { + super(options); + } + + /** Sets a hook that runs right after a database client has been unregistered as prime owner. */ + public void setAfterUnregisterHook(Runnable hook) { + this.afterUnregisterHook = hook; + } + + /** Returns the number of multiplexed CreateSession calls that have returned, with any outcome. */ + public int getCompletedMultiplexedCreateSessions() { + return completedMultiplexedCreateSessions.get(); + } + + /** Returns the session that the next priming attempt uses, or {@code null}. */ + @Nullable + public String getPrimeSessionName() { + DynamicChannelPoolPrimer primer = getChannelPrimer(); + return primer == null ? null : primer.getPrimeSessionName(); + } + + /** Returns the names of all registered prime sessions, most recently created first. */ + public List getPrimeSessionNames() { + List names = new ArrayList<>(); + DynamicChannelPoolPrimer primer = getChannelPrimer(); + if (primer != null) { + for (PrimeSession entry : primer.getPrimeSessions()) { + names.add(entry.getSessionName()); + } + } + return names; + } + + @Override + public Session createSession( + String databaseName, + @Nullable String databaseRole, + @Nullable Map labels, + @Nullable Map options, + boolean isMultiplexed) { + try { + return super.createSession(databaseName, databaseRole, labels, options, isMultiplexed); + } finally { + if (isMultiplexed) { + completedMultiplexedCreateSessions.incrementAndGet(); + } + } + } + + /** Returns the names of the databases that currently have a registered prime owner. */ + public Set getPrimeOwnerDatabases() { + DynamicChannelPoolPrimer primer = getChannelPrimer(); + return primer == null ? Collections.emptySet() : primer.getPrimeOwners().keySet(); + } + + @Override + public void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) { + super.unregisterChannelPrimeOwner(databaseName, ownerTicket); + afterUnregisterHook.run(); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java new file mode 100644 index 000000000000..fccbaa4a14bd --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java @@ -0,0 +1,998 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.spi.v1; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.auth.oauth2.AccessToken; +import com.google.auth.oauth2.OAuth2Credentials; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; +import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.cloud.spanner.spi.v1.DynamicChannelPoolPrimer.PrimeSession; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ResultSet; +import com.google.spanner.v1.SpannerGrpc; +import io.grpc.CallCredentials; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.Context; +import io.grpc.Contexts; +import io.grpc.ForwardingClientCall.SimpleForwardingClientCall; +import io.grpc.ManagedChannel; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Server; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.auth.MoreCallCredentials; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.ServerCallStreamObserver; +import io.grpc.stub.StreamObserver; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import javax.annotation.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class DynamicChannelPoolPrimerTest { + private static final String PROJECT_NAME = "projects/my-project"; + private static final String DATABASE_NAME = + "projects/my-project/instances/my-instance/databases/my-database"; + private static final String OTHER_DATABASE_NAME = + "projects/my-project/instances/my-instance/databases/other-database"; + private static final String SESSION_NAME = DATABASE_NAME + "/sessions/multiplexed-session"; + private static final String REFRESHED_SESSION_NAME = + DATABASE_NAME + "/sessions/refreshed-session"; + private static final String OTHER_SESSION_NAME = OTHER_DATABASE_NAME + "/sessions/other-session"; + private static final String RESOURCE_HEADER_KEY = "google-cloud-resource-prefix"; + private static final String DEFAULT_TOKEN = "default-token"; + private static final String PROVIDER_TOKEN = "provider-token"; + + private static final Metadata.Key AUTHORIZATION_KEY = + Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key API_CLIENT_KEY = + Metadata.Key.of("x-goog-api-client", Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key RESOURCE_PREFIX_KEY = + Metadata.Key.of(RESOURCE_HEADER_KEY, Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key REQUEST_PARAMS_KEY = + Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key ROUTE_TO_LEADER_KEY = + Metadata.Key.of("x-goog-spanner-route-to-leader", Metadata.ASCII_STRING_MARSHALLER); + private static final Metadata.Key REQUEST_ID_KEY = + Metadata.Key.of("x-goog-spanner-request-id", Metadata.ASCII_STRING_MARSHALLER); + + /** Minimal Spanner service that only serves ExecuteSql and can hold or fail calls. */ + private static final class PrimeService extends SpannerGrpc.SpannerImplBase { + final List requests = new CopyOnWriteArrayList<>(); + final List headers = new CopyOnWriteArrayList<>(); + volatile boolean holdResponses; + volatile CountDownLatch callStarted = new CountDownLatch(1); + volatile CountDownLatch callCancelled = new CountDownLatch(1); + @Nullable volatile Status failWith; + + /** Failures for specific sessions, which take precedence over {@link #failWith}. */ + final Map failSessionsWith = new ConcurrentHashMap<>(); + + @Override + public void executeSql(ExecuteSqlRequest request, StreamObserver responseObserver) { + requests.add(request); + ServerCallStreamObserver serverObserver = + (ServerCallStreamObserver) responseObserver; + serverObserver.setOnCancelHandler(() -> callCancelled.countDown()); + callStarted.countDown(); + if (holdResponses) { + return; + } + Status failure = failSessionsWith.getOrDefault(request.getSession(), failWith); + if (failure != null) { + responseObserver.onError(failure.asRuntimeException()); + return; + } + responseObserver.onNext(ResultSet.getDefaultInstance()); + responseObserver.onCompleted(); + } + } + + private PrimeService service; + private Server server; + private String serverNameForChannels; + private final List channels = new ArrayList<>(); + private final RequestIdCreatorImpl requestIdCreator = new RequestIdCreatorImpl(); + + @Before + public void setUp() throws Exception { + service = new PrimeService(); + String serverName = InProcessServerBuilder.generateName(); + server = + InProcessServerBuilder.forName(serverName) + .addService(service) + .intercept( + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, + Metadata headers, + ServerCallHandler next) { + service.headers.add(headers); + return Contexts.interceptCall(Context.current(), call, headers, next); + } + }) + .build() + .start(); + serverNameForChannels = serverName; + } + + @After + public void tearDown() throws Exception { + for (ManagedChannel channel : channels) { + channel.shutdownNow(); + } + server.shutdownNow(); + server.awaitTermination(5, TimeUnit.SECONDS); + } + + private ManagedChannel newChannel() { + return newChannel(InProcessChannelBuilder.forName(serverNameForChannels)); + } + + private ManagedChannel newChannel(InProcessChannelBuilder builder) { + ManagedChannel channel = builder.build(); + channels.add(channel); + return channel; + } + + /** + * Returns a channel that carries what the delegate channel of the dynamic channel pool carries: + * the fixed headers of the client through the GAX header interceptor, and the request id + * interceptor of the client. + */ + private ManagedChannel newDelegateLikeChannel() { + Metadata fixedHeaders = new Metadata(); + fixedHeaders.put(API_CLIENT_KEY, "test-client"); + return newChannel( + InProcessChannelBuilder.forName(serverNameForChannels) + .intercept( + new RequestIdInterceptor(), + new ClientInterceptor() { + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + return new SimpleForwardingClientCall( + next.newCall(method, callOptions)) { + @Override + public void start(Listener responseListener, Metadata headers) { + headers.merge(fixedHeaders); + super.start(responseListener, headers); + } + }; + } + })); + } + + private static CallCredentials credentialsWithToken(String token) { + return MoreCallCredentials.from( + OAuth2Credentials.create( + new AccessToken(token, new Date(System.currentTimeMillis() + 3_600_000L)))); + } + + private DynamicChannelPoolPrimer newPrimer( + @Nullable CallCredentialsProvider callCredentialsProvider, + boolean routeToLeader, + boolean multiplexedSessionsEnabled, + Duration rpcDeadline) { + // The fixed headers of the client are configured on the channel and must not be sent by the + // primer: it only sends the headers that a normal call adds per call. + SpannerMetadataProvider metadataProvider = + SpannerMetadataProvider.create( + ImmutableMap.of("x-goog-api-client", "test-client", "user-agent", "test-agent"), + RESOURCE_HEADER_KEY); + return new DynamicChannelPoolPrimer( + metadataProvider, + PROJECT_NAME, + requestIdCreator, + credentialsWithToken(DEFAULT_TOKEN), + callCredentialsProvider, + routeToLeader, + multiplexedSessionsEnabled, + rpcDeadline); + } + + private DynamicChannelPoolPrimer newPrimer() { + return newPrimer( + /* callCredentialsProvider= */ null, + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ true, + Duration.ofSeconds(5)); + } + + private static final AtomicLong OWNER_TICKETS = new AtomicLong(); + + /** + * Registers the session under the current owner of the database, registering a fresh owner first + * if the database has none, like a database client that creates or refreshes its session. + */ + private static long registerPrimeSession( + DynamicChannelPoolPrimer primer, String databaseName, String sessionName) { + Long ownerTicket = primer.getPrimeOwners().get(databaseName); + if (ownerTicket == null) { + ownerTicket = OWNER_TICKETS.incrementAndGet(); + primer.registerPrimeOwner(databaseName, ownerTicket); + } + assertThat(primer.registerPrimeSession(databaseName, sessionName, ownerTicket)).isTrue(); + return ownerTicket; + } + + private static T getWithin(ListenableFuture future, Duration timeout) throws Exception { + return future.get(timeout.toMillis(), TimeUnit.MILLISECONDS); + } + + private static Throwable failureOf(ListenableFuture future) throws Exception { + ExecutionException exception = + assertThrows(ExecutionException.class, () -> getWithin(future, Duration.ofSeconds(10))); + return exception.getCause(); + } + + /** Returns the request id header of the given call, asserting that there is exactly one. */ + private static XGoogSpannerRequestId requestIdOf(Metadata headers) { + List values = ImmutableList.copyOf(headers.getAll(REQUEST_ID_KEY)); + assertThat(values).hasSize(1); + return XGoogSpannerRequestId.of(values.get(0)); + } + + /** Asserts the per-call headers that every priming RPC carries. */ + private void assertPrimeHeaders(Metadata headers, String expectedToken) throws Exception { + assertThat(headers.getAll(AUTHORIZATION_KEY)).containsExactly("Bearer " + expectedToken); + assertThat(headers.getAll(RESOURCE_PREFIX_KEY)).containsExactly(DATABASE_NAME); + assertThat(headers.getAll(REQUEST_PARAMS_KEY)) + .containsExactly("session=" + urlEncode(SESSION_NAME)); + XGoogSpannerRequestId requestId = requestIdOf(headers); + assertThat(requestId.getHeaderValue()) + .isEqualTo( + XGoogSpannerRequestId.of( + requestIdCreator.getClientId(), + DynamicChannelPoolPrimer.REQUEST_ID_CHANNEL, + requestNumberOf(requestId), + /* attempt= */ 1) + .getHeaderValue()); + } + + /** Returns the request number of the given request id. */ + private static long requestNumberOf(XGoogSpannerRequestId requestId) { + // ..... + String[] parts = requestId.getHeaderValue().split("\\."); + assertThat(parts).hasLength(6); + return Long.parseLong(parts[4]); + } + + private static String urlEncode(String value) throws Exception { + return java.net.URLEncoder.encode(value, "UTF-8"); + } + + @Test + public void primeExecutesSelectOneWithSessionAndHeaders() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + ListenableFuture future = primer.prime(newChannel()); + + assertThat(getWithin(future, Duration.ofSeconds(10))).isNull(); + assertThat(service.requests).hasSize(1); + ExecuteSqlRequest request = service.requests.get(0); + assertThat(request.getSession()).isEqualTo(SESSION_NAME); + assertThat(request.getSql()).isEqualTo("SELECT 1"); + // A single-use read-only query: no transaction selector. + assertThat(request.hasTransaction()).isFalse(); + + assertThat(service.headers).hasSize(1); + Metadata headers = service.headers.get(0); + assertPrimeHeaders(headers, DEFAULT_TOKEN); + assertThat(headers.getAll(ROUTE_TO_LEADER_KEY)).containsExactly("true"); + // The fixed headers of the client are carried by the delegate channel itself and must not be + // sent by the primer, or the delegate would send them twice. + assertThat(headers.containsKey(API_CLIENT_KEY)).isFalse(); + } + + @Test + public void primeSendsEveryHeaderOnceOnDelegateThatCarriesFixedHeaders() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + getWithin(primer.prime(newDelegateLikeChannel()), Duration.ofSeconds(10)); + + Metadata headers = service.headers.get(0); + assertPrimeHeaders(headers, DEFAULT_TOKEN); + assertThat(headers.getAll(API_CLIENT_KEY)).containsExactly("test-client"); + assertThat(headers.getAll(ROUTE_TO_LEADER_KEY)).containsExactly("true"); + } + + @Test + public void everyPrimeCarriesFreshRequestIdWithFirstAttempt() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + ManagedChannel channel = newChannel(); + + // The pool invokes prime() once per attempt, so each attempt is a new request id. + getWithin(primer.prime(channel), Duration.ofSeconds(10)); + getWithin(primer.prime(channel), Duration.ofSeconds(10)); + + assertThat(service.headers).hasSize(2); + XGoogSpannerRequestId first = requestIdOf(service.headers.get(0)); + XGoogSpannerRequestId second = requestIdOf(service.headers.get(1)); + assertThat(second).isNotEqualTo(first); + assertThat(requestNumberOf(second)).isGreaterThan(requestNumberOf(first)); + assertThat(second.getLogicalRequestKey()).isNotEqualTo(first.getLogicalRequestKey()); + } + + @Test + public void primeOmitsRouteToLeaderHeaderWhenDisabled() throws Exception { + DynamicChannelPoolPrimer primer = + newPrimer( + /* callCredentialsProvider= */ null, + /* routeToLeader= */ false, + /* multiplexedSessionsEnabled= */ true, + Duration.ofSeconds(5)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + assertThat(service.headers.get(0).containsKey(ROUTE_TO_LEADER_KEY)).isFalse(); + } + + @Test + public void primePrefersCallCredentialsProviderOverDefaultCredentials() throws Exception { + DynamicChannelPoolPrimer primer = + newPrimer( + () -> credentialsWithToken(PROVIDER_TOKEN), + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ true, + Duration.ofSeconds(5)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + assertThat(service.headers.get(0).getAll(AUTHORIZATION_KEY)) + .containsExactly("Bearer " + PROVIDER_TOKEN); + } + + @Test + public void primeFallsBackToDefaultCredentialsWhenProviderReturnsNull() throws Exception { + DynamicChannelPoolPrimer primer = + newPrimer( + () -> null, + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ true, + Duration.ofSeconds(5)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + assertThat(service.headers.get(0).getAll(AUTHORIZATION_KEY)) + .containsExactly("Bearer " + DEFAULT_TOKEN); + } + + @Test + public void primeFailsWhenRpcFails() throws Exception { + service.failWith = Status.UNAVAILABLE.withDescription("backend unavailable"); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + Throwable failure = failureOf(primer.prime(newChannel())); + + assertThat(failure).isInstanceOf(StatusRuntimeException.class); + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.UNAVAILABLE); + // A transient failure says nothing about the session, so it stays registered. + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + } + + @Test + public void notFoundFailureEvictsSessionAndNextPrimeFallsBackToOlderSession() throws Exception { + service.failSessionsWith.put( + OTHER_SESSION_NAME, Status.NOT_FOUND.withDescription("Session not found")); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + ManagedChannel channel = newChannel(); + + // The pool invokes prime() once per attempt: the first attempt uses the newest session, which + // is gone, and evicts it; the next attempt falls back to the older session. + Throwable failure = failureOf(primer.prime(channel)); + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.NOT_FOUND); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + assertThat(primer.getPrimeSessions()).hasSize(1); + + assertThat(getWithin(primer.prime(channel), Duration.ofSeconds(10))).isNull(); + + assertThat(service.requests).hasSize(2); + assertThat(service.requests.get(0).getSession()).isEqualTo(OTHER_SESSION_NAME); + assertThat(service.requests.get(1).getSession()).isEqualTo(SESSION_NAME); + } + + @Test + public void failedPreconditionAboutSessionEvictsSession() throws Exception { + service.failWith = Status.FAILED_PRECONDITION.withDescription("Invalid session: expired"); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + Throwable failure = failureOf(primer.prime(newChannel())); + + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.FAILED_PRECONDITION); + assertThat(primer.getPrimeSession()).isNull(); + } + + @Test + public void unrelatedFailedPreconditionKeepsSession() throws Exception { + service.failWith = Status.FAILED_PRECONDITION.withDescription("Database is not ready"); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + failureOf(primer.prime(newChannel())); + + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + } + + @Test + public void permissionDeniedFailureEvictsSessionAndNextPrimeFallsBackToOlderSession() + throws Exception { + service.failSessionsWith.put( + OTHER_SESSION_NAME, + Status.PERMISSION_DENIED.withDescription("Caller is missing IAM permission")); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); + ManagedChannel channel = newChannel(); + + // The caller cannot use the newest session, so the attempt rotates to the older candidate. + Throwable failure = failureOf(primer.prime(channel)); + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.PERMISSION_DENIED); + assertThat(primer.getPrimeSessions()).hasSize(1); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + + assertThat(getWithin(primer.prime(channel), Duration.ofSeconds(10))).isNull(); + assertThat(service.requests).hasSize(2); + assertThat(service.requests.get(1).getSession()).isEqualTo(SESSION_NAME); + } + + @Test + public void invalidDatabaseRoleFailureEvictsSession() throws Exception { + service.failWith = + Status.INVALID_ARGUMENT.withDescription("Database role my-role is not valid"); + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + Throwable failure = failureOf(primer.prime(newChannel())); + + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.INVALID_ARGUMENT); + assertThat(primer.getPrimeSession()).isNull(); + } + + @Test + public void unavailableAndDeadlineExceededKeepSession() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + service.failWith = Status.UNAVAILABLE.withDescription("Connection reset"); + failureOf(primer.prime(newChannel())); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + + service.failWith = Status.DEADLINE_EXCEEDED.withDescription("Deadline exceeded"); + failureOf(primer.prime(newChannel())); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + } + + @Test + public void candidateSpecificFailureClassification() { + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.PERMISSION_DENIED.withDescription("Permission denied").asRuntimeException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.PERMISSION_DENIED.asException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.INVALID_ARGUMENT + .withDescription("Database role my-role is not valid") + .asRuntimeException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.FAILED_PRECONDITION + .withDescription("Role my-role does not exist") + .asRuntimeException())) + .isTrue(); + // A message that merely contains the letters of the word role is not about a database role. + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.INVALID_ARGUMENT + .withDescription("Controller rejected the statement") + .asRuntimeException())) + .isFalse(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.INVALID_ARGUMENT.withDescription("Syntax error").asRuntimeException())) + .isFalse(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.UNAUTHENTICATED.withDescription("Invalid role").asRuntimeException())) + .isFalse(); + } + + @Test + public void lateRegistrationAfterUnregisterIsDropped() { + DynamicChannelPoolPrimer primer = newPrimer(); + // The initial session of a database client, and a refresh whose CreateSession is in flight + // with the same owner ticket. + long retiredTicket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + // The database client is retired while the refresh is in flight. + primer.unregisterPrimeOwner(DATABASE_NAME, retiredTicket); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.getPrimeOwners()).isEmpty(); + + // The refresh returns after the retirement and is dropped: the database has no owner. + assertThat(primer.registerPrimeSession(DATABASE_NAME, REFRESHED_SESSION_NAME, retiredTicket)) + .isFalse(); + assertThat(primer.getPrimeSession()).isNull(); + + // The replacement client registers itself as the new owner and registers normally. + long replacementTicket = OWNER_TICKETS.incrementAndGet(); + primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); + String replacementSession = DATABASE_NAME + "/sessions/replacement-session"; + assertThat(primer.registerPrimeSession(DATABASE_NAME, replacementSession, replacementTicket)) + .isTrue(); + assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); + + // An even later response of the retired client still does not replace the new session: the + // database is owned by the replacement. + assertThat(primer.registerPrimeSession(DATABASE_NAME, REFRESHED_SESSION_NAME, retiredTicket)) + .isFalse(); + assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); + assertThat(primer.getPrimeSessions()).hasSize(1); + assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); + } + + @Test + public void registrationWithoutOwnerIsDropped() { + DynamicChannelPoolPrimer primer = newPrimer(); + + // No owner at all, for example a session that was created without an owner ticket. + assertThat(primer.registerPrimeSession(DATABASE_NAME, SESSION_NAME, 42L)).isFalse(); + assertThat(primer.getPrimeSession()).isNull(); + + // An owner of another database does not own this database. + primer.registerPrimeOwner(OTHER_DATABASE_NAME, 42L); + assertThat(primer.registerPrimeSession(DATABASE_NAME, SESSION_NAME, 42L)).isFalse(); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.registerPrimeSession(OTHER_DATABASE_NAME, OTHER_SESSION_NAME, 42L)).isTrue(); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + } + + @Test + public void replacementOwnerTakesOverTheDatabase() { + DynamicChannelPoolPrimer primer = newPrimer(); + long retiredTicket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + // The replacement registers itself before the retired client has been closed, for example + // because the close of the retired client is still running. The retired client's session is + // dropped right away, because the retired client no longer maintains it. + long replacementTicket = OWNER_TICKETS.incrementAndGet(); + primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); + String replacementSession = DATABASE_NAME + "/sessions/replacement-session"; + assertThat(primer.registerPrimeSession(DATABASE_NAME, replacementSession, replacementTicket)) + .isTrue(); + + // The late close of the retired client leaves the replacement and its session untouched. + primer.unregisterPrimeOwner(DATABASE_NAME, retiredTicket); + assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); + assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); + + // Registering the same owner again is a no-op that keeps its session. + primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); + assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); + } + + @Test + public void retirementOfOneDatabaseDoesNotAffectOtherDatabases() { + DynamicChannelPoolPrimer primer = newPrimer(); + long ticket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + long otherTicket = OWNER_TICKETS.incrementAndGet(); + primer.registerPrimeOwner(OTHER_DATABASE_NAME, otherTicket); + + primer.unregisterPrimeOwner(DATABASE_NAME, ticket); + + assertThat(primer.registerPrimeSession(OTHER_DATABASE_NAME, OTHER_SESSION_NAME, otherTicket)) + .isTrue(); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + assertThat(primer.getPrimeOwners()).containsExactly(OTHER_DATABASE_NAME, otherTicket); + } + + @Test + public void churnOfManyDatabaseClientsLeavesNothingBehind() { + DynamicChannelPoolPrimer primer = newPrimer(); + // A long-lived client that churns through many database names, including repeated clients of + // the same database, must not retain anything for retired clients. + for (int i = 0; i < 10_000; i++) { + String databaseName = "projects/p/instances/i/databases/d" + (i % 100); + long ticket = OWNER_TICKETS.incrementAndGet(); + primer.registerPrimeOwner(databaseName, ticket); + assertThat( + primer.registerPrimeSession(databaseName, databaseName + "/sessions/s" + i, ticket)) + .isTrue(); + assertThat(primer.getPrimeOwners()).hasSize(1); + assertThat(primer.getPrimeSessions()).hasSize(1); + primer.unregisterPrimeOwner(databaseName, ticket); + // A late CreateSession response of the retired client finds no owner and is dropped. + assertThat( + primer.registerPrimeSession( + databaseName, databaseName + "/sessions/late" + i, ticket)) + .isFalse(); + assertThat(primer.getPrimeOwners()).isEmpty(); + assertThat(primer.getPrimeSessions()).isEmpty(); + } + } + + @Test + public void invalidSessionFailureClassification() { + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.NOT_FOUND.withDescription("Session not found").asRuntimeException())) + .isTrue(); + assertThat(DynamicChannelPoolPrimer.isCandidateSpecificFailure(Status.NOT_FOUND.asException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.FAILED_PRECONDITION + .withDescription("Session does not exist") + .asRuntimeException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.FAILED_PRECONDITION.withDescription("invalid session").asRuntimeException())) + .isTrue(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.FAILED_PRECONDITION.asRuntimeException())) + .isFalse(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.FAILED_PRECONDITION + .withDescription("Transaction was already committed") + .asRuntimeException())) + .isFalse(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.UNAVAILABLE.withDescription("Session not found").asRuntimeException())) + .isFalse(); + assertThat( + DynamicChannelPoolPrimer.isCandidateSpecificFailure( + Status.DEADLINE_EXCEEDED.asRuntimeException())) + .isFalse(); + assertThat(DynamicChannelPoolPrimer.isCandidateSpecificFailure(new IllegalStateException())) + .isFalse(); + } + + @Test + public void primeFailsWithDeadlineExceededWhenServerDoesNotRespond() throws Exception { + service.holdResponses = true; + DynamicChannelPoolPrimer primer = + newPrimer( + /* callCredentialsProvider= */ null, + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ true, + Duration.ofMillis(200)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + Throwable failure = failureOf(primer.prime(newChannel())); + + assertThat(failure).isInstanceOf(StatusRuntimeException.class); + assertThat(((StatusRuntimeException) failure).getStatus().getCode()) + .isEqualTo(Status.Code.DEADLINE_EXCEEDED); + // The deadline cancels the server-side call as well. + assertThat(service.callCancelled.await(5, TimeUnit.SECONDS)).isTrue(); + // A deadline says nothing about the session, so it stays registered. + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + } + + @Test + public void cancellingPrimeCancelsRpc() throws Exception { + service.holdResponses = true; + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + ListenableFuture future = primer.prime(newChannel()); + assertThat(service.callStarted.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(future.isDone()).isFalse(); + + assertThat(future.cancel(true)).isTrue(); + + assertThat(future.isCancelled()).isTrue(); + assertThat(service.callCancelled.await(5, TimeUnit.SECONDS)).isTrue(); + } + + @Test + public void primeFailsWhenChannelIsShutDownDuringRpc() throws Exception { + service.holdResponses = true; + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + ManagedChannel channel = newChannel(); + + ListenableFuture future = primer.prime(channel); + assertThat(service.callStarted.await(5, TimeUnit.SECONDS)).isTrue(); + channel.shutdownNow(); + + Throwable failure = failureOf(future); + assertThat(failure).isInstanceOf(StatusRuntimeException.class); + assertThat(future.isCancelled()).isFalse(); + } + + @Test + public void primeOnShutDownChannelFailsPromptly() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + ManagedChannel channel = newChannel(); + channel.shutdownNow(); + + Throwable failure = failureOf(primer.prime(channel)); + + assertThat(failure).isInstanceOf(StatusRuntimeException.class); + } + + @Test + public void primeFailsWithFailedPreconditionWithoutSession() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + + ListenableFuture future = primer.prime(newChannel()); + + assertThat(future.isDone()).isTrue(); + Throwable failure = failureOf(future); + assertThat(failure).isInstanceOf(SpannerException.class); + assertThat(((SpannerException) failure).getErrorCode()) + .isEqualTo(ErrorCode.FAILED_PRECONDITION); + assertThat(service.requests).isEmpty(); + } + + @Test + public void primeSucceedsWithoutRpcWhenMultiplexedSessionsAreDisabled() throws Exception { + DynamicChannelPoolPrimer primer = + newPrimer( + /* callCredentialsProvider= */ null, + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ false, + Duration.ofSeconds(5)); + + ListenableFuture future = primer.prime(newChannel()); + + assertThat(future.isDone()).isTrue(); + assertThat(getWithin(future, Duration.ofSeconds(1))).isNull(); + assertThat(service.requests).isEmpty(); + } + + @Test + public void concurrentPrimesOnOneChannelAllSucceed() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + ManagedChannel channel = newChannel(); + + List> futures = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + futures.add(primer.prime(channel)); + } + for (ListenableFuture future : futures) { + assertThat(getWithin(future, Duration.ofSeconds(10))).isNull(); + } + assertThat(service.requests).hasSize(5); + } + + @Test + public void latestRegisteredSessionIsUsedForPriming() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.getPrimeSessionName()).isNull(); + + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + // A session that is created later, for example the refresh of the multiplexed session, + // replaces the previous prime session of that database. + registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); + assertThat(primer.getPrimeSessions()).hasSize(1); + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + assertThat(service.requests).hasSize(2); + assertThat(service.requests.get(0).getSession()).isEqualTo(SESSION_NAME); + assertThat(service.requests.get(1).getSession()).isEqualTo(REFRESHED_SESSION_NAME); + } + + @Test + public void newestGenerationIsSelectedAcrossDatabases() { + DynamicChannelPoolPrimer primer = newPrimer(); + + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); + + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + assertThat(sessionNames(primer)).containsExactly(OTHER_SESSION_NAME, SESSION_NAME).inOrder(); + assertThat(primer.getPrimeSessions().get(0).getGeneration()) + .isGreaterThan(primer.getPrimeSessions().get(1).getGeneration()); + + // The refresh of the first database replaces its entry and makes it the newest again. + registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); + + assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); + assertThat(sessionNames(primer)) + .containsExactly(REFRESHED_SESSION_NAME, OTHER_SESSION_NAME) + .inOrder(); + assertThat(primer.getPrimeSessions().get(0).getDatabaseName()).isEqualTo(DATABASE_NAME); + assertThat(primer.getPrimeSessions().get(0).getGeneration()) + .isGreaterThan(primer.getPrimeSessions().get(1).getGeneration()); + } + + @Test + public void evictionNeverClobbersNewerEntryOfSameDatabase() { + DynamicChannelPoolPrimer primer = newPrimer(); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + PrimeSession stale = primer.getPrimeSession(); + assertThat(stale).isNotNull(); + + // A newer session of the same database was registered before the stale one is evicted. + registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); + PrimeSession current = primer.getPrimeSession(); + assertThat(current).isNotEqualTo(stale); + + assertThat(primer.evictPrimeSession(stale)).isFalse(); + assertThat(primer.getPrimeSession()).isEqualTo(current); + assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); + + assertThat(primer.evictPrimeSession(current)).isTrue(); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.evictPrimeSession(current)).isFalse(); + } + + @Test + public void unregisterRemovesOnlyThatDatabaseClient() { + DynamicChannelPoolPrimer primer = newPrimer(); + long ticket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + long otherTicket = registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); + + primer.unregisterPrimeOwner(OTHER_DATABASE_NAME, otherTicket); + assertThat(sessionNames(primer)).containsExactly(SESSION_NAME); + assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, ticket); + + // Unknown databases and unknown tickets are ignored. + primer.unregisterPrimeOwner( + "projects/my-project/instances/my-instance/databases/unknown", ticket); + primer.unregisterPrimeOwner(DATABASE_NAME, ticket + 1000); + assertThat(sessionNames(primer)).containsExactly(SESSION_NAME); + assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, ticket); + + primer.unregisterPrimeOwner(DATABASE_NAME, ticket); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.getPrimeSessions()).isEmpty(); + assertThat(primer.getPrimeOwners()).isEmpty(); + } + + private static List sessionNames(DynamicChannelPoolPrimer primer) { + List names = new ArrayList<>(); + for (PrimeSession entry : primer.getPrimeSessions()) { + names.add(entry.getSessionName()); + } + return names; + } + + @Test + public void rpcDeadlineNormallyStaysBelowPoolPrimeTimeout() { + // The default prime timeout of 10 seconds is capped at the maximum RPC deadline. + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofSeconds(10))) + .isEqualTo(Duration.ofSeconds(5)); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofSeconds(6))) + .isEqualTo(Duration.ofSeconds(5)); + // Below that, the deadline keeps the safety margin below the prime timeout. + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofSeconds(4))) + .isEqualTo(Duration.ofSeconds(3)); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofSeconds(2))) + .isEqualTo(Duration.ofSeconds(1)); + // Very short prime timeouts never drop the deadline below half of the prime timeout. + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofMillis(1500))) + .isEqualTo(Duration.ofMillis(750)); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofMillis(200))) + .isEqualTo(Duration.ofMillis(100)); + // Prime timeouts of at most twice the minimum deadline are clamped to the minimum deadline, + // which then reaches (2 ms) or exceeds (1 ms, 1 ns) the prime timeout itself; the pool's own + // prime timeout bounds the attempt in that case. + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofMillis(3))) + .isEqualTo(Duration.ofNanos(1_500_000L)); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofMillis(2))) + .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofMillis(1))) + .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); + assertThat(DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofNanos(1))) + .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); + // The clamped deadline is accepted by the primer. + assertThat( + newPrimer( + /* callCredentialsProvider= */ null, + /* routeToLeader= */ true, + /* multiplexedSessionsEnabled= */ true, + DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofNanos(1))) + .getRpcDeadline()) + .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); + assertThrows( + IllegalArgumentException.class, + () -> DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ZERO)); + assertThrows( + IllegalArgumentException.class, + () -> DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofSeconds(-1))); + assertThrows( + IllegalArgumentException.class, () -> DynamicChannelPoolPrimer.rpcDeadlineFor(null)); + } + + @Test + public void emptyNamesAreRejected() { + DynamicChannelPoolPrimer primer = newPrimer(); + + assertThrows( + IllegalArgumentException.class, () -> primer.registerPrimeSession(DATABASE_NAME, "", 1L)); + assertThrows( + IllegalArgumentException.class, () -> primer.registerPrimeSession(DATABASE_NAME, null, 1L)); + assertThrows( + IllegalArgumentException.class, () -> primer.registerPrimeSession("", SESSION_NAME, 1L)); + assertThrows( + IllegalArgumentException.class, () -> primer.registerPrimeSession(null, SESSION_NAME, 1L)); + assertThrows(IllegalArgumentException.class, () -> primer.registerPrimeOwner("", 1L)); + assertThrows(IllegalArgumentException.class, () -> primer.registerPrimeOwner(null, 1L)); + assertThrows(IllegalArgumentException.class, () -> primer.unregisterPrimeOwner("", 1L)); + assertThrows(IllegalArgumentException.class, () -> primer.unregisterPrimeOwner(null, 1L)); + assertThat(primer.getPrimeSession()).isNull(); + assertThat(primer.getPrimeOwners()).isEmpty(); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java index de21347e0609..8c31cfc350ea 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java @@ -20,8 +20,10 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeTrue; @@ -41,6 +43,7 @@ import com.google.auth.oauth2.OAuth2Credentials; import com.google.cloud.NoCredentials; import com.google.cloud.ServiceOptions; +import com.google.cloud.grpc.GcpChannelPrimer; import com.google.cloud.grpc.GcpManagedChannel; import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; import com.google.cloud.grpc.GcpManagedChannelOptions; @@ -64,11 +67,15 @@ import com.google.cloud.spanner.SpannerOptionsHelper; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.TransactionRunner; +import com.google.cloud.spanner.XGoogSpannerRequestId; import com.google.cloud.spanner.spi.v1.GapicSpannerRpc.AdminRequestsLimitExceededRetryAlgorithm; import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Futures; import com.google.protobuf.ListValue; import com.google.rpc.ErrorInfo; +import com.google.spanner.v1.CreateSessionRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.GetSessionRequest; import com.google.spanner.v1.ResultSetMetadata; @@ -92,6 +99,7 @@ import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder; import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; import io.grpc.protobuf.lite.ProtoLiteUtils; +import io.grpc.stub.StreamObserver; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; import io.opentelemetry.context.propagation.ContextPropagators; @@ -105,16 +113,19 @@ import java.lang.reflect.Array; import java.lang.reflect.Modifier; import java.net.InetSocketAddress; +import java.net.URLEncoder; import java.time.Duration; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.IdentityHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; @@ -123,6 +134,8 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BooleanSupplier; +import javax.annotation.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -204,6 +217,9 @@ String debugString() { private static InetSocketAddress address; private static final Map optionsMap = new HashMap<>(); private static Metadata lastSeenHeaders; + // Headers of every unary ExecuteSql call, which is the RPC that channel priming uses. + private static final List>> executeSqlHeaders = + new CopyOnWriteArrayList<>(); private static String defaultUserAgent; private static Spanner spanner; private static boolean isRouteToLeader; @@ -248,6 +264,9 @@ public ServerCall.Listener interceptCall( Metadata headers, ServerCallHandler next) { lastSeenHeaders = headers; + if (call.getMethodDescriptor().equals(SpannerGrpc.getExecuteSqlMethod())) { + executeSqlHeaders.add(copyAsciiHeaders(headers)); + } String auth = headers.get(Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER)); assertThat(auth).isEqualTo("Bearer " + VARIABLE_OAUTH_TOKEN); @@ -302,6 +321,354 @@ public void reset() throws InterruptedException { isRouteToLeader = false; isEndToEndTracing = false; isTraceContextPresent = false; + executeSqlHeaders.clear(); + } + + /** Copies all values of every ASCII header, so duplicated headers remain visible. */ + private static Map> copyAsciiHeaders(Metadata headers) { + Map> copy = new HashMap<>(); + for (String key : headers.keys()) { + if (!key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { + Iterable values = headers.getAll(Key.of(key, Metadata.ASCII_STRING_MARSHALLER)); + copy.put(key, values == null ? ImmutableList.of() : ImmutableList.copyOf(values)); + } + } + return copy; + } + + private static void awaitCondition(BooleanSupplier condition, Duration timeout) + throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (!condition.getAsBoolean()) { + if (System.nanoTime() > deadline) { + throw new AssertionError("Condition not met within " + timeout); + } + Thread.sleep(10L); + } + } + + private static GapicSpannerRpc getRpc(Spanner spanner) throws Exception { + java.lang.reflect.Method method = SpannerOptions.class.getDeclaredMethod("getSpannerRpcV1"); + method.setAccessible(true); + return (GapicSpannerRpc) method.invoke(spanner.getOptions()); + } + + /** The CreateSession options of a database client with the given prime owner ticket. */ + private static Map primeOwner(long ownerTicket) { + return ImmutableMap.of(Option.CHANNEL_PRIME_OWNER, ownerTicket); + } + + private static void registerPrimeStatementResult() { + mockSpanner.putStatementResult( + StatementResult.query( + Statement.of(DynamicChannelPoolPrimer.PRIME_SQL), + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues( + com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata(SELECT1AND2_METADATA) + .build())); + } + + private static List primeRequests() { + List requests = new ArrayList<>(); + for (ExecuteSqlRequest request : mockSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { + if (request.getSql().equals(DynamicChannelPoolPrimer.PRIME_SQL)) { + requests.add(request); + } + } + return requests; + } + + @Test + public void testChannelPoolOptionsRegisterPrimerOnlyWithDynamicChannelPool() { + GcpChannelPrimer primer = channel -> Futures.immediateFuture(null); + + SpannerOptions dcpOptions = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + GcpChannelPoolOptions withPrimer = + GapicSpannerRpc.getGrpcGcpChannelPoolOptions(dcpOptions, primer); + assertSame(primer, withPrimer.getChannelPrimer()); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT, + withPrimer.getChannelPrimeTimeout()); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS, + withPrimer.getChannelPrimeMaxAttempts()); + // The other dynamic pool settings are retained. + assertEquals( + dcpOptions.getGcpChannelPoolOptions().getMaxRpcPerChannel(), + withPrimer.getMaxRpcPerChannel()); + assertNull(GapicSpannerRpc.getGrpcGcpChannelPoolOptions(dcpOptions, null).getChannelPrimer()); + assertNull(GapicSpannerRpc.getGrpcGcpChannelPoolOptions(dcpOptions).getChannelPrimer()); + + SpannerOptions staticOptions = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .enableGrpcGcpExtension() + .disableDynamicChannelPool() + .setNumChannels(2) + .build(); + assertNull( + GapicSpannerRpc.getGrpcGcpChannelPoolOptions(staticOptions, primer).getChannelPrimer()); + } + + @Test + public void testChannelPoolOptionsNeverOverrideUserProvidedPrimer() { + GcpChannelPrimer userPrimer = channel -> Futures.immediateFuture(null); + GcpChannelPrimer spannerPrimer = channel -> Futures.immediateFuture(null); + Duration userTimeout = Duration.ofSeconds(3); + int userAttempts = 7; + SpannerOptions options = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setChannelPrimer(userPrimer) + .setChannelPrimeTimeout(userTimeout) + .setChannelPrimeMaxAttempts(userAttempts) + .build()) + .build(); + + GcpChannelPoolOptions poolOptions = + GapicSpannerRpc.getGrpcGcpChannelPoolOptions(options, spannerPrimer); + + assertSame(userPrimer, poolOptions.getChannelPrimer()); + assertEquals(userTimeout, poolOptions.getChannelPrimeTimeout()); + assertEquals(userAttempts, poolOptions.getChannelPrimeMaxAttempts()); + } + + @Test + public void testUserProvidedPrimeSettingsSurviveWithoutUserPrimer() { + GcpChannelPrimer spannerPrimer = channel -> Futures.immediateFuture(null); + Duration userTimeout = Duration.ofSeconds(4); + int userAttempts = 2; + SpannerOptions options = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setChannelPrimeTimeout(userTimeout) + .setChannelPrimeMaxAttempts(userAttempts) + .build()) + .build(); + + GcpChannelPoolOptions poolOptions = + GapicSpannerRpc.getGrpcGcpChannelPoolOptions(options, spannerPrimer); + + assertSame(spannerPrimer, poolOptions.getChannelPrimer()); + assertEquals(userTimeout, poolOptions.getChannelPrimeTimeout()); + assertEquals(userAttempts, poolOptions.getChannelPrimeMaxAttempts()); + } + + @Test + public void testRpcCreatesPrimerOnlyWithDynamicChannelPool() { + GapicSpannerRpc dcpRpc = + new GapicSpannerRpc( + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(), + true); + try { + assertNotNull(dcpRpc.getChannelPrimer()); + assertNotNull(findGrpcGcpChannel(dcpRpc)); + } finally { + dcpRpc.shutdown(); + } + + GapicSpannerRpc staticPoolRpc = + new GapicSpannerRpc( + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .disableDynamicChannelPool() + .setNumChannels(1) + .build(), + true); + try { + assertNull(staticPoolRpc.getChannelPrimer()); + } finally { + staticPoolRpc.shutdown(); + } + + GapicSpannerRpc gaxPoolRpc = + new GapicSpannerRpc( + createSpannerOptions().toBuilder() + .disableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(), + true); + try { + assertNull(gaxPoolRpc.getChannelPrimer()); + } finally { + gaxPoolRpc.shutdown(); + } + } + + @Test + public void testPrimerRpcDeadlineFollowsPoolPrimeTimeout() { + SpannerOptions defaultTimeoutOptions = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + assertEquals( + SpannerOptions.DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT, + defaultTimeoutOptions.getGcpChannelPoolOptions().getChannelPrimeTimeout()); + GapicSpannerRpc defaultTimeoutRpc = new GapicSpannerRpc(defaultTimeoutOptions, true); + try { + assertEquals( + DynamicChannelPoolPrimer.MAX_RPC_DEADLINE, + defaultTimeoutRpc.getChannelPrimer().getRpcDeadline()); + } finally { + defaultTimeoutRpc.shutdown(); + } + + // A user-provided prime timeout below the maximum RPC deadline pulls the RPC deadline down + // with it, so the RPC always fails before the pool times out the attempt. + Duration userTimeout = Duration.ofSeconds(3); + SpannerOptions shortTimeoutOptions = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder().setChannelPrimeTimeout(userTimeout).build()) + .build(); + assertEquals( + userTimeout, shortTimeoutOptions.getGcpChannelPoolOptions().getChannelPrimeTimeout()); + GapicSpannerRpc shortTimeoutRpc = new GapicSpannerRpc(shortTimeoutOptions, true); + try { + GcpManagedChannel pool = findGrpcGcpChannel(shortTimeoutRpc); + assertNotNull(pool); + Duration rpcDeadline = shortTimeoutRpc.getChannelPrimer().getRpcDeadline(); + assertEquals(userTimeout.minus(DynamicChannelPoolPrimer.RPC_DEADLINE_MARGIN), rpcDeadline); + assertTrue(rpcDeadline.compareTo(userTimeout) < 0); + } finally { + shortTimeoutRpc.shutdown(); + } + } + + @Test + public void testDynamicChannelPoolPrimesScaledUpChannelsWithSelectOne() throws Exception { + registerPrimeStatementResult(); + // Keep the load queries open long enough to trigger a scale-up, and make the priming query + // slow enough to observe that the channel is not published before it succeeds. + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(3000, 0)); + mockSpanner.setExecuteSqlExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(1000, 0)); + SpannerOptions options = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .enableLeaderAwareRouting() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setInitSize(1) + .setMinSize(1) + .setMaxSize(4) + .setDynamicScaling(1, 2, Duration.ofMinutes(3)) + .build()) + .build(); + ExecutorService executor = Executors.newFixedThreadPool(8); + try (Spanner spanner = options.getService()) { + GapicSpannerRpc rpc = getRpc(spanner); + assertNotNull(rpc.getChannelPrimer()); + GcpManagedChannel pool = findGrpcGcpChannel(rpc); + assertNotNull(pool); + assertEquals(1, pool.getNumberOfChannels()); + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + + // Six concurrent streams on one channel with maxRpcPerChannel=2 force a scale-up. + List> loads = new ArrayList<>(); + for (int i = 0; i < 6; i++) { + loads.add( + executor.submit( + () -> { + long rows = 0; + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1AND2)) { + while (resultSet.next()) { + rows++; + } + } + return rows; + })); + } + + awaitCondition(() -> !primeRequests().isEmpty(), Duration.ofSeconds(15)); + // The priming query is still running (it takes at least one second), so the scaled-up + // channel has not been published yet. + assertEquals(1, pool.getNumberOfChannels()); + awaitCondition(() -> pool.getNumberOfChannels() > 1, Duration.ofSeconds(15)); + for (Future load : loads) { + assertEquals(2L, load.get().longValue()); + } + + // Every priming query used the multiplexed session, which is the session that the + // single-use load queries used as well. + String multiplexedSession = null; + for (ExecuteSqlRequest request : mockSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { + if (request.getSql().equals(SELECT1AND2.getSql())) { + multiplexedSession = request.getSession(); + break; + } + } + assertNotNull(multiplexedSession); + // The mock keeps multiplexed sessions apart from regular sessions. + assertFalse(mockSpanner.getSessions().containsKey(multiplexedSession)); + boolean multiplexedSessionCreated = false; + for (CreateSessionRequest request : + mockSpanner.getRequestsOfType(CreateSessionRequest.class)) { + multiplexedSessionCreated |= request.getSession().getMultiplexed(); + } + assertTrue(multiplexedSessionCreated); + List primes = primeRequests(); + assertThat(primes).isNotEmpty(); + for (ExecuteSqlRequest prime : primes) { + assertEquals(multiplexedSession, prime.getSession()); + assertFalse(prime.hasTransaction()); + } + assertEquals(primes.size(), pool.getNumberOfChannels() - 1); + + // The priming query carries the same credentials and headers as a normal call, and every + // header exactly once: the fixed headers come from the delegate channel and the per-call + // headers from the primer. + assertEquals(primes.size(), executeSqlHeaders.size()); + Set requestIds = new HashSet<>(); + for (Map> headers : executeSqlHeaders) { + assertThat(headers.get("authorization")).containsExactly("Bearer " + VARIABLE_OAUTH_TOKEN); + assertThat(headers.get("x-goog-api-client")).hasSize(1); + assertThat(headers.get("x-goog-api-client").get(0)) + .contains(ServiceOptions.getGoogApiClientLibName() + "/"); + assertThat(headers.get(ApiClientHeaderProvider.getDefaultResourceHeaderKey())) + .containsExactly("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE]"); + assertThat(headers.get("x-goog-request-params")) + .containsExactly("session=" + URLEncoder.encode(multiplexedSession, "UTF-8")); + assertThat(headers.get("x-goog-spanner-route-to-leader")).containsExactly("true"); + // The request id uses the client id of the rpc, the unknown channel 0 because the channel + // is not part of the pool yet, and attempt 1 because the pool invokes prime() per attempt. + assertThat(headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME)).hasSize(1); + String requestId = headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME).get(0); + assertTrue(requestId, requestIds.add(requestId)); + String[] parts = requestId.split("\\."); + assertEquals(requestId, 6, parts.length); + assertEquals(requestId, String.valueOf(rpc.getRequestIdCreator().getClientId()), parts[2]); + assertEquals(requestId, "0", parts[3]); + assertEquals(requestId, "1", parts[5]); + } + } finally { + executor.shutdownNow(); + } } @Test @@ -1019,6 +1386,433 @@ public void testCreateSession_whenMultiplexedSessionIsFalse_assertSessionProto() rpc.shutdown(); } + @Test + public void testMultiplexedCreateSessionRecordsLatestPrimeSession() { + SpannerOptions options = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); + try { + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + assertNull(primer.getPrimeSessionName()); + rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); + + // A regular session never becomes the prime session. + Session regular = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), false); + assertFalse(regular.getMultiplexed()); + assertNull(primer.getPrimeSessionName()); + + // A multiplexed session that is created without an owner ticket, for example by a batch + // client, never becomes the prime session either. + Session unowned = rpc.createSession("DATABASE_NAME", null, null, null, true); + assertTrue(unowned.getMultiplexed()); + assertNull(primer.getPrimeSessionName()); + + // The first successfully created multiplexed session of the owner becomes the prime session. + Session first = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); + assertTrue(first.getMultiplexed()); + assertEquals(first.getName(), primer.getPrimeSessionName()); + + // A failed multiplexed CreateSession leaves the prime session untouched. + mockSpanner.setCreateSessionExecutionTime( + SimulatedExecutionTime.ofException( + Status.PERMISSION_DENIED.withDescription("test").asRuntimeException())); + SpannerException exception = + assertThrows( + SpannerException.class, + () -> rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true)); + assertEquals(ErrorCode.PERMISSION_DENIED, exception.getErrorCode()); + assertEquals(first.getName(), primer.getPrimeSessionName()); + + // The periodic refresh of the multiplexed session carries the same ticket and replaces the + // prime session of its database. + Session refreshed = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); + assertTrue(refreshed.getMultiplexed()); + assertNotEquals(first.getName(), refreshed.getName()); + assertEquals(refreshed.getName(), primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(1); + + // The multiplexed session of another database client replaces the prime session as well. + rpc.registerChannelPrimeOwner("OTHER_DATABASE_NAME", 2L); + Session second = rpc.createSession("OTHER_DATABASE_NAME", null, null, primeOwner(2L), true); + assertTrue(second.getMultiplexed()); + assertNotEquals(refreshed.getName(), second.getName()); + assertEquals(second.getName(), primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(2); + } finally { + rpc.shutdown(); + } + } + + @Test + public void testChurnOfDatabaseClientsLeavesNoPrimeOwnersOrSessionsBehind() { + SpannerOptions options = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); + try { + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + for (int i = 0; i < 50; i++) { + String databaseName = "projects/p/instances/i/databases/d" + i; + long ticket = 100L + i; + rpc.registerChannelPrimeOwner(databaseName, ticket); + Session session = rpc.createSession(databaseName, null, null, primeOwner(ticket), true); + assertEquals(session.getName(), primer.getPrimeSessionName()); + assertThat(primer.getPrimeOwners()).containsExactly(databaseName, ticket); + rpc.unregisterChannelPrimeOwner(databaseName, ticket); + assertThat(primer.getPrimeOwners()).isEmpty(); + assertThat(primer.getPrimeSessions()).isEmpty(); + } + } finally { + rpc.shutdown(); + } + } + + @Test + public void testDynamicChannelPoolPrimingFallsBackToOlderSessionWhenNewestIsInvalid() + throws Exception { + registerPrimeStatementResult(); + // Keep the load queries open long enough to trigger a scale-up. + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(3000, 0)); + SpannerOptions options = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setInitSize(1) + .setMinSize(1) + .setMaxSize(4) + .setDynamicScaling(1, 2, Duration.ofMinutes(3)) + .build()) + .build(); + ExecutorService executor = Executors.newFixedThreadPool(8); + try (Spanner spanner = options.getService()) { + GapicSpannerRpc rpc = getRpc(spanner); + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + GcpManagedChannel pool = findGrpcGcpChannel(rpc); + assertNotNull(pool); + + // The older database client registers its multiplexed session first. + DatabaseClient older = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + awaitCondition(() -> primer.getPrimeSessionName() != null, Duration.ofSeconds(15)); + String olderSession = primer.getPrimeSessionName(); + assertThat(olderSession) + .startsWith("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE]/"); + // The newer database client registers the session that priming prefers. + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE2]")); + awaitCondition(() -> primer.getPrimeSessions().size() == 2, Duration.ofSeconds(15)); + String newestSession = primer.getPrimeSessionName(); + assertThat(newestSession) + .startsWith("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE2]/"); + // The newest session becomes invalid on the backend: the mock returns NOT_FOUND for it. + rpc.deleteSession(newestSession, null); + + // Six concurrent streams on one channel with maxRpcPerChannel=2 force a scale-up. + List> loads = new ArrayList<>(); + for (int i = 0; i < 6; i++) { + loads.add( + executor.submit( + () -> { + long rows = 0; + try (ResultSet resultSet = older.singleUse().executeQuery(SELECT1AND2)) { + while (resultSet.next()) { + rows++; + } + } + return rows; + })); + } + + // Priming recovers on the older session within the pool's own retries: the first attempt + // of every channel of the scale-up event uses the newest session, fails with NOT_FOUND and + // evicts it, and the retries after the pool's backoff all use the older session. + awaitCondition(() -> pool.getNumberOfChannels() > 1, Duration.ofSeconds(15)); + for (Future load : loads) { + assertEquals(2L, load.get().longValue()); + } + List primes = primeRequests(); + assertThat(primes.size()).isAtLeast(2); + assertEquals(newestSession, primes.get(0).getSession()); + int newestSessionPrimes = 0; + int olderSessionPrimes = 0; + for (ExecuteSqlRequest prime : primes) { + if (prime.getSession().equals(newestSession)) { + // The pool may add more than one channel per scale-up event, and their first attempts + // all start before the first failure evicts the newest session. No attempt uses the + // newest session once a retry has started. + assertEquals(0, olderSessionPrimes); + newestSessionPrimes++; + } else { + assertEquals(olderSession, prime.getSession()); + olderSessionPrimes++; + } + } + assertThat(newestSessionPrimes).isAtLeast(1); + // Every scaled-up channel was published after one successful retry on the older session. + assertEquals(pool.getNumberOfChannels() - 1, olderSessionPrimes); + assertEquals(olderSession, primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(1); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testLateMultiplexedCreateSessionAfterUnregisterNeverRecordsPrimeSession() + throws Exception { + SpannerOptions options = + createSpannerOptions().toBuilder() + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); + Session initial = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); + assertEquals(initial.getName(), primer.getPrimeSessionName()); + + // A refresh of the multiplexed session is still in flight when the database client that + // owns it is retired, which unregisters it as the owner of its database. + mockSpanner.freeze(); + Future refresh = + executor.submit( + () -> rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true)); + awaitCondition( + () -> mockSpanner.getRequestsOfType(CreateSessionRequest.class).size() == 2, + Duration.ofSeconds(15)); + rpc.unregisterChannelPrimeOwner("DATABASE_NAME", 1L); + assertNull(primer.getPrimeSessionName()); + assertThat(primer.getPrimeOwners()).isEmpty(); + + // The late response of the retired client is dropped. + mockSpanner.unfreeze(); + Session late = refresh.get(15, TimeUnit.SECONDS); + assertTrue(late.getMultiplexed()); + assertNull(primer.getPrimeSessionName()); + + // The replacement client registers its session normally, and it stays selected. + rpc.registerChannelPrimeOwner("DATABASE_NAME", 2L); + Session replacement = rpc.createSession("DATABASE_NAME", null, null, primeOwner(2L), true); + assertEquals(replacement.getName(), primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(1); + // An even later response of the retired client, for example a retried refresh, finds the + // replacement as the owner and is dropped as well. + Session evenLater = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); + assertTrue(evenLater.getMultiplexed()); + assertEquals(replacement.getName(), primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(1); + } finally { + mockSpanner.unfreeze(); + executor.shutdownNow(); + rpc.shutdown(); + } + } + + @Test + public void testDynamicChannelPoolPrimingRotatesToOlderSessionWhenNewestIsDenied() + throws Exception { + // The newest database rejects the priming query with PERMISSION_DENIED, for example because + // the caller lost access to it. The mock records the denied primes itself, because it does not + // pass them on to the regular handler. + String deniedDatabasePrefix = "projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE2]/"; + List deniedPrimes = new CopyOnWriteArrayList<>(); + MockSpannerServiceImpl denyingSpanner = + new MockSpannerServiceImpl() { + @Override + public void executeSql( + ExecuteSqlRequest request, + StreamObserver responseObserver) { + if (request.getSql().equals(DynamicChannelPoolPrimer.PRIME_SQL) + && request.getSession().startsWith(deniedDatabasePrefix)) { + deniedPrimes.add(request); + responseObserver.onError( + Status.PERMISSION_DENIED + .withDescription("Caller is missing IAM permission spanner.databases.select") + .asRuntimeException()); + return; + } + super.executeSql(request, responseObserver); + } + }; + denyingSpanner.setAbortProbability(0.0D); + denyingSpanner.putStatementResult(StatementResult.query(SELECT1AND2, SELECT1_RESULTSET)); + denyingSpanner.putStatementResult( + StatementResult.query( + Statement.of(DynamicChannelPoolPrimer.PRIME_SQL), + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues( + com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata(SELECT1AND2_METADATA) + .build())); + // Keep the load queries open long enough to trigger a scale-up. + denyingSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(3000, 0)); + Server denyingServer = + NettyServerBuilder.forAddress(new InetSocketAddress("localhost", 0)) + .addService(denyingSpanner) + .build() + .start(); + ExecutorService executor = Executors.newFixedThreadPool(8); + try { + SpannerOptions options = + createSpannerOptions().toBuilder() + .setHost("http://localhost:" + denyingServer.getPort()) + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .setGcpChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setInitSize(1) + .setMinSize(1) + .setMaxSize(4) + .setDynamicScaling(1, 2, Duration.ofMinutes(3)) + .build()) + .build(); + try (Spanner spanner = options.getService()) { + GapicSpannerRpc rpc = getRpc(spanner); + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + GcpManagedChannel pool = findGrpcGcpChannel(rpc); + assertNotNull(pool); + + DatabaseClient older = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + awaitCondition(() -> primer.getPrimeSessionName() != null, Duration.ofSeconds(15)); + String olderSession = primer.getPrimeSessionName(); + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE2]")); + awaitCondition(() -> primer.getPrimeSessions().size() == 2, Duration.ofSeconds(15)); + String newestSession = primer.getPrimeSessionName(); + assertThat(newestSession).startsWith(deniedDatabasePrefix); + + // Six concurrent streams on one channel with maxRpcPerChannel=2 force a scale-up. + List> loads = new ArrayList<>(); + for (int i = 0; i < 6; i++) { + loads.add( + executor.submit( + () -> { + long rows = 0; + try (ResultSet resultSet = older.singleUse().executeQuery(SELECT1AND2)) { + while (resultSet.next()) { + rows++; + } + } + return rows; + })); + } + + // The first attempt uses the newest session, is denied and evicts exactly that entry; + // the pool's retry succeeds on the older candidate. + awaitCondition(() -> pool.getNumberOfChannels() > 1, Duration.ofSeconds(15)); + for (Future load : loads) { + assertEquals(2L, load.get().longValue()); + } + assertThat(deniedPrimes).isNotEmpty(); + for (ExecuteSqlRequest denied : deniedPrimes) { + assertEquals(newestSession, denied.getSession()); + } + List succeededPrimes = new ArrayList<>(); + for (ExecuteSqlRequest request : + denyingSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { + if (request.getSql().equals(DynamicChannelPoolPrimer.PRIME_SQL)) { + succeededPrimes.add(request); + } + } + assertEquals(pool.getNumberOfChannels() - 1, succeededPrimes.size()); + for (ExecuteSqlRequest prime : succeededPrimes) { + assertEquals(olderSession, prime.getSession()); + } + assertEquals(olderSession, primer.getPrimeSessionName()); + assertThat(primer.getPrimeSessions()).hasSize(1); + } + } finally { + executor.shutdownNow(); + denyingServer.shutdownNow(); + denyingServer.awaitTermination(); + } + } + + @Test + public void testMultiplexedCreateSessionWithoutEchoedFlagNeverRecordsPrimeSession() + throws Exception { + // A backend that ignores the multiplexed flag of the request creates a regular session, which + // is not safe for the concurrent use that priming implies and must never become the prime + // session. + MockSpannerServiceImpl nonEchoingSpanner = + new MockSpannerServiceImpl() { + @Override + public void createSession( + CreateSessionRequest request, StreamObserver responseObserver) { + super.createSession( + request, + new StreamObserver() { + @Override + public void onNext(Session session) { + responseObserver.onNext(session.toBuilder().setMultiplexed(false).build()); + } + + @Override + public void onError(Throwable t) { + responseObserver.onError(t); + } + + @Override + public void onCompleted() { + responseObserver.onCompleted(); + } + }); + } + }; + Server nonEchoingServer = + NettyServerBuilder.forAddress(new InetSocketAddress("localhost", 0)) + .addService(nonEchoingSpanner) + .build() + .start(); + try { + SpannerOptions options = + createSpannerOptions().toBuilder() + .setHost("http://localhost:" + nonEchoingServer.getPort()) + .enableGrpcGcpExtension() + .enableDynamicChannelPool() + .build(); + GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); + try { + DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); + assertNotNull(primer); + rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); + + Session session = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); + assertFalse(session.getMultiplexed()); + assertTrue( + nonEchoingSpanner + .getRequestsOfType(CreateSessionRequest.class) + .get(0) + .getSession() + .getMultiplexed()); + assertNull(primer.getPrimeSessionName()); + } finally { + rpc.shutdown(); + } + } finally { + nonEchoingServer.shutdownNow(); + nonEchoingServer.awaitTermination(); + } + } + @Test public void testChannelEndpointCacheFactoryUsedWhenLocationApiEnabled() { AtomicBoolean factoryCalled = new AtomicBoolean(false); @@ -1221,9 +2015,10 @@ public void testGrpcGcpExtensionPreservesChannelConfigurator() throws Exception GapicSpannerRpc.class.getDeclaredMethod( "maybeEnableGrpcGcpExtension", InstantiatingGrpcChannelProvider.Builder.class, - SpannerOptions.class); + SpannerOptions.class, + DynamicChannelPoolPrimer.class); method.setAccessible(true); - method.invoke(null, channelProviderBuilder, options); + method.invoke(null, channelProviderBuilder, options, null); ApiFunction chainedConfigurator = channelProviderBuilder.getChannelConfigurator(); @@ -1583,33 +2378,69 @@ private static void countGrpcGcpObjectsFromChannelzField( private static void countGrpcGcpObjects( Object object, Set visited, GrpcGcpObjectCounts counts) { + visitObjectGraph( + object, + visited, + visited1 -> { + if (visited1 instanceof GcpManagedChannel) { + counts.gcpManagedChannels++; + } + if (visited1.getClass().getName().equals(GRPC_GCP_CHANNEL_REF_CLASS_NAME)) { + counts.channelRefs++; + } + }); + } + + /** + * Returns the grpc-gcp channel pool that serves the Spanner stub of the given rpc. GAX wraps the + * pool in its own channel pool and interceptor channels, so it is located by walking the object + * graph of the rpc. + */ + @Nullable + private static GcpManagedChannel findGrpcGcpChannel(GapicSpannerRpc rpc) { + java.util.concurrent.atomic.AtomicReference found = + new java.util.concurrent.atomic.AtomicReference<>(); + visitObjectGraph( + rpc, + Collections.newSetFromMap(new IdentityHashMap<>()), + object -> { + if (object instanceof GcpManagedChannel) { + found.compareAndSet(null, (GcpManagedChannel) object); + } + }); + return found.get(); + } + + private static void visitObjectGraph( + Object object, Set visited, java.util.function.Consumer visitor) { if (object == null || !visited.add(object)) { return; } - if (object instanceof GcpManagedChannel) { - counts.gcpManagedChannels++; - } + visitor.accept(object); Class clazz = object.getClass(); - if (clazz.getName().equals(GRPC_GCP_CHANNEL_REF_CLASS_NAME)) { - counts.channelRefs++; + if (object instanceof java.util.concurrent.atomic.AtomicReference) { + // JDK internals are not reflectively accessible; unwrap the value instead. + visitObjectGraph( + ((java.util.concurrent.atomic.AtomicReference) object).get(), visited, visitor); + return; } if (object instanceof Collection) { for (Object value : (Collection) object) { - countGrpcGcpObjects(value, visited, counts); + visitObjectGraph(value, visited, visitor); } return; } if (object instanceof Map) { for (Map.Entry entry : ((Map) object).entrySet()) { - countGrpcGcpObjects(entry.getKey(), visited, counts); - countGrpcGcpObjects(entry.getValue(), visited, counts); + visitObjectGraph(entry.getKey(), visited, visitor); + visitObjectGraph(entry.getValue(), visited, visitor); } return; } if (clazz.isArray()) { int length = Array.getLength(object); for (int i = 0; i < length; i++) { - countGrpcGcpObjects(Array.get(object, i), visited, counts); + visitObjectGraph(Array.get(object, i), visited, visitor); } return; } @@ -1623,7 +2454,7 @@ private static void countGrpcGcpObjects( } try { field.setAccessible(true); - countGrpcGcpObjects(field.get(object), visited, counts); + visitObjectGraph(field.get(object), visited, visitor); } catch (RuntimeException | IllegalAccessException ignored) { // Ignore fields that are not reflectively accessible in this runtime. } From f25a550da84c683715cb33312b01a1db73dac159 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 12:31:08 +0530 Subject: [PATCH 2/7] fix tests --- .../google/cloud/spanner/SpannerOptions.java | 42 +++++++++---------- .../spi/v1/DynamicChannelPoolPrimer.java | 17 +------- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 1 - .../spanner/TransactionManagerImplTest.java | 4 +- .../spanner/TransactionRunnerImplTest.java | 2 +- .../spi/v1/DynamicChannelPoolPrimerTest.java | 41 ++---------------- 6 files changed, 27 insertions(+), 80 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index d2b9d2c5b4b8..7e4116e50a23 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -180,8 +180,8 @@ public class SpannerOptions extends ServiceOptions { /** * Default maximum time for one attempt to prime a channel that the dynamic channel pool adds - * during scale-up. When multiplexed sessions are enabled, scaled-up channels are primed by - * executing {@code SELECT 1} with the multiplexed session before they are published to the pool. + * during scale-up. Scaled-up channels are primed by executing {@code SELECT 1} with a multiplexed + * session before they are published to the pool. */ public static final Duration DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_TIMEOUT = Duration.ofSeconds(10); @@ -210,13 +210,12 @@ public class SpannerOptions extends ServiceOptions { *
  • Channel prime max attempts: {@value #DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS} * * - *

    When multiplexed sessions are enabled (the default), channels that the pool adds during - * scale-up are primed with {@code SELECT 1} on the multiplexed session before they are published. - * The primer is registered by the Spanner client when dynamic channel pooling is enabled, unless - * these options already contain a primer. Priming uses the most recently created multiplexed - * session of any database client of the {@link Spanner} instance; a session that turns out to be - * invalid is dropped and the next most recent one is used. When multiplexed sessions are - * disabled, scaled-up channels are published without priming. + *

    Channels that the pool adds during scale-up are primed with {@code SELECT 1} on a + * multiplexed session before they are published. The primer is registered by the Spanner client + * when dynamic channel pooling is enabled, unless these options already contain a primer. Priming + * uses the most recently created multiplexed session of any database client of the {@link + * Spanner} instance; a session that turns out to be invalid is dropped and the next most recent + * one is used. * * @return a new {@link GcpChannelPoolOptions} instance with Spanner defaults */ @@ -2057,13 +2056,11 @@ public Builder disableGrpcGcpExtension() { * Enables dynamic channel pooling. When enabled, the client will automatically scale the number * of channels based on load. This requires the gRPC-GCP extension to be enabled. * - *

    When multiplexed sessions are enabled (the default), channels that the pool adds during - * scale-up are primed before they serve traffic: the client executes {@code SELECT 1} with the - * multiplexed session on the new channel, and the pool only publishes the channel once that - * succeeds. When multiplexed sessions are disabled, scaled-up channels are published without - * priming. See {@link #createDefaultDynamicChannelPoolOptions()} for the prime timeout and - * attempt defaults, and {@link #setGcpChannelPoolOptions(GcpChannelPoolOptions)} to customize - * them. + *

    Channels that the pool adds during scale-up are primed before they serve traffic: the + * client executes {@code SELECT 1} with a multiplexed session on the new channel, and the pool + * only publishes the channel once that succeeds. See {@link + * #createDefaultDynamicChannelPoolOptions()} for the prime timeout and attempt defaults, and + * {@link #setGcpChannelPoolOptions(GcpChannelPoolOptions)} to customize them. * *

    Dynamic channel pooling is disabled by default. Use this method to explicitly enable it. * Note that calling {@link #setNumChannels(int)} will disable dynamic channel pooling even if @@ -2105,13 +2102,12 @@ public Builder setGrpcGcpOtelMetricsEnabled(boolean enableGrpcGcpOtelMetrics) { * #createDefaultDynamicChannelPoolOptions()}). Values that are left unset in the given options * are filled in from those defaults. * - *

    When multiplexed sessions are enabled (the default), channels that the pool adds during - * scale-up are primed with {@code SELECT 1} on the multiplexed session before they are - * published; otherwise they are published without priming. A channel primer, prime timeout, or - * prime attempt count that is set in the given options takes precedence over the Spanner primer - * and its defaults. The deadline of the Spanner primer's RPC is derived from the prime timeout - * and normally stays below it; only a prime timeout of two milliseconds or less yields the - * primer's minimum deadline of one millisecond, and the prime timeout then bounds the attempt. + *

    Channels that the pool adds during scale-up are primed with {@code SELECT 1} on a + * multiplexed session before they are published. A channel primer, prime timeout, or prime + * attempt count that is set in the given options takes precedence over the Spanner primer and + * its defaults. The deadline of the Spanner primer's RPC is derived from the prime timeout and + * normally stays below it; only a prime timeout of two milliseconds or less yields the primer's + * minimum deadline of one millisecond, and the prime timeout then bounds the attempt. * *

    Example usage: * diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java index 543d6a2e3b22..c69b2269155a 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -20,7 +20,6 @@ import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Spanner; import com.google.cloud.spanner.SpannerExceptionFactory; -import com.google.cloud.spanner.SpannerOptions; import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; import com.google.cloud.spanner.XGoogSpannerRequestId; import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; @@ -62,10 +61,7 @@ * Primes channels that the grpc-gcp dynamic channel pool adds during scale-up by executing {@code * SELECT 1} on the new channel with a multiplexed session, so the transport, TLS session, and * server-side connection are established before the channel serves live traffic. Mirrors the - * priming that the Go Spanner client performs for its dynamic channel pool. Priming requires - * multiplexed sessions: when they are disabled on the {@link SpannerOptions} there is nothing to - * prime with, and scaled-up channels are published unprimed, which is the behavior of the pool - * without a primer. + * priming that the Go Spanner client performs for its dynamic channel pool. * *

    The channel pool is shared by all database clients of a {@link Spanner} instance, while * multiplexed sessions are created per database. The primer therefore keeps a small registry of @@ -143,7 +139,6 @@ final class DynamicChannelPoolPrimer implements GcpChannelPrimer { @Nullable private final CallCredentials defaultCallCredentials; @Nullable private final CallCredentialsProvider callCredentialsProvider; private final boolean routeToLeader; - private final boolean multiplexedSessionsEnabled; private final Duration rpcDeadline; /** Source of the creation generations of {@link PrimeSession} entries. */ @@ -233,7 +228,6 @@ public String toString() { * @param callCredentialsProvider the optional user-supplied provider that takes precedence over * {@code defaultCallCredentials} for each call, exactly like for normal Spanner calls * @param routeToLeader whether normal Spanner calls carry the route-to-leader header - * @param multiplexedSessionsEnabled whether the client creates multiplexed sessions at all * @param rpcDeadline the deadline of a single priming RPC */ DynamicChannelPoolPrimer( @@ -243,7 +237,6 @@ public String toString() { @Nullable CallCredentials defaultCallCredentials, @Nullable CallCredentialsProvider callCredentialsProvider, boolean routeToLeader, - boolean multiplexedSessionsEnabled, Duration rpcDeadline) { this.metadataProvider = Preconditions.checkNotNull(metadataProvider); this.projectName = Preconditions.checkNotNull(projectName); @@ -251,7 +244,6 @@ public String toString() { this.defaultCallCredentials = defaultCallCredentials; this.callCredentialsProvider = callCredentialsProvider; this.routeToLeader = routeToLeader; - this.multiplexedSessionsEnabled = multiplexedSessionsEnabled; Preconditions.checkArgument( rpcDeadline != null && !rpcDeadline.isZero() && !rpcDeadline.isNegative(), "rpcDeadline must be positive"); @@ -443,13 +435,6 @@ Map getPrimeOwners() { @Override public ListenableFuture prime(ManagedChannel channel) { - if (!multiplexedSessionsEnabled) { - // Priming executes SELECT 1 with a multiplexed session, because that is the only session - // that lives at the level of the shared channel pool rather than inside a session pool. - // Without multiplexed sessions there is nothing to prime with, so the channel is published - // unprimed, which is the behavior of a pool without a primer. - return Futures.immediateFuture(null); - } PrimeSession entry = getPrimeSession(); if (entry == null) { // The Go client refuses to scale up at all before a multiplexed session exists. The Java diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index 48278c552515..a48f9fef3aeb 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -834,7 +834,6 @@ private DynamicChannelPoolPrimer createChannelPrimer( defaultCallCredentials, callCredentialsProvider, leaderAwareRoutingEnabled, - options.getSessionPoolOptions().getUseMultiplexedSession(), // The options already contain the merged prime timeout, including a user-provided one. DynamicChannelPoolPrimer.rpcDeadlineFor( options.getGcpChannelPoolOptions().getChannelPrimeTimeout())); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index 547f6b70a22d..19c03859e669 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -241,7 +241,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.eq(null), + Mockito.anyMap(), Mockito.eq(true))) .thenAnswer( invocation -> @@ -324,7 +324,7 @@ public void inlineBegin() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.eq(null), + Mockito.anyMap(), Mockito.eq(true))) .thenAnswer( invocation -> diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 1dd2418aa05a..c0bd9d1f3169 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -232,7 +232,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.eq(null), + Mockito.anyMap(), Mockito.eq(true))) .thenAnswer( invocation -> diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java index fccbaa4a14bd..833414352150 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java @@ -217,7 +217,6 @@ private static CallCredentials credentialsWithToken(String token) { private DynamicChannelPoolPrimer newPrimer( @Nullable CallCredentialsProvider callCredentialsProvider, boolean routeToLeader, - boolean multiplexedSessionsEnabled, Duration rpcDeadline) { // The fixed headers of the client are configured on the channel and must not be sent by the // primer: it only sends the headers that a normal call adds per call. @@ -232,16 +231,12 @@ private DynamicChannelPoolPrimer newPrimer( credentialsWithToken(DEFAULT_TOKEN), callCredentialsProvider, routeToLeader, - multiplexedSessionsEnabled, rpcDeadline); } private DynamicChannelPoolPrimer newPrimer() { return newPrimer( - /* callCredentialsProvider= */ null, - /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ true, - Duration.ofSeconds(5)); + /* callCredentialsProvider= */ null, /* routeToLeader= */ true, Duration.ofSeconds(5)); } private static final AtomicLong OWNER_TICKETS = new AtomicLong(); @@ -366,10 +361,7 @@ public void everyPrimeCarriesFreshRequestIdWithFirstAttempt() throws Exception { public void primeOmitsRouteToLeaderHeaderWhenDisabled() throws Exception { DynamicChannelPoolPrimer primer = newPrimer( - /* callCredentialsProvider= */ null, - /* routeToLeader= */ false, - /* multiplexedSessionsEnabled= */ true, - Duration.ofSeconds(5)); + /* callCredentialsProvider= */ null, /* routeToLeader= */ false, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); @@ -383,7 +375,6 @@ public void primePrefersCallCredentialsProviderOverDefaultCredentials() throws E newPrimer( () -> credentialsWithToken(PROVIDER_TOKEN), /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ true, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); @@ -396,11 +387,7 @@ public void primePrefersCallCredentialsProviderOverDefaultCredentials() throws E @Test public void primeFallsBackToDefaultCredentialsWhenProviderReturnsNull() throws Exception { DynamicChannelPoolPrimer primer = - newPrimer( - () -> null, - /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ true, - Duration.ofSeconds(5)); + newPrimer(() -> null, /* routeToLeader= */ true, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); @@ -724,10 +711,7 @@ public void primeFailsWithDeadlineExceededWhenServerDoesNotRespond() throws Exce service.holdResponses = true; DynamicChannelPoolPrimer primer = newPrimer( - /* callCredentialsProvider= */ null, - /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ true, - Duration.ofMillis(200)); + /* callCredentialsProvider= */ null, /* routeToLeader= */ true, Duration.ofMillis(200)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); Throwable failure = failureOf(primer.prime(newChannel())); @@ -799,22 +783,6 @@ public void primeFailsWithFailedPreconditionWithoutSession() throws Exception { assertThat(service.requests).isEmpty(); } - @Test - public void primeSucceedsWithoutRpcWhenMultiplexedSessionsAreDisabled() throws Exception { - DynamicChannelPoolPrimer primer = - newPrimer( - /* callCredentialsProvider= */ null, - /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ false, - Duration.ofSeconds(5)); - - ListenableFuture future = primer.prime(newChannel()); - - assertThat(future.isDone()).isTrue(); - assertThat(getWithin(future, Duration.ofSeconds(1))).isNull(); - assertThat(service.requests).isEmpty(); - } - @Test public void concurrentPrimesOnOneChannelAllSucceed() throws Exception { DynamicChannelPoolPrimer primer = newPrimer(); @@ -962,7 +930,6 @@ public void rpcDeadlineNormallyStaysBelowPoolPrimeTimeout() { newPrimer( /* callCredentialsProvider= */ null, /* routeToLeader= */ true, - /* multiplexedSessionsEnabled= */ true, DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofNanos(1))) .getRpcDeadline()) .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); From 80bc81a83fb5d76ab035fc739815281d0ffbbae8 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 13:46:53 +0530 Subject: [PATCH 3/7] fix built-in metric test --- .../spanner/BuiltInOpenTelemetryMetricsProviderTest.java | 4 ++++ .../cloud/spanner/MultiplexedSessionDatabaseClientTest.java | 3 +++ 2 files changed, 7 insertions(+) diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java index 1f2adb41a218..deb5ad1511e0 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java @@ -110,6 +110,10 @@ private void verifyHash(String hash) { private SpannerOptions newTestOptions() { return SpannerOptions.newBuilder() + // The builder picks up SPANNER_EMULATOR_HOST from the environment, and building with an + // emulator host replaces the credentials below with NoCredentials. Built-in metrics are + // disabled for a client without credentials, so leave the emulator out of these tests. + .setEmulatorHost(null) .setProjectId("host-project") .setCredentials( OAuth2Credentials.create(new AccessToken("test-token", new Date(Long.MAX_VALUE)))) diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index 8a99746084de..f9b4e19cb699 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -145,6 +145,7 @@ public void testMaintainer() { @Test public void testClosedClientIgnoresInitialSessionThatArrivesAfterClose() { + assumeTrue(isJava8()); Clock clock = mock(Clock.class); when(clock.instant()).thenReturn(Instant.now()); SessionClient sessionClient = mock(SessionClient.class); @@ -192,6 +193,7 @@ public void testClosedClientIgnoresInitialSessionThatArrivesAfterClose() { @Test public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() { + assumeTrue(isJava8()); Instant now = Instant.now(); Clock clock = mock(Clock.class); when(clock.instant()).thenReturn(now); @@ -473,6 +475,7 @@ public void testCloseKeepsChannelUsageEntryWhileAnotherClientIsUsingSameSpanner( @Test public void testChannelPrimeOwnerTicketIsRegisteredCarriedAndUnregistered() { + assumeTrue(isJava8()); Clock clock = mock(Clock.class); when(clock.instant()).thenReturn(Instant.now()); SessionClient sessionClient = mock(SessionClient.class); From 2fcd4158b1764677acb95742b2eaaeea045d1836 Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 17:30:21 +0530 Subject: [PATCH 4/7] incorporate suggestions --- .../MultiplexedSessionDatabaseClient.java | 105 ++-- .../google/cloud/spanner/SessionClient.java | 44 +- .../com/google/cloud/spanner/SpannerImpl.java | 5 +- .../google/cloud/spanner/SpannerOptions.java | 12 +- .../spi/v1/DynamicChannelPoolPrimer.java | 436 +++----------- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 77 ++- .../spi/v1/SpannerMetadataProvider.java | 3 + .../cloud/spanner/spi/v1/SpannerRpc.java | 36 +- .../DynamicChannelPoolPrimeSessionTest.java | 33 +- .../MultiplexedSessionDatabaseClientTest.java | 101 ++-- .../google/cloud/spanner/SpannerImplTest.java | 44 +- .../spanner/TransactionManagerImplTest.java | 4 +- .../spanner/TransactionRunnerImplTest.java | 2 +- .../spanner/spi/v1/ChannelPrimerTestRpc.java | 30 +- .../spi/v1/DynamicChannelPoolPrimerTest.java | 533 ++++-------------- .../spanner/spi/v1/GapicSpannerRpcTest.java | 437 +------------- 16 files changed, 421 insertions(+), 1481 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 8b1e768d680e..9d3d507a0355 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -27,11 +27,11 @@ import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionClient.SessionOption; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; -import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.cloud.spanner.spi.v1.SpannerRpc.ChannelPrimeSessionSource; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.Futures; import com.google.spanner.v1.BatchWriteResponse; import java.time.Clock; import java.time.Duration; @@ -40,6 +40,7 @@ import java.util.EnumSet; import java.util.HashMap; import java.util.Map; +import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -50,12 +51,14 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; /** * {@link DatabaseClient} implementation that uses a single multiplexed session to execute * transactions. */ -final class MultiplexedSessionDatabaseClient extends AbstractMultiplexedSessionDatabaseClient { +final class MultiplexedSessionDatabaseClient extends AbstractMultiplexedSessionDatabaseClient + implements ChannelPrimeSessionSource { /** * The maximum number of attempts that the client will try to execute CreateSession for the * initial multiplexed session. This value is only used for the very first multiplexed session @@ -188,7 +191,7 @@ private SharedChannelUsage(int numChannels) { */ private final AtomicInteger numCurrentSingleUseTransactions = new AtomicInteger(); - private boolean isClosed; + private volatile boolean isClosed; /** The duration before we try to replace the multiplexed session. The default is 7 days. */ private final Duration sessionExpirationDuration; @@ -222,18 +225,6 @@ private SharedChannelUsage(int numChannels) { private final AtomicLong numSessionsReleased = new AtomicLong(); - /** Source of the {@link #channelPrimeOwnerTicket owner tickets} of all clients of the process. */ - private static final AtomicLong CHANNEL_PRIME_OWNER_TICKETS = new AtomicLong(); - - /** - * The ticket under which this client owns its database for priming channels that the dynamic - * channel pool adds during scale-up. The ticket is registered with the rpc before the first - * multiplexed session is created, carried by every CreateSession of this client, and unregistered - * when this client is closed, so a session that arrives after the close is never used for - * priming. - */ - private final long channelPrimeOwnerTicket = CHANNEL_PRIME_OWNER_TICKETS.incrementAndGet(); - MultiplexedSessionDatabaseClient(SessionClient sessionClient) { this(sessionClient, Clock.systemUTC()); } @@ -267,12 +258,14 @@ private SharedChannelUsage(int numChannels) { final SettableApiFuture initialSessionReferenceFuture = SettableApiFuture.create(); this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); - // Register as the owner of the database before the first CreateSession, so that session and - // every refreshed session can be attributed to this client for channel priming. - spanner - .getRpc() - .registerChannelPrimeOwner( - sessionClient.getDatabaseId().getName(), channelPrimeOwnerTicket); + // This escape is safe: resourceNotFoundException is initialized at declaration, + // multiplexedSessionReference above, and isClosed is volatile: the fields + // getChannelPrimeSessionName() reads. + // Registration adds to a CopyOnWriteArrayList under a lock, whose volatile array write happens + // before a primer reads that array, so this is safe publication, not a racy final-field escape. + // Keep registration before the first CreateSession and inside this constructor: clients are + // constructed at multiple call sites, and a missed registration would silently disable priming. + spanner.getRpc().registerChannelPrimeSessionSource(this); try { Duration waitDuration = @@ -287,18 +280,13 @@ private SharedChannelUsage(int numChannels) { // The constructor did not complete, so the caller never gets a reference to this client and // will never close it. Undo everything that was registered above: mark the client closed so // that a CreateSession that is still in flight neither starts the maintainer nor is handed - // to a waiter, stop the maintainer, drop the channel prime owner ticket so that a late - // session is not recorded for priming, and release the shared channel usage. + // to a waiter, stop the maintainer, unregister the prime-session source, and release the + // shared channel usage. close(); throw t; } } - /** The options of every CreateSession of this client: they carry the owner ticket. */ - private Map createSessionOptions() { - return SessionClient.optionMap(SessionOption.channelPrimeOwner(channelPrimeOwnerTicket)); - } - private void asyncCreateMultiplexedSession( SettableApiFuture sessionReferenceFuture, int remainingAttempts) { this.sessionClient.asyncCreateMultiplexedSession( @@ -351,8 +339,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount asyncCreateMultiplexedSession(future, remainingAttempts - 1); } } - }, - createSessionOptions()); + }); } private void maybeWaitForSessionCreation( @@ -427,6 +414,27 @@ AtomicLong getNumSessionsReleased() { return this.numSessionsReleased; } + @Override + @Nullable + public String getChannelPrimeSessionName() { + // This client already uses the returned session as multiplexed for all real traffic. If the + // backend failed to honor the multiplexed request flag, withholding it from the primer would + // not make the client safe. + if (isClosed || !isValid()) { + return null; + } + ApiFuture future = multiplexedSessionReference.get(); + if (!future.isDone() || future.isCancelled()) { + return null; + } + try { + SessionReference session = Futures.getDone(future); + return isClosed || !isValid() ? null : session.getName(); + } catch (CancellationException | ExecutionException e) { + return null; + } + } + void close() { boolean releaseChannelUsage = false; synchronized (this) { @@ -437,13 +445,9 @@ void close() { } } if (releaseChannelUsage) { - // The multiplexed session is no longer maintained, so it must no longer be used for priming - // dynamic channel pool channels, and a CreateSession of this client that is still in flight - // must not register its session either. - spanner - .getRpc() - .unregisterChannelPrimeOwner( - sessionClient.getDatabaseId().getName(), channelPrimeOwnerTicket); + // The multiplexed session is no longer maintained, so this client must no longer be asked + // for a session when priming dynamic channel pool channels. + spanner.getRpc().unregisterChannelPrimeSessionSource(this); synchronized (CHANNEL_USAGE) { SharedChannelUsage sharedChannelUsage = CHANNEL_USAGE.get(this.spanner); if (sharedChannelUsage != null) { @@ -722,17 +726,19 @@ void maintain() { new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { - if (isClientClosed()) { - // The client was closed while the session was being refreshed. The refreshed - // session belongs to a client that no longer exists and is ignored. - return; + synchronized (MultiplexedSessionDatabaseClient.this) { + if (isClosed) { + // The client was closed while the session was being refreshed. The refreshed + // session belongs to a client that no longer exists and is ignored. + return; + } + multiplexedSessionReference.set( + ApiFutures.immediateFuture(session.getSessionReference())); + expirationDate.set( + clock + .instant() + .plus(MultiplexedSessionDatabaseClient.this.sessionExpirationDuration)); } - multiplexedSessionReference.set( - ApiFutures.immediateFuture(session.getSessionReference())); - expirationDate.set( - clock - .instant() - .plus(MultiplexedSessionDatabaseClient.this.sessionExpirationDuration)); } @Override @@ -741,8 +747,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount // we continue to use the session that has passed its expiration date for now, and // that a new attempt at creating a new session will be done in 10 minutes from now. } - }, - createSessionOptions()); + }); } } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java index e055faa6b636..234675a221bc 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java @@ -34,7 +34,6 @@ import java.util.Map; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; -import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; /** Client for creating single sessions and batches of sessions. */ @@ -95,10 +94,6 @@ static SessionOption channelAffinityRef(ChannelAffinityRef channelAffinityRef) { return new SessionOption(SpannerRpc.Option.CHANNEL_ID_AFFINITY, channelAffinityRef); } - static SessionOption channelPrimeOwner(long ownerTicket) { - return new SessionOption(SpannerRpc.Option.CHANNEL_PRIME_OWNER, ownerTicket); - } - SpannerRpc.Option rpcOption() { return rpcOption; } @@ -269,19 +264,8 @@ SessionImpl createSession() { * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available. */ void createMultiplexedSession(SessionConsumer consumer) { - createMultiplexedSession(consumer, /* options= */ null); - } - - /** - * Creates a multiplexed session with the given {@link SpannerRpc.Option options} for the {@code - * CreateSession} call and returns it to the given {@link SessionConsumer}. - * - * @see #createMultiplexedSession(SessionConsumer) - */ - void createMultiplexedSession( - SessionConsumer consumer, @Nullable Map options) { try { - SessionImpl sessionImpl = createMultiplexedSession(options); + SessionImpl sessionImpl = createMultiplexedSession(); consumer.onSessionReady(sessionImpl); } catch (Throwable t) { consumer.onSessionCreateFailure(t, 1); @@ -293,16 +277,6 @@ void createMultiplexedSession( * GRPC channel. In case of an error during the gRPC calls, an exception will be thrown. */ SessionImpl createMultiplexedSession() { - return createMultiplexedSession((Map) null); - } - - /** - * Creates a multiplexed session with the given {@link SpannerRpc.Option options} for the {@code - * CreateSession} call and returns it. - * - * @see #createMultiplexedSession() - */ - SessionImpl createMultiplexedSession(@Nullable Map options) { ISpan span = spanner .getTracer() @@ -315,7 +289,7 @@ SessionImpl createMultiplexedSession(@Nullable Map options db.getName(), spanner.getOptions().getDatabaseRole(), spanner.getOptions().getSessionLabels(), - options, + null, true); SessionImpl sessionImpl = new SessionImpl( @@ -345,13 +319,10 @@ SessionImpl createMultiplexedSession(@Nullable Map options * SessionConsumer#onSessionCreateFailure(Throwable, int)} call with the error. * * @param consumer The {@link SessionConsumer} to use for callbacks when sessions are available. - * @param options The {@link SpannerRpc.Option options} of the {@code CreateSession} call, or - * {@code null} for none. */ - void asyncCreateMultiplexedSession( - SessionConsumer consumer, @Nullable Map options) { + void asyncCreateMultiplexedSession(SessionConsumer consumer) { try { - executor.submit(new CreateMultiplexedSessionsRunnable(consumer, options)); + executor.submit(new CreateMultiplexedSessionsRunnable(consumer)); } catch (Throwable t) { consumer.onSessionCreateFailure(t, 1); } @@ -359,18 +330,15 @@ void asyncCreateMultiplexedSession( private final class CreateMultiplexedSessionsRunnable implements Runnable { private final SessionConsumer consumer; - @Nullable private final Map options; - private CreateMultiplexedSessionsRunnable( - SessionConsumer consumer, @Nullable Map options) { + private CreateMultiplexedSessionsRunnable(SessionConsumer consumer) { Preconditions.checkNotNull(consumer); this.consumer = consumer; - this.options = options; } @Override public void run() { - createMultiplexedSession(consumer, options); + createMultiplexedSession(consumer); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index a3950ba3f382..755b1acc8290 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -298,9 +298,8 @@ public DatabaseClient getDatabaseClient(DatabaseId db) { checkClosed(); String clientId = null; if (dbClients.containsKey(db) && !dbClients.get(db).isValid()) { - // Close the invalidated client and remove it. Closing it unregisters it as the owner of - // its database for priming dynamic channel pool channels, because its multiplexed session - // is no longer maintained. + // Close the invalidated client and remove it. Closing it unregisters its multiplexed + // session source from the dynamic channel pool primer. dbClients.get(db).closeAsync(new ClosedException()); clientId = dbClients.get(db).clientId; dbClients.remove(db); diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 7e4116e50a23..6e4261d21ed0 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -213,9 +213,8 @@ public class SpannerOptions extends ServiceOptions { *

    Channels that the pool adds during scale-up are primed with {@code SELECT 1} on a * multiplexed session before they are published. The primer is registered by the Spanner client * when dynamic channel pooling is enabled, unless these options already contain a primer. Priming - * uses the most recently created multiplexed session of any database client of the {@link - * Spanner} instance; a session that turns out to be invalid is dropped and the next most recent - * one is used. + * rotates across available multiplexed sessions owned by live database clients of the {@link + * Spanner} instance. Closed or invalid database clients do not supply sessions for priming. * * @return a new {@link GcpChannelPoolOptions} instance with Spanner defaults */ @@ -276,6 +275,13 @@ static GcpChannelPoolOptions mergeWithDefaultChannelPoolOptions( if (userOptions.getCleanupInterval() == null || userOptions.getCleanupInterval().isZero()) { merged.setCleanupInterval(defaults.getCleanupInterval()); } + if (userOptions.getChannelPrimeTimeout() == null + || userOptions.getChannelPrimeTimeout().isZero()) { + merged.setChannelPrimeTimeout(defaults.getChannelPrimeTimeout()); + } + if (userOptions.getChannelPrimeMaxAttempts() <= 0) { + merged.setChannelPrimeMaxAttempts(defaults.getChannelPrimeMaxAttempts()); + } return merged.build(); } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java index c69b2269155a..0ec2f9ca2958 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -23,6 +23,7 @@ import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; import com.google.cloud.spanner.XGoogSpannerRequestId; import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; +import com.google.cloud.spanner.spi.v1.SpannerRpc.ChannelPrimeSessionSource; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; @@ -42,58 +43,37 @@ import io.grpc.ManagedChannel; import io.grpc.Metadata; import io.grpc.MethodDescriptor; -import io.grpc.Status; import io.grpc.stub.ClientCalls; import java.io.UnsupportedEncodingException; import java.net.URLEncoder; import java.time.Duration; -import java.util.HashMap; import java.util.List; -import java.util.Locale; import java.util.Map; -import java.util.Objects; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; -import java.util.regex.Pattern; +import java.util.concurrent.atomic.AtomicInteger; import javax.annotation.Nullable; /** * Primes channels that the grpc-gcp dynamic channel pool adds during scale-up by executing {@code * SELECT 1} on the new channel with a multiplexed session, so the transport, TLS session, and - * server-side connection are established before the channel serves live traffic. Mirrors the - * priming that the Go Spanner client performs for its dynamic channel pool. + * server-side connection are established before the channel serves live traffic. * *

    The channel pool is shared by all database clients of a {@link Spanner} instance, while - * multiplexed sessions are created per database. The primer therefore keeps a small registry of - * prime sessions keyed by database name. {@link GapicSpannerRpc} registers every multiplexed - * session that it successfully creates, which replaces the previous entry of that database, and the - * registry hands out the entry with the highest creation generation, that is the most recently - * created multiplexed session. A multiplexed {@code CreateSession} is always the first RPC of a - * database client before any data RPC is possible, and the periodic refresh of a multiplexed - * session goes through the same path, so the most recently created multiplexed session is normally - * a valid session to prime with. {@code SELECT 1} against any database is sufficient to prime a - * channel. + * multiplexed sessions are created per database. Each live multiplexed-session database client + * registers itself as a session source. Consecutive priming attempts rotate across those sources, + * and unavailable sources are skipped. A source exposes its current session only after creation + * completes successfully, so selecting a session never blocks the channel-pool scale-up path. + * Priming with any database's session establishes the transport, TLS session, and server-side + * connection for the whole channel, which benefits every database that later uses it. It does not + * warm database-specific server-side state, so a request served by a freshly scaled-up channel for + * a database other than the one used to prime it can still see slightly higher latency than on a + * channel that database has already used. This trade-off is accepted. * - *

    A registered session can nevertheless become invalid, for example when its database is - * dropped. The registry is therefore maintained on rare paths only: the entries of a database are - * removed when its database client is invalidated or closed, and a priming attempt that fails for a - * reason that is specific to the session it used, such as the session not being found, the session - * or its database role being invalid, or permission being denied, evicts exactly the entry that it - * used, so the pool's next attempt falls back to the next most recent session. Nothing is written - * from any data RPC, and the scale-up path performs a single volatile read to select the session. - * - *

    A {@code CreateSession} that is still in flight when its database client is retired must not - * register its session after the client has been retired. Every database client that creates - * multiplexed sessions therefore registers itself as the owner of its database through {@link - * #registerPrimeOwner(String, long)} with a unique owner ticket before its first {@code - * CreateSession}, includes that ticket in every {@code CreateSession} it issues, and unregisters - * itself through {@link #unregisterPrimeOwner(String, long)} when it is invalidated or closed. A - * session is registered only when the ticket of its {@code CreateSession} is the current owner - * ticket of its database, so a late response of a retired client finds no owner, or the ticket of - * the replacement client, and is dropped. The owner map holds one entry per live database client - * and nothing for retired ones, so it is bounded by the number of live clients. This mirrors the Go - * client, which installs a newly created multiplexed session only while its creation is still the - * current one and the session manager is still valid. + *

    A database client unregisters itself when invalidated or closed. The primer therefore never + * retains a retired client's session name, and a {@code CreateSession} that completes after its + * client was retired cannot become available for priming. Registry writes happen only when a client + * is created or closed; session selection iterates an immutable snapshot. */ final class DynamicChannelPoolPrimer implements GcpChannelPrimer { /** The statement executed on every scaled-up channel. */ @@ -119,14 +99,6 @@ final class DynamicChannelPoolPrimer implements GcpChannelPrimer { */ static final Duration MIN_RPC_DEADLINE = Duration.ofMillis(1); - private static final Metadata.Key ROUTE_TO_LEADER_KEY = - Metadata.Key.of("x-goog-spanner-route-to-leader", Metadata.ASCII_STRING_MARSHALLER); - private static final Metadata.Key REQUEST_PARAMS_KEY = - Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); - - /** Matches failure messages about the database role of the session, such as an invalid role. */ - private static final Pattern ROLE_PATTERN = Pattern.compile("\\brole\\b"); - /** * Channel number of the request id of a priming RPC. The channel that is being primed is not part * of the pool yet, so it has no pool channel id and the request id uses 0 (unknown). @@ -136,86 +108,18 @@ final class DynamicChannelPoolPrimer implements GcpChannelPrimer { private final SpannerMetadataProvider metadataProvider; private final String projectName; private final RequestIdCreator requestIdCreator; - @Nullable private final CallCredentials defaultCallCredentials; @Nullable private final CallCredentialsProvider callCredentialsProvider; - private final boolean routeToLeader; private final Duration rpcDeadline; - /** Source of the creation generations of {@link PrimeSession} entries. */ - private final AtomicLong generations = new AtomicLong(); - - /** - * The registered prime sessions, at most one per database, sorted by descending generation. The - * list is replaced as a whole on every write, so a read is a single volatile read. - */ - private volatile ImmutableList primeSessions = ImmutableList.of(); - /** - * The owner ticket of the live database client of each database. Guarded by {@link #writeLock}. - * An entry is put by {@link #registerPrimeOwner(String, long)} and removed by {@link - * #unregisterPrimeOwner(String, long)}, so the map holds exactly one entry per live database - * client that creates multiplexed sessions and nothing for retired clients. + * Registered sources. Reference identity keeps overlapping client instances independent, and + * writers synchronize because the identity scan and mutation must be atomic. */ - private final Map primeOwners = new HashMap<>(); - - /** Guards the writers of {@link #primeSessions}; all of them run on rare paths only. */ - private final Object writeLock = new Object(); - - /** A multiplexed session that is registered for priming. */ - static final class PrimeSession { - private final String databaseName; - private final String sessionName; - private final long ownerTicket; - private final long generation; - - private PrimeSession( - String databaseName, String sessionName, long ownerTicket, long generation) { - this.databaseName = databaseName; - this.sessionName = sessionName; - this.ownerTicket = ownerTicket; - this.generation = generation; - } - - String getDatabaseName() { - return databaseName; - } - - String getSessionName() { - return sessionName; - } - - /** The owner ticket of the database client that created the session. */ - long getOwnerTicket() { - return ownerTicket; - } - - /** The creation generation; a higher generation means a more recently created session. */ - long getGeneration() { - return generation; - } - - @Override - public boolean equals(Object other) { - if (!(other instanceof PrimeSession)) { - return false; - } - PrimeSession that = (PrimeSession) other; - return generation == that.generation - && ownerTicket == that.ownerTicket - && databaseName.equals(that.databaseName) - && sessionName.equals(that.sessionName); - } + private final CopyOnWriteArrayList primeSessionSources = + new CopyOnWriteArrayList<>(); - @Override - public int hashCode() { - return Objects.hash(databaseName, sessionName, ownerTicket, generation); - } - - @Override - public String toString() { - return sessionName + "#" + generation; - } - } + /** Index from which the next priming attempt starts searching for an available source. */ + private final AtomicInteger nextSourceIndex = new AtomicInteger(); /** * @param metadataProvider provides the fixed headers and the resource-prefix header of a normal @@ -223,27 +127,20 @@ public String toString() { * @param projectName the project resource name that serves as the default resource-prefix value * @param requestIdCreator the request id creator of the owning rpc, so priming RPCs carry a * request id with the same client id as every other call of the rpc - * @param defaultCallCredentials the credentials that GAX attaches to normal calls, or {@code + * @param callCredentialsProvider provides the credentials of a normal Spanner call, or {@code * null} when the client runs without credentials - * @param callCredentialsProvider the optional user-supplied provider that takes precedence over - * {@code defaultCallCredentials} for each call, exactly like for normal Spanner calls - * @param routeToLeader whether normal Spanner calls carry the route-to-leader header * @param rpcDeadline the deadline of a single priming RPC */ DynamicChannelPoolPrimer( SpannerMetadataProvider metadataProvider, String projectName, RequestIdCreator requestIdCreator, - @Nullable CallCredentials defaultCallCredentials, @Nullable CallCredentialsProvider callCredentialsProvider, - boolean routeToLeader, Duration rpcDeadline) { this.metadataProvider = Preconditions.checkNotNull(metadataProvider); this.projectName = Preconditions.checkNotNull(projectName); this.requestIdCreator = Preconditions.checkNotNull(requestIdCreator); - this.defaultCallCredentials = defaultCallCredentials; this.callCredentialsProvider = callCredentialsProvider; - this.routeToLeader = routeToLeader; Preconditions.checkArgument( rpcDeadline != null && !rpcDeadline.isZero() && !rpcDeadline.isNegative(), "rpcDeadline must be positive"); @@ -280,189 +177,93 @@ Duration getRpcDeadline() { return rpcDeadline; } - /** - * Registers the database client with the given owner ticket as the current owner of the given - * database. Must be called before the client issues its first {@code CreateSession}, so every - * session that the client creates can be attributed to it through {@link - * #registerPrimeSession(String, String, long)}. A previous owner of the database, if any, is - * replaced and its registered session, if any, is removed, because the previous client has been - * or is being retired. - * - * @param ownerTicket a ticket that is unique among all database clients of the process - */ - void registerPrimeOwner(String databaseName, long ownerTicket) { - checkDatabaseName(databaseName); - synchronized (writeLock) { - Long previous = primeOwners.put(databaseName, ownerTicket); - if (previous != null && previous != ownerTicket) { - removeEntries(databaseName, previous); - } - } - } - - /** - * Unregisters the database client with the given owner ticket as the owner of the given database, - * and removes the session that it registered, if any. Called when the client is invalidated or - * closed, because its multiplexed session is then no longer maintained. A late {@code - * CreateSession} response of the client is dropped afterwards, because its ticket is no longer - * the current owner ticket of the database. If the database has already been taken over by a - * replacement client, the replacement and its session are left untouched. - */ - void unregisterPrimeOwner(String databaseName, long ownerTicket) { - checkDatabaseName(databaseName); - synchronized (writeLock) { - Long current = primeOwners.get(databaseName); - if (current != null && current == ownerTicket) { - primeOwners.remove(databaseName); - } - removeEntries(databaseName, ownerTicket); - } - } - - /** - * Registers a multiplexed session that was just created successfully for the given database as a - * prime session. The entry replaces any previous entry of the same database, for example when a - * database client refreshes its multiplexed session, and receives the highest generation, so it - * becomes the session that the next priming attempt uses. - * - *

    The registration is dropped unless the given owner ticket, which the {@code CreateSession} - * request carried, is the current owner ticket of the database: the session then belongs to a - * database client that has been invalidated or closed while its {@code CreateSession} was in - * flight, and must not replace the session of the client that succeeded it. - * - * @param ownerTicket the owner ticket that the {@code CreateSession} request carried - * @return whether the session was registered - */ - boolean registerPrimeSession(String databaseName, String sessionName, long ownerTicket) { - checkDatabaseName(databaseName); - Preconditions.checkArgument( - sessionName != null && !sessionName.isEmpty(), "sessionName must not be empty"); - synchronized (writeLock) { - Long current = primeOwners.get(databaseName); - if (current == null || current != ownerTicket) { - return false; - } - PrimeSession entry = - new PrimeSession(databaseName, sessionName, ownerTicket, generations.incrementAndGet()); - ImmutableList.Builder builder = ImmutableList.builder(); - // The new entry has the highest generation, so it goes first to keep the descending order. - builder.add(entry); - for (PrimeSession existing : primeSessions) { - if (!existing.databaseName.equals(databaseName)) { - builder.add(existing); + /** Registers a session source by reference identity. Repeated registration is a no-op. */ + void registerPrimeSessionSource(ChannelPrimeSessionSource source) { + Preconditions.checkNotNull(source); + synchronized (primeSessionSources) { + for (ChannelPrimeSessionSource existing : primeSessionSources) { + if (existing == source) { + return; } } - primeSessions = builder.build(); - return true; - } - } - - /** Removes the entries of the given database that were registered with the given ticket. */ - private void removeEntries(String databaseName, long ownerTicket) { - ImmutableList.Builder builder = ImmutableList.builder(); - boolean removed = false; - for (PrimeSession existing : primeSessions) { - if (existing.databaseName.equals(databaseName) && existing.ownerTicket == ownerTicket) { - removed = true; - } else { - builder.add(existing); - } + primeSessionSources.add(source); } - if (removed) { - primeSessions = builder.build(); - } - } - - private static void checkDatabaseName(String databaseName) { - Preconditions.checkArgument( - databaseName != null && !databaseName.isEmpty(), "databaseName must not be empty"); } - /** - * Removes exactly the given entry, and returns whether it was still registered. An entry that has - * already been replaced by a newer session of the same database, or that has already been - * removed, is left alone, so an eviction never clobbers a newer registration. - */ - @VisibleForTesting - boolean evictPrimeSession(PrimeSession entry) { - synchronized (writeLock) { - if (!primeSessions.contains(entry)) { - return false; - } - ImmutableList.Builder builder = ImmutableList.builder(); - for (PrimeSession existing : primeSessions) { - if (!existing.equals(entry)) { - builder.add(existing); + /** Deregisters a session source. Repeated deregistration is a no-op. */ + void unregisterPrimeSessionSource(ChannelPrimeSessionSource source) { + Preconditions.checkNotNull(source); + synchronized (primeSessionSources) { + for (int i = 0; i < primeSessionSources.size(); i++) { + if (primeSessionSources.get(i) == source) { + primeSessionSources.remove(i); + return; } } - primeSessions = builder.build(); - return true; } } - /** - * Returns the entry with the highest generation, that is the most recently created registered - * multiplexed session, or {@code null} if none is registered. - */ + /** Returns a snapshot of registered sources in registration order. */ @VisibleForTesting - @Nullable - PrimeSession getPrimeSession() { - ImmutableList current = primeSessions; - return current.isEmpty() ? null : current.get(0); + List getPrimeSessionSources() { + return ImmutableList.copyOf(primeSessionSources); } - /** Returns the name of the session that the next priming attempt uses, or {@code null}. */ + /** Returns a currently available session name without blocking, rotating the starting source. */ @VisibleForTesting @Nullable String getPrimeSessionName() { - PrimeSession entry = getPrimeSession(); - return entry == null ? null : entry.sessionName; - } - - /** Returns all registered entries in descending generation order. */ - @VisibleForTesting - ImmutableList getPrimeSessions() { - return primeSessions; - } - - /** Returns a copy of the current owner ticket of every database with a live database client. */ - @VisibleForTesting - Map getPrimeOwners() { - synchronized (writeLock) { - return new HashMap<>(primeOwners); + // The primer neither classifies failures nor evicts sessions. A dead session remains until its + // owner observes it and isValid() stops offering it. Client traffic normally reaches + // MultiplexedSessionTransaction.onError quickly; periodic refresh backs up an idle client. + // Attempts may thus use a dead session, and the shared cursor cannot ensure consecutive + // concurrent attempts on one channel use different sources. Deciding whether a session is still + // usable deliberately stays with the owning client, which already tracks it for its own + // traffic. + List sources = ImmutableList.copyOf(primeSessionSources); + int size = sources.size(); + if (size == 0) { + return null; } + int start = Math.floorMod(nextSourceIndex.getAndIncrement(), size); + for (int offset = 0; offset < size; offset++) { + ChannelPrimeSessionSource source = sources.get((start + offset) % size); + String sessionName = source.getChannelPrimeSessionName(); + if (sessionName != null) { + return sessionName; + } + } + return null; } @Override public ListenableFuture prime(ManagedChannel channel) { - PrimeSession entry = getPrimeSession(); - if (entry == null) { - // The Go client refuses to scale up at all before a multiplexed session exists. The Java - // primer cannot gate the pool's scale-up decision, so the attempt fails fast instead and the - // pool's retry with backoff and its close-on-failure behaviour handle it. The wasted dial - // and the prime-failure metrics of such an attempt are the known divergence from Go. + String sessionName = getPrimeSessionName(); + if (sessionName == null) { + // The primer cannot gate the pool's scale-up decision, so the attempt fails fast. The pool's + // retry with backoff and its close-on-failure behaviour handle the unavailable session. return Futures.immediateFailedFuture( SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Cannot prime a dynamic channel pool channel before a multiplexed session is" + " available")); } - return executePrimeStatement(channel, entry); + return executePrimeStatement(channel, sessionName); } - private ListenableFuture executePrimeStatement(ManagedChannel channel, PrimeSession entry) { - String sessionName = entry.sessionName; + private ListenableFuture executePrimeStatement(ManagedChannel channel, String sessionName) { ExecuteSqlRequest request = ExecuteSqlRequest.newBuilder().setSession(sessionName).setSql(PRIME_SQL).build(); - // The priming query always uses the unary ExecuteSql method with an explicit short deadline, - // as the Go client does. ExecuteStreamingSql must never be used here: its configured default - // deadline is one hour, which would let a hung priming call outlive the pool's prime timeout - // by far. + // ExecuteStreamingSql must never be used here: its configured default deadline is one hour, + // which would let a hung priming call outlive the pool's prime timeout by far. The unary + // ExecuteSql method therefore uses an explicit short deadline. CallOptions callOptions = CallOptions.DEFAULT.withDeadlineAfter(rpcDeadline.toNanos(), TimeUnit.NANOSECONDS); - CallCredentials callCredentials = resolveCallCredentials(); - if (callCredentials != null) { - callOptions = callOptions.withCallCredentials(callCredentials); + if (callCredentialsProvider != null) { + CallCredentials callCredentials = callCredentialsProvider.getCallCredentials(); + if (callCredentials != null) { + callOptions = callOptions.withCallCredentials(callCredentials); + } } Channel channelWithHeaders = ClientInterceptors.intercept( @@ -470,82 +271,21 @@ private ListenableFuture executePrimeStatement(ManagedChannel channel, Pri ClientCall call = channelWithHeaders.newCall(SpannerGrpc.getExecuteSqlMethod(), callOptions); // Cancelling the future returned by futureUnaryCall cancels the underlying ClientCall, and - // cancelling the derived futures cancels their input. A channel that is shut down under the + // cancelling the derived future cancels its input. A channel that is shut down under the // running call fails the call with UNAVAILABLE or CANCELLED, which fails this future. - ListenableFuture result = - Futures.transform( - ClientCalls.futureUnaryCall(call, request), - resultSet -> null, - MoreExecutors.directExecutor()); - return Futures.catchingAsync( - result, - Throwable.class, - failure -> { - if (isCandidateSpecificFailure(failure)) { - // The failure is specific to the session that was used, for example because its - // database was dropped or the caller lost access to it. Evict exactly the entry that - // was used, so the pool's next attempt falls back to the next most recent session. A - // newer session of the same database that was registered in the meantime is never - // clobbered. - evictPrimeSession(entry); - } - return Futures.immediateFailedFuture(failure); - }, + return Futures.transform( + ClientCalls.futureUnaryCall(call, request), + resultSet -> null, MoreExecutors.directExecutor()); } - /** - * Returns whether a priming failure is terminal for the session that was used, so priming should - * rotate to the next candidate: the session or its database is not found, permission to use it is - * denied, or its session or database role is reported as invalid. Only such failures evict the - * session; transient failures such as UNAVAILABLE or DEADLINE_EXCEEDED say nothing about the - * session and keep it registered. - */ - @VisibleForTesting - static boolean isCandidateSpecificFailure(Throwable failure) { - Status status = Status.fromThrowable(failure); - switch (status.getCode()) { - case NOT_FOUND: - case PERMISSION_DENIED: - return true; - case FAILED_PRECONDITION: - case INVALID_ARGUMENT: - break; - default: - return false; - } - String description = status.getDescription(); - if (description == null) { - return false; - } - String message = description.toLowerCase(Locale.ENGLISH); - if (ROLE_PATTERN.matcher(message).find()) { - return true; - } - return message.contains("session") - && (message.contains("not found") - || message.contains("invalid") - || message.contains("does not exist")); - } - - @Nullable - private CallCredentials resolveCallCredentials() { - if (callCredentialsProvider != null) { - CallCredentials callCredentials = callCredentialsProvider.getCallCredentials(); - if (callCredentials != null) { - return callCredentials; - } - } - return defaultCallCredentials; - } - /** * Returns the per-call headers of a priming RPC. The delegate channel that the pool hands to the * primer is built by the same channel provider as every other channel of the pool, so it already * carries the fixed headers of the client (x-goog-api-client, the user agent, and any custom - * headers) through the GAX header interceptor. Only the headers that a normal call adds per call - * are attached here: the resource-prefix header, x-goog-request-params, the route-to-leader - * header when leader-aware routing is enabled, and x-goog-spanner-request-id. + * headers) through the GAX header interceptor. Only the headers that a normal single-use + * read-only query adds per call are attached here: the resource-prefix header, + * x-goog-request-params, and x-goog-spanner-request-id. */ @VisibleForTesting Metadata newHeaders(String sessionName) { @@ -560,10 +300,8 @@ Metadata newHeaders(String sessionName) { headers.put(key, value); } } - headers.put(REQUEST_PARAMS_KEY, "session=" + urlEncode(sessionName)); - if (routeToLeader) { - headers.put(ROUTE_TO_LEADER_KEY, "true"); - } + headers.put( + SpannerMetadataProvider.REQUEST_PARAMS_HEADER_KEY, "session=" + urlEncode(sessionName)); // The pool invokes prime() once per attempt, so every attempt gets a fresh request id with // attempt number 1. The header is written directly and the request id is deliberately not set // as a call option: the RequestIdInterceptor on the delegate channel only acts on the call diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index a48f9fef3aeb..05a1e0baa3f1 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -91,6 +91,7 @@ import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub; import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings; import com.google.cloud.spanner.encryption.EncryptionConfigProtoMapper; +import com.google.cloud.spanner.spi.v1.SpannerRpc.ChannelPrimeSessionSource; import com.google.cloud.spanner.v1.stub.SpannerStub; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.annotations.VisibleForTesting; @@ -805,12 +806,8 @@ private InstantiatingGrpcChannelProvider.Builder createBaseChannelProviderBuilde * of priming RPCs come from the request id creator of this rpc. The deadline of the priming RPC * is derived from the pool's prime timeout and normally stays below it; see {@link * DynamicChannelPoolPrimer#rpcDeadlineFor(Duration)} for the minimum deadline that a prime - * timeout of two milliseconds or less yields. The sessions that the primer uses are registered by - * {@link #createSession(String, String, Map, Map, boolean)} whenever a multiplexed session is - * created successfully for the current owner of its database, which the database clients register - * through {@link #registerChannelPrimeOwner(String, long)} and remove through {@link - * #unregisterChannelPrimeOwner(String, long)}, so a session whose creation was still in flight - * when its database client was retired is never registered. + * timeout of two milliseconds or less yields. Live multiplexed-session database clients register + * themselves as session sources and unregister themselves when closed. */ @Nullable private DynamicChannelPoolPrimer createChannelPrimer( @@ -818,27 +815,37 @@ private DynamicChannelPoolPrimer createChannelPrimer( if (!options.isGrpcGcpExtensionEnabled() || !options.isDynamicChannelPoolEnabled()) { return null; } - CallCredentials defaultCallCredentials = null; - try { - Credentials credentials = credentialsProvider.getCredentials(); - if (credentials != null) { - defaultCallCredentials = MoreCallCredentials.from(credentials); - } - } catch (IOException e) { - throw newSpannerException(e); - } return new DynamicChannelPoolPrimer( metadataProvider, projectName, requestIdCreator, - defaultCallCredentials, - callCredentialsProvider, - leaderAwareRoutingEnabled, + createChannelPrimeCallCredentialsProvider(credentialsProvider, callCredentialsProvider), // The options already contain the merged prime timeout, including a user-provided one. DynamicChannelPoolPrimer.rpcDeadlineFor( options.getGcpChannelPoolOptions().getChannelPrimeTimeout())); } + @VisibleForTesting + @Nullable + static CallCredentialsProvider createChannelPrimeCallCredentialsProvider( + CredentialsProvider credentialsProvider, + @Nullable CallCredentialsProvider callCredentialsProvider) { + final CallCredentials defaultCallCredentials; + try { + Credentials credentials = credentialsProvider.getCredentials(); + defaultCallCredentials = credentials == null ? null : MoreCallCredentials.from(credentials); + } catch (IOException e) { + throw newSpannerException(e); + } + if (callCredentialsProvider == null) { + return defaultCallCredentials == null ? null : () -> defaultCallCredentials; + } + return () -> { + CallCredentials callCredentials = callCredentialsProvider.getCallCredentials(); + return callCredentials != null ? callCredentials : defaultCallCredentials; + }; + } + /** Returns the primer of scaled-up dynamic channel pool channels, or {@code null} if none. */ @VisibleForTesting @Nullable @@ -2052,38 +2059,20 @@ public Session createSession( CreateSessionRequest request = requestBuilder.build(); GrpcCallContext context = newCallContext(options, databaseName, request, SpannerGrpc.getCreateSessionMethod(), true); - Session session = get(spannerStub.createSessionCallable().futureCall(request, context)); - // Every multiplexed session, including the periodic refresh of an existing one, is created - // through this method, so the most recently created multiplexed session is normally a valid - // session for priming channels that the dynamic channel pool adds during scale-up. The - // registration is keyed on both the request flag and the response flag: a backend that does - // not echo the multiplexed flag, for example because it ignores the request flag and creates a - // regular session, must never turn that session into a prime session. Regular sessions are - // not safe for the concurrent use that priming implies. The owner ticket of the database - // client that issued the request is passed on, so the primer can drop the session if the - // client has been retired while the request was in flight; a request without an owner ticket - // never registers a prime session. - Long primeOwnerTicket = Option.CHANNEL_PRIME_OWNER.getLong(options); - if (channelPrimer != null - && isMultiplexed - && session.getMultiplexed() - && primeOwnerTicket != null) { - channelPrimer.registerPrimeSession(databaseName, session.getName(), primeOwnerTicket); - } - return session; - } - - @Override - public void registerChannelPrimeOwner(String databaseName, long ownerTicket) { + return get(spannerStub.createSessionCallable().futureCall(request, context)); + } + + @Override + public void registerChannelPrimeSessionSource(ChannelPrimeSessionSource source) { if (channelPrimer != null) { - channelPrimer.registerPrimeOwner(databaseName, ownerTicket); + channelPrimer.registerPrimeSessionSource(source); } } @Override - public void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) { + public void unregisterChannelPrimeSessionSource(ChannelPrimeSessionSource source) { if (channelPrimer != null) { - channelPrimer.unregisterPrimeOwner(databaseName, ownerTicket); + channelPrimer.unregisterPrimeSessionSource(source); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java index e9c748472754..7a27fd14da56 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java @@ -32,6 +32,9 @@ /** For internal use only. */ class SpannerMetadataProvider { + static final Metadata.Key REQUEST_PARAMS_HEADER_KEY = + Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); + private final Cache>> extraHeadersCache = CacheBuilder.newBuilder().maximumSize(100).build(); private final Map, String> headers; diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 4e07cfa1d835..defdaf444584 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -82,14 +82,7 @@ public interface SpannerRpc extends ServiceRpc { /** Options passed in {@link SpannerRpc} methods to control how an RPC is issued. */ enum Option { CHANNEL_HINT("Channel Hint"), - CHANNEL_ID_AFFINITY("Channel ID Affinity"), - /** - * The owner ticket of the database client that issues a multiplexed {@code CreateSession}. The - * created session is used for priming channels that the dynamic channel pool adds during - * scale-up only while the ticket is the current one that was registered through {@link - * SpannerRpc#registerChannelPrimeOwner(String, long)} for the database. Internal. - */ - CHANNEL_PRIME_OWNER("Channel Prime Owner"); + CHANNEL_ID_AFFINITY("Channel ID Affinity"); private final String value; @@ -121,6 +114,13 @@ public String toString() { } } + /** Source of a multiplexed session for priming dynamic channel pool channels. */ + interface ChannelPrimeSessionSource { + /** Returns a session name without blocking, or {@code null} if none is currently available. */ + @Nullable + String getChannelPrimeSessionName(); + } + /** * Represents results from paginated RPCs, i.e., those where up to a maximum number of items is * returned from each call and a followup call must be made to fetch more. @@ -383,23 +383,11 @@ default Session createSession( void deleteSession(String sessionName, @Nullable Map options) throws SpannerException; - /** - * Registers the database client with the given owner ticket as the current owner of the given - * database for priming channels that the dynamic channel pool adds during scale-up. A multiplexed - * session that is created with {@link Option#CHANNEL_PRIME_OWNER} set to the ticket is used for - * priming only while the ticket is the current owner ticket of its database. Called before the - * client issues its first multiplexed {@code CreateSession}. The default implementation does - * nothing. - */ - default void registerChannelPrimeOwner(String databaseName, long ownerTicket) {} + /** Registers a source of a multiplexed session for priming dynamic channel pool channels. */ + default void registerChannelPrimeSessionSource(ChannelPrimeSessionSource source) {} - /** - * Unregisters the database client with the given owner ticket as the owner of the given database - * and removes the multiplexed session that it created from the sessions that are used for priming - * channels that the dynamic channel pool adds during scale-up. Called when the client is - * invalidated or closed. The default implementation does nothing. - */ - default void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) {} + /** Unregisters a source previously passed to {@link #registerChannelPrimeSessionSource}. */ + default void unregisterChannelPrimeSessionSource(ChannelPrimeSessionSource source) {} ApiFuture asyncDeleteSession(String sessionName, @Nullable Map options) throws SpannerException; diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java index 9b48d55d42be..62c38511ee3b 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java @@ -231,6 +231,9 @@ public void refreshInFlightWhenClientIsReplacedNeverBecomesPrimeSession() throws DatabaseClientImpl retired = createClientWithPrimeSession(); String retiredSession = multiplexedClient(retired).getCurrentSessionReference().getName(); invalidate(retired); + // Invalid sources stay registered until replacement, but no longer expose their session. + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(1); + assertThat(rpc.getPrimeSessionNames()).isEmpty(); blockRefreshInFlight(retired); // Getting the client again retires the invalid client and creates its replacement. The @@ -245,7 +248,7 @@ public void refreshInFlightWhenClientIsReplacedNeverBecomesPrimeSession() throws assertThat(replacementSession).isNotEqualTo(retiredSession); // The late refresh of the retired client is dropped, and the replacement stays selected. assertThat(rpc.getPrimeSessionNames()).containsExactly(replacementSession); - assertThat(rpc.getPrimeOwnerDatabases()).containsExactly(DATABASE_ID.getName()); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(1); // The retired client ignored its refreshed session as well. assertEquals(retiredSession, multiplexedClient(retired).getCurrentSessionReference().getName()); assertTrue(replacement.isValid()); @@ -270,7 +273,7 @@ public void refreshInFlightWhenSpannerIsClosedNeverBecomesPrimeSession() throws spanner.close(); assertThat(rpc.getPrimeSessionNames()).isEmpty(); - assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(0); assertEquals(session, multiplexedClient(client).getCurrentSessionReference().getName()); } @@ -293,7 +296,7 @@ public void initialCreateSessionInFlightWhenSpannerIsClosedNeverBecomesPrimeSess spanner.close(); assertThat(rpc.getPrimeSessionNames()).isEmpty(); - assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(0); // The closed client ignored the session that arrived after it was closed. SpannerException exception = assertThrows( @@ -302,10 +305,10 @@ public void initialCreateSessionInFlightWhenSpannerIsClosedNeverBecomesPrimeSess } @Test - public void churnOfDatabaseClientsLeavesNoPrimeOwnersBehind() throws Exception { + public void churnOfDatabaseClientsLeavesNoPrimeSessionSourcesBehind() throws Exception { // A long-lived Spanner instance that churns through many database names must not retain - // anything for the retired database clients: the registry holds exactly one owner and one - // session per live database client. + // anything for retired database clients: the registry holds exactly one source per live + // database client. List liveDatabases = new ArrayList<>(); for (int i = 0; i < 20; i++) { DatabaseId databaseId = DatabaseId.of("[PROJECT]", "[INSTANCE]", "churn-" + i); @@ -322,20 +325,21 @@ public void churnOfDatabaseClientsLeavesNoPrimeOwnersBehind() throws Exception { awaitCondition(() -> rpc.getPrimeSessionNames().contains(replacementSession)); assertThat(rpc.getPrimeSessionNames()).doesNotContain(session); assertThat(rpc.getPrimeSessionNames()).hasSize(liveDatabases.size()); - assertThat(rpc.getPrimeOwnerDatabases()).containsExactlyElementsIn(liveDatabases); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(liveDatabases.size()); } spanner.close(); assertThat(rpc.getPrimeSessionNames()).isEmpty(); - assertThat(rpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(0); } @Test - public void failedClientConstructionLeavesNoPrimeOwnerBehind() throws Exception { + public void failedClientConstructionLeavesNoPrimeSessionSourceBehind() throws Exception { // A client that waits for its first multiplexed session throws from its constructor when the // CreateSession fails, and is never cached by the Spanner instance, so nothing will ever close - // it. It must still have released the owner ticket that it registered before the CreateSession. + // it. It must still have released the session source that it registered before the + // CreateSession. createWaitingSpanner(Duration.ofSeconds(5)); mockSpanner.setCreateSessionExecutionTime( SimulatedExecutionTime.ofStickyException( @@ -346,7 +350,7 @@ public void failedClientConstructionLeavesNoPrimeOwnerBehind() throws Exception SpannerException exception = assertThrows(SpannerException.class, () -> waitingSpanner.getDatabaseClient(databaseId)); assertEquals(ErrorCode.PERMISSION_DENIED, exception.getErrorCode()); - assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(waitingRpc.getPrimeSessionSourceCount()).isEqualTo(0); assertThat(waitingRpc.getPrimeSessionNames()).isEmpty(); } @@ -371,15 +375,16 @@ public void sessionArrivingAfterFailedClientConstructionNeverBecomesPrimeSession SpannerException exception = assertThrows(SpannerException.class, () -> waitingSpanner.getDatabaseClient(databaseId)); assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode()); - // The owner ticket is released as soon as the constructor fails, before the session arrives. - assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + // The session source is unregistered as soon as the constructor fails, before the session + // arrives. + assertThat(waitingRpc.getPrimeSessionSourceCount()).isEqualTo(0); int arrivedBeforeRelease = mockSpanner.arrivedCreateSessions.get(); mockSpanner.releaseCreateSessions(); awaitCondition(() -> waitingRpc.getCompletedMultiplexedCreateSessions() >= 1); assertThat(waitingRpc.getPrimeSessionNames()).isEmpty(); - assertThat(waitingRpc.getPrimeOwnerDatabases()).isEmpty(); + assertThat(waitingRpc.getPrimeSessionSourceCount()).isEqualTo(0); // No maintainer was started for the abandoned client, so no session is ever refreshed. Thread.sleep(200L); assertEquals(arrivedBeforeRelease, mockSpanner.arrivedCreateSessions.get()); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index f9b4e19cb699..dfbe267bfdcd 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -19,15 +19,12 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeFalse; import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -35,7 +32,9 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.cloud.NoCredentials; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionConsumer; @@ -54,7 +53,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; -import javax.annotation.Nullable; import org.junit.After; import org.junit.Test; import org.junit.runner.RunWith; @@ -119,7 +117,7 @@ public void testMaintainer() { return null; }) .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + .asyncCreateMultiplexedSession(any(SessionConsumer.class)); // Create a client. This should get session1. MultiplexedSessionDatabaseClient client = @@ -174,7 +172,7 @@ public void testClosedClientIgnoresInitialSessionThatArrivesAfterClose() { return null; }) .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + .asyncCreateMultiplexedSession(any(SessionConsumer.class)); MultiplexedSessionDatabaseClient client = new MultiplexedSessionDatabaseClient(sessionClient, clock); assertNotNull(consumer.get()); @@ -233,7 +231,7 @@ public void testClosedClientIgnoresRefreshedSessionThatArrivesAfterClose() { return null; }) .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + .asyncCreateMultiplexedSession(any(SessionConsumer.class)); MultiplexedSessionDatabaseClient client = new MultiplexedSessionDatabaseClient(sessionClient, clock); assertEquals(sessionReference1, client.getCurrentSessionReference()); @@ -474,7 +472,7 @@ public void testCloseKeepsChannelUsageEntryWhileAnotherClientIsUsingSameSpanner( } @Test - public void testChannelPrimeOwnerTicketIsRegisteredCarriedAndUnregistered() { + public void testChannelPrimeSessionSourceLifecycleAndNonBlockingAccessor() throws Exception { assumeTrue(isJava8()); Clock clock = mock(Clock.class); when(clock.instant()).thenReturn(Instant.now()); @@ -484,7 +482,6 @@ public void testChannelPrimeOwnerTicketIsRegisteredCarriedAndUnregistered() { SpannerOptions spannerOptions = mock(SpannerOptions.class); SessionPoolOptions sessionPoolOptions = mock(SessionPoolOptions.class); when(sessionClient.getSpanner()).thenReturn(spanner); - when(sessionClient.getDatabaseId()).thenReturn(TEST_DATABASE_ID); when(spanner.getRpc()).thenReturn(rpc); when(spanner.getOptions()).thenReturn(spannerOptions); when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); @@ -492,60 +489,71 @@ public void testChannelPrimeOwnerTicketIsRegisteredCarriedAndUnregistered() { .thenReturn(Duration.ofDays(7)); when(sessionPoolOptions.getMultiplexedSessionMaintenanceLoopFrequency()) .thenReturn(Duration.ofMinutes(10)); - SessionImpl session = mock(SessionImpl.class); - when(session.getSessionReference()).thenReturn(mock(SessionReference.class)); - // Capture the options of every CreateSession without delivering a session yet. - List> createSessionOptions = new ArrayList<>(); - AtomicReference consumer = new AtomicReference<>(); + List consumers = new ArrayList<>(); doAnswer( (Answer) invocationOnMock -> { - consumer.set(invocationOnMock.getArgument(0)); - createSessionOptions.add(invocationOnMock.getArgument(1)); + consumers.add(invocationOnMock.getArgument(0)); return null; }) .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class), any()); + .asyncCreateMultiplexedSession(any(SessionConsumer.class)); MultiplexedSessionDatabaseClient first = new MultiplexedSessionDatabaseClient(sessionClient, clock); MultiplexedSessionDatabaseClient second = new MultiplexedSessionDatabaseClient(sessionClient, clock); - // Every client registers itself as an owner with its own ticket before its first - // CreateSession, and that CreateSession carries the ticket. - ArgumentCaptor tickets = ArgumentCaptor.forClass(Long.class); - verify(rpc, times(2)) - .registerChannelPrimeOwner(eq(TEST_DATABASE_ID.getName()), tickets.capture()); - long firstTicket = tickets.getAllValues().get(0); - long secondTicket = tickets.getAllValues().get(1); - assertNotEquals(firstTicket, secondTicket); - assertEquals(2, createSessionOptions.size()); - assertEquals( - Long.valueOf(firstTicket), - SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(0))); - assertEquals( - Long.valueOf(secondTicket), - SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(1))); - verify(rpc, never()).unregisterChannelPrimeOwner(any(), anyLong()); - - // A refresh of the session carries the same ticket as the initial CreateSession. - consumer.get().onSessionReady(session); + ArgumentCaptor sources = + ArgumentCaptor.forClass(SpannerRpc.ChannelPrimeSessionSource.class); + verify(rpc, times(2)).registerChannelPrimeSessionSource(sources.capture()); + assertThat(sources.getAllValues()).containsExactly(first, second).inOrder(); + // Pending futures yield no session immediately; the accessor never waits for completion. + assertThat(first.getChannelPrimeSessionName()).isNull(); + assertThat(second.getChannelPrimeSessionName()).isNull(); + + consumers + .get(0) + .onSessionCreateFailure( + SpannerExceptionFactory.newSpannerException(ErrorCode.PERMISSION_DENIED, "denied"), 1); + assertThat(first.getChannelPrimeSessionName()).isNull(); + + Field referenceField = + MultiplexedSessionDatabaseClient.class.getDeclaredField("multiplexedSessionReference"); + referenceField.setAccessible(true); + @SuppressWarnings("unchecked") + AtomicReference> reference = + (AtomicReference>) referenceField.get(first); + SettableApiFuture cancelled = SettableApiFuture.create(); + cancelled.cancel(false); + reference.set(cancelled); + assertThat(first.getChannelPrimeSessionName()).isNull(); + + SessionReference initialReference = mock(SessionReference.class); + when(initialReference.getName()).thenReturn("initial-session"); + SessionImpl initialSession = mock(SessionImpl.class); + when(initialSession.getSessionReference()).thenReturn(initialReference); + consumers.get(1).onSessionReady(initialSession); + assertThat(second.getChannelPrimeSessionName()).isEqualTo("initial-session"); + verify(rpc, times(2)).registerChannelPrimeSessionSource(any()); + when(clock.instant()).thenReturn(Instant.now().plus(Duration.ofDays(8))); second.getMaintainer().maintain(); - assertEquals(3, createSessionOptions.size()); - assertEquals( - Long.valueOf(secondTicket), - SpannerRpc.Option.CHANNEL_PRIME_OWNER.getLong(createSessionOptions.get(2))); + SessionReference refreshedReference = mock(SessionReference.class); + when(refreshedReference.getName()).thenReturn("refreshed-session"); + SessionImpl refreshedSession = mock(SessionImpl.class); + when(refreshedSession.getSessionReference()).thenReturn(refreshedReference); + consumers.get(2).onSessionReady(refreshedSession); + assertThat(second.getChannelPrimeSessionName()).isEqualTo("refreshed-session"); + verify(rpc, times(2)).registerChannelPrimeSessionSource(any()); - // Closing a client unregisters exactly its own ticket, once. first.close(); - verify(rpc).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), firstTicket); - verify(rpc, never()).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), secondTicket); first.close(); - verify(rpc, times(1)).unregisterChannelPrimeOwner(any(), anyLong()); + verify(rpc).unregisterChannelPrimeSessionSource(first); + verify(rpc, never()).unregisterChannelPrimeSessionSource(second); second.close(); - verify(rpc).unregisterChannelPrimeOwner(TEST_DATABASE_ID.getName(), secondTicket); + verify(rpc).unregisterChannelPrimeSessionSource(second); + assertThat(second.getChannelPrimeSessionName()).isNull(); } private SessionClient createSessionClient(SpannerImpl spanner) { @@ -611,8 +619,7 @@ private FailingMultiplexedSessionClient(SpannerImpl spanner) { } @Override - void asyncCreateMultiplexedSession( - SessionConsumer consumer, @Nullable Map options) { + void asyncCreateMultiplexedSession(SessionConsumer consumer) { consumer.onSessionCreateFailure( SpannerExceptionFactory.newSpannerException(ErrorCode.UNAUTHENTICATED, "test"), 1); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java index cc919756b581..7438ac6b1aba 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java @@ -21,9 +21,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -134,7 +132,7 @@ public void getDbclientAgainGivesSame() { } @Test - public void invalidatedDatabaseClientUnregistersItsChannelPrimeOwnerTicket() { + public void invalidatedDatabaseClientUnregistersItsChannelPrimeSessionSource() { DatabaseId db = DatabaseId.of("projects/p1/instances/i1/databases/d1"); Mockito.when(spannerOptions.getTransportOptions()) .thenReturn(GrpcTransportOptions.newBuilder().build()); @@ -159,29 +157,28 @@ boolean isValid() { }; try { DatabaseClient first = spanner.getDatabaseClient(db); - ArgumentCaptor tickets = ArgumentCaptor.forClass(Long.class); - verify(rpc).registerChannelPrimeOwner(eq(db.getName()), tickets.capture()); - long firstTicket = tickets.getValue(); - verify(rpc, never()).unregisterChannelPrimeOwner(anyString(), anyLong()); + ArgumentCaptor sources = + ArgumentCaptor.forClass(SpannerRpc.ChannelPrimeSessionSource.class); + verify(rpc).registerChannelPrimeSessionSource(sources.capture()); + SpannerRpc.ChannelPrimeSessionSource firstSource = sources.getValue(); + verify(rpc, never()).unregisterChannelPrimeSessionSource(firstSource); - // The invalidated client is replaced. Closing it unregisters its ticket, and the - // replacement registers a new ticket of its own. valid.set(false); DatabaseClient second = spanner.getDatabaseClient(db); assertThat(second).isNotSameInstanceAs(first); - verify(rpc).unregisterChannelPrimeOwner(db.getName(), firstTicket); - verify(rpc, times(2)).registerChannelPrimeOwner(eq(db.getName()), tickets.capture()); - long secondTicket = tickets.getValue(); - assertThat(secondTicket).isNotEqualTo(firstTicket); - verify(rpc, never()).unregisterChannelPrimeOwner(db.getName(), secondTicket); + verify(rpc).unregisterChannelPrimeSessionSource(firstSource); + verify(rpc, times(2)).registerChannelPrimeSessionSource(sources.capture()); + SpannerRpc.ChannelPrimeSessionSource secondSource = sources.getValue(); + assertThat(secondSource).isNotSameInstanceAs(firstSource); + verify(rpc, never()).unregisterChannelPrimeSessionSource(secondSource); } finally { spanner.close(); } } @Test - public void closeUnregistersChannelPrimeOwnerTicketOfEveryDatabaseClient() { + public void closeUnregistersChannelPrimeSessionSourceOfEveryDatabaseClient() { DatabaseId db1 = DatabaseId.of("projects/p1/instances/i1/databases/d1"); DatabaseId db2 = DatabaseId.of("projects/p1/instances/i1/databases/d2"); Mockito.when(spannerOptions.getTransportOptions()) @@ -191,17 +188,16 @@ public void closeUnregistersChannelPrimeOwnerTicketOfEveryDatabaseClient() { SpannerImpl spanner = new SpannerImpl(rpc, spannerOptions); spanner.getDatabaseClient(db1); spanner.getDatabaseClient(db2); - ArgumentCaptor ticket1 = ArgumentCaptor.forClass(Long.class); - ArgumentCaptor ticket2 = ArgumentCaptor.forClass(Long.class); - verify(rpc).registerChannelPrimeOwner(eq(db1.getName()), ticket1.capture()); - verify(rpc).registerChannelPrimeOwner(eq(db2.getName()), ticket2.capture()); - verify(rpc, never()).unregisterChannelPrimeOwner(anyString(), anyLong()); + ArgumentCaptor sources = + ArgumentCaptor.forClass(SpannerRpc.ChannelPrimeSessionSource.class); + verify(rpc, times(2)).registerChannelPrimeSessionSource(sources.capture()); + verify(rpc, never()).unregisterChannelPrimeSessionSource(any()); spanner.close(); - verify(rpc).unregisterChannelPrimeOwner(db1.getName(), ticket1.getValue()); - verify(rpc).unregisterChannelPrimeOwner(db2.getName(), ticket2.getValue()); - verify(rpc, times(2)).unregisterChannelPrimeOwner(anyString(), anyLong()); + verify(rpc).unregisterChannelPrimeSessionSource(sources.getAllValues().get(0)); + verify(rpc).unregisterChannelPrimeSessionSource(sources.getAllValues().get(1)); + verify(rpc, times(2)).unregisterChannelPrimeSessionSource(any()); } @Test diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index 19c03859e669..547f6b70a22d 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -241,7 +241,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> @@ -324,7 +324,7 @@ public void inlineBegin() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index c0bd9d1f3169..1dd2418aa05a 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -232,7 +232,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java index 71bd66b11756..7b46950746e8 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/ChannelPrimerTestRpc.java @@ -17,20 +17,17 @@ package com.google.cloud.spanner.spi.v1; import com.google.cloud.spanner.SpannerOptions; -import com.google.cloud.spanner.spi.v1.DynamicChannelPoolPrimer.PrimeSession; +import com.google.cloud.spanner.spi.v1.SpannerRpc.ChannelPrimeSessionSource; import com.google.spanner.v1.Session; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; import javax.annotation.Nullable; /** - * A {@link GapicSpannerRpc} for tests outside this package that need to observe the prime sessions - * and prime owners of the dynamic channel pool primer and to hook into the retirement of a database - * client. + * A {@link GapicSpannerRpc} for tests outside this package that need to observe prime-session + * sources and hook into the retirement of a database client. */ public class ChannelPrimerTestRpc extends GapicSpannerRpc { private final AtomicInteger completedMultiplexedCreateSessions = new AtomicInteger(); @@ -40,7 +37,7 @@ public ChannelPrimerTestRpc(SpannerOptions options) { super(options); } - /** Sets a hook that runs right after a database client has been unregistered as prime owner. */ + /** Sets a hook that runs right after a database client's session source is unregistered. */ public void setAfterUnregisterHook(Runnable hook) { this.afterUnregisterHook = hook; } @@ -57,13 +54,16 @@ public String getPrimeSessionName() { return primer == null ? null : primer.getPrimeSessionName(); } - /** Returns the names of all registered prime sessions, most recently created first. */ + /** Returns currently available session names in source-preference order. */ public List getPrimeSessionNames() { List names = new ArrayList<>(); DynamicChannelPoolPrimer primer = getChannelPrimer(); if (primer != null) { - for (PrimeSession entry : primer.getPrimeSessions()) { - names.add(entry.getSessionName()); + for (ChannelPrimeSessionSource source : primer.getPrimeSessionSources()) { + String sessionName = source.getChannelPrimeSessionName(); + if (sessionName != null) { + names.add(sessionName); + } } } return names; @@ -85,15 +85,15 @@ public Session createSession( } } - /** Returns the names of the databases that currently have a registered prime owner. */ - public Set getPrimeOwnerDatabases() { + /** Returns the number of registered prime-session sources. */ + public int getPrimeSessionSourceCount() { DynamicChannelPoolPrimer primer = getChannelPrimer(); - return primer == null ? Collections.emptySet() : primer.getPrimeOwners().keySet(); + return primer == null ? 0 : primer.getPrimeSessionSources().size(); } @Override - public void unregisterChannelPrimeOwner(String databaseName, long ownerTicket) { - super.unregisterChannelPrimeOwner(databaseName, ownerTicket); + public void unregisterChannelPrimeSessionSource(ChannelPrimeSessionSource source) { + super.unregisterChannelPrimeSessionSource(source); afterUnregisterHook.run(); } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java index 833414352150..bf2547fdbb38 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java @@ -25,7 +25,7 @@ import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; import com.google.cloud.spanner.XGoogSpannerRequestId; -import com.google.cloud.spanner.spi.v1.DynamicChannelPoolPrimer.PrimeSession; +import com.google.cloud.spanner.spi.v1.SpannerRpc.ChannelPrimeSessionSource; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.ListenableFuture; @@ -58,13 +58,10 @@ import java.util.ArrayList; import java.util.Date; import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; import org.junit.After; import org.junit.Before; @@ -109,9 +106,6 @@ private static final class PrimeService extends SpannerGrpc.SpannerImplBase { volatile CountDownLatch callCancelled = new CountDownLatch(1); @Nullable volatile Status failWith; - /** Failures for specific sessions, which take precedence over {@link #failWith}. */ - final Map failSessionsWith = new ConcurrentHashMap<>(); - @Override public void executeSql(ExecuteSqlRequest request, StreamObserver responseObserver) { requests.add(request); @@ -122,7 +116,7 @@ public void executeSql(ExecuteSqlRequest request, StreamObserver resp if (holdResponses) { return; } - Status failure = failSessionsWith.getOrDefault(request.getSession(), failWith); + Status failure = failWith; if (failure != null) { responseObserver.onError(failure.asRuntimeException()); return; @@ -209,15 +203,16 @@ public void start(Listener responseListener, Metadata headers) { } private static CallCredentials credentialsWithToken(String token) { - return MoreCallCredentials.from( - OAuth2Credentials.create( - new AccessToken(token, new Date(System.currentTimeMillis() + 3_600_000L)))); + return MoreCallCredentials.from(oauthCredentialsWithToken(token)); + } + + private static OAuth2Credentials oauthCredentialsWithToken(String token) { + return OAuth2Credentials.create( + new AccessToken(token, new Date(System.currentTimeMillis() + 3_600_000L))); } private DynamicChannelPoolPrimer newPrimer( - @Nullable CallCredentialsProvider callCredentialsProvider, - boolean routeToLeader, - Duration rpcDeadline) { + @Nullable CallCredentialsProvider callCredentialsProvider, Duration rpcDeadline) { // The fixed headers of the client are configured on the channel and must not be sent by the // primer: it only sends the headers that a normal call adds per call. SpannerMetadataProvider metadataProvider = @@ -225,35 +220,48 @@ private DynamicChannelPoolPrimer newPrimer( ImmutableMap.of("x-goog-api-client", "test-client", "user-agent", "test-agent"), RESOURCE_HEADER_KEY); return new DynamicChannelPoolPrimer( - metadataProvider, - PROJECT_NAME, - requestIdCreator, - credentialsWithToken(DEFAULT_TOKEN), - callCredentialsProvider, - routeToLeader, - rpcDeadline); + metadataProvider, PROJECT_NAME, requestIdCreator, callCredentialsProvider, rpcDeadline); } private DynamicChannelPoolPrimer newPrimer() { - return newPrimer( - /* callCredentialsProvider= */ null, /* routeToLeader= */ true, Duration.ofSeconds(5)); + return newPrimer(() -> credentialsWithToken(DEFAULT_TOKEN), Duration.ofSeconds(5)); } - private static final AtomicLong OWNER_TICKETS = new AtomicLong(); + private static class MutableSessionSource implements ChannelPrimeSessionSource { + @Nullable volatile String sessionName; - /** - * Registers the session under the current owner of the database, registering a fresh owner first - * if the database has none, like a database client that creates or refreshes its session. - */ - private static long registerPrimeSession( - DynamicChannelPoolPrimer primer, String databaseName, String sessionName) { - Long ownerTicket = primer.getPrimeOwners().get(databaseName); - if (ownerTicket == null) { - ownerTicket = OWNER_TICKETS.incrementAndGet(); - primer.registerPrimeOwner(databaseName, ownerTicket); + private MutableSessionSource(@Nullable String sessionName) { + this.sessionName = sessionName; } - assertThat(primer.registerPrimeSession(databaseName, sessionName, ownerTicket)).isTrue(); - return ownerTicket; + + @Override + @Nullable + public String getChannelPrimeSessionName() { + return sessionName; + } + } + + private static final class EqualSessionSource extends MutableSessionSource { + private EqualSessionSource(String sessionName) { + super(sessionName); + } + + @Override + public boolean equals(Object other) { + return other instanceof EqualSessionSource; + } + + @Override + public int hashCode() { + return 1; + } + } + + private static MutableSessionSource registerPrimeSession( + DynamicChannelPoolPrimer primer, String databaseName, String sessionName) { + MutableSessionSource source = new MutableSessionSource(sessionName); + primer.registerPrimeSessionSource(source); + return source; } private static T getWithin(ListenableFuture future, Duration timeout) throws Exception { @@ -279,6 +287,7 @@ private void assertPrimeHeaders(Metadata headers, String expectedToken) throws E assertThat(headers.getAll(RESOURCE_PREFIX_KEY)).containsExactly(DATABASE_NAME); assertThat(headers.getAll(REQUEST_PARAMS_KEY)) .containsExactly("session=" + urlEncode(SESSION_NAME)); + assertThat(headers.containsKey(ROUTE_TO_LEADER_KEY)).isFalse(); XGoogSpannerRequestId requestId = requestIdOf(headers); assertThat(requestId.getHeaderValue()) .isEqualTo( @@ -320,7 +329,6 @@ public void primeExecutesSelectOneWithSessionAndHeaders() throws Exception { assertThat(service.headers).hasSize(1); Metadata headers = service.headers.get(0); assertPrimeHeaders(headers, DEFAULT_TOKEN); - assertThat(headers.getAll(ROUTE_TO_LEADER_KEY)).containsExactly("true"); // The fixed headers of the client are carried by the delegate channel itself and must not be // sent by the primer, or the delegate would send them twice. assertThat(headers.containsKey(API_CLIENT_KEY)).isFalse(); @@ -336,7 +344,6 @@ public void primeSendsEveryHeaderOnceOnDelegateThatCarriesFixedHeaders() throws Metadata headers = service.headers.get(0); assertPrimeHeaders(headers, DEFAULT_TOKEN); assertThat(headers.getAll(API_CLIENT_KEY)).containsExactly("test-client"); - assertThat(headers.getAll(ROUTE_TO_LEADER_KEY)).containsExactly("true"); } @Test @@ -358,10 +365,8 @@ public void everyPrimeCarriesFreshRequestIdWithFirstAttempt() throws Exception { } @Test - public void primeOmitsRouteToLeaderHeaderWhenDisabled() throws Exception { - DynamicChannelPoolPrimer primer = - newPrimer( - /* callCredentialsProvider= */ null, /* routeToLeader= */ false, Duration.ofSeconds(5)); + public void primeOmitsRouteToLeaderHeader() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); @@ -370,12 +375,12 @@ public void primeOmitsRouteToLeaderHeaderWhenDisabled() throws Exception { } @Test - public void primePrefersCallCredentialsProviderOverDefaultCredentials() throws Exception { - DynamicChannelPoolPrimer primer = - newPrimer( - () -> credentialsWithToken(PROVIDER_TOKEN), - /* routeToLeader= */ true, - Duration.ofSeconds(5)); + public void primeUsesCallCredentialsProviderOverDefaultCredentials() throws Exception { + CallCredentialsProvider credentialsProvider = + GapicSpannerRpc.createChannelPrimeCallCredentialsProvider( + () -> oauthCredentialsWithToken(DEFAULT_TOKEN), + () -> credentialsWithToken(PROVIDER_TOKEN)); + DynamicChannelPoolPrimer primer = newPrimer(credentialsProvider, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); @@ -386,8 +391,10 @@ public void primePrefersCallCredentialsProviderOverDefaultCredentials() throws E @Test public void primeFallsBackToDefaultCredentialsWhenProviderReturnsNull() throws Exception { - DynamicChannelPoolPrimer primer = - newPrimer(() -> null, /* routeToLeader= */ true, Duration.ofSeconds(5)); + CallCredentialsProvider credentialsProvider = + GapicSpannerRpc.createChannelPrimeCallCredentialsProvider( + () -> oauthCredentialsWithToken(DEFAULT_TOKEN), () -> null); + DynamicChannelPoolPrimer primer = newPrimer(credentialsProvider, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); @@ -397,321 +404,37 @@ public void primeFallsBackToDefaultCredentialsWhenProviderReturnsNull() throws E } @Test - public void primeFailsWhenRpcFails() throws Exception { - service.failWith = Status.UNAVAILABLE.withDescription("backend unavailable"); - DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - - Throwable failure = failureOf(primer.prime(newChannel())); - - assertThat(failure).isInstanceOf(StatusRuntimeException.class); - assertThat(((StatusRuntimeException) failure).getStatus().getCode()) - .isEqualTo(Status.Code.UNAVAILABLE); - // A transient failure says nothing about the session, so it stays registered. - assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); - } - - @Test - public void notFoundFailureEvictsSessionAndNextPrimeFallsBackToOlderSession() throws Exception { - service.failSessionsWith.put( - OTHER_SESSION_NAME, Status.NOT_FOUND.withDescription("Session not found")); - DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); - assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); - ManagedChannel channel = newChannel(); - - // The pool invokes prime() once per attempt: the first attempt uses the newest session, which - // is gone, and evicts it; the next attempt falls back to the older session. - Throwable failure = failureOf(primer.prime(channel)); - assertThat(((StatusRuntimeException) failure).getStatus().getCode()) - .isEqualTo(Status.Code.NOT_FOUND); - assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); - assertThat(primer.getPrimeSessions()).hasSize(1); - - assertThat(getWithin(primer.prime(channel), Duration.ofSeconds(10))).isNull(); - - assertThat(service.requests).hasSize(2); - assertThat(service.requests.get(0).getSession()).isEqualTo(OTHER_SESSION_NAME); - assertThat(service.requests.get(1).getSession()).isEqualTo(SESSION_NAME); - } - - @Test - public void failedPreconditionAboutSessionEvictsSession() throws Exception { - service.failWith = Status.FAILED_PRECONDITION.withDescription("Invalid session: expired"); - DynamicChannelPoolPrimer primer = newPrimer(); + public void primeOmitsCredentialsWithoutProviderOrDefaultCredentials() throws Exception { + CallCredentialsProvider credentialsProvider = + GapicSpannerRpc.createChannelPrimeCallCredentialsProvider(() -> null, null); + DynamicChannelPoolPrimer primer = newPrimer(credentialsProvider, Duration.ofSeconds(5)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - Throwable failure = failureOf(primer.prime(newChannel())); - - assertThat(((StatusRuntimeException) failure).getStatus().getCode()) - .isEqualTo(Status.Code.FAILED_PRECONDITION); - assertThat(primer.getPrimeSession()).isNull(); - } - - @Test - public void unrelatedFailedPreconditionKeepsSession() throws Exception { - service.failWith = Status.FAILED_PRECONDITION.withDescription("Database is not ready"); - DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - - failureOf(primer.prime(newChannel())); - - assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); - } - - @Test - public void permissionDeniedFailureEvictsSessionAndNextPrimeFallsBackToOlderSession() - throws Exception { - service.failSessionsWith.put( - OTHER_SESSION_NAME, - Status.PERMISSION_DENIED.withDescription("Caller is missing IAM permission")); - DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); - ManagedChannel channel = newChannel(); - - // The caller cannot use the newest session, so the attempt rotates to the older candidate. - Throwable failure = failureOf(primer.prime(channel)); - assertThat(((StatusRuntimeException) failure).getStatus().getCode()) - .isEqualTo(Status.Code.PERMISSION_DENIED); - assertThat(primer.getPrimeSessions()).hasSize(1); - assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); - assertThat(getWithin(primer.prime(channel), Duration.ofSeconds(10))).isNull(); - assertThat(service.requests).hasSize(2); - assertThat(service.requests.get(1).getSession()).isEqualTo(SESSION_NAME); + assertThat(service.headers.get(0).containsKey(AUTHORIZATION_KEY)).isFalse(); } @Test - public void invalidDatabaseRoleFailureEvictsSession() throws Exception { - service.failWith = - Status.INVALID_ARGUMENT.withDescription("Database role my-role is not valid"); + public void primeFailsWhenRpcFails() throws Exception { + service.failWith = Status.UNAVAILABLE.withDescription("backend unavailable"); DynamicChannelPoolPrimer primer = newPrimer(); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); Throwable failure = failureOf(primer.prime(newChannel())); + assertThat(failure).isInstanceOf(StatusRuntimeException.class); assertThat(((StatusRuntimeException) failure).getStatus().getCode()) - .isEqualTo(Status.Code.INVALID_ARGUMENT); - assertThat(primer.getPrimeSession()).isNull(); - } - - @Test - public void unavailableAndDeadlineExceededKeepSession() throws Exception { - DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - - service.failWith = Status.UNAVAILABLE.withDescription("Connection reset"); - failureOf(primer.prime(newChannel())); - assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); - - service.failWith = Status.DEADLINE_EXCEEDED.withDescription("Deadline exceeded"); - failureOf(primer.prime(newChannel())); + .isEqualTo(Status.Code.UNAVAILABLE); + // A transient failure says nothing about the session, so it stays registered. assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); } - @Test - public void candidateSpecificFailureClassification() { - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.PERMISSION_DENIED.withDescription("Permission denied").asRuntimeException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.PERMISSION_DENIED.asException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.INVALID_ARGUMENT - .withDescription("Database role my-role is not valid") - .asRuntimeException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.FAILED_PRECONDITION - .withDescription("Role my-role does not exist") - .asRuntimeException())) - .isTrue(); - // A message that merely contains the letters of the word role is not about a database role. - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.INVALID_ARGUMENT - .withDescription("Controller rejected the statement") - .asRuntimeException())) - .isFalse(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.INVALID_ARGUMENT.withDescription("Syntax error").asRuntimeException())) - .isFalse(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.UNAUTHENTICATED.withDescription("Invalid role").asRuntimeException())) - .isFalse(); - } - - @Test - public void lateRegistrationAfterUnregisterIsDropped() { - DynamicChannelPoolPrimer primer = newPrimer(); - // The initial session of a database client, and a refresh whose CreateSession is in flight - // with the same owner ticket. - long retiredTicket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - - // The database client is retired while the refresh is in flight. - primer.unregisterPrimeOwner(DATABASE_NAME, retiredTicket); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.getPrimeOwners()).isEmpty(); - - // The refresh returns after the retirement and is dropped: the database has no owner. - assertThat(primer.registerPrimeSession(DATABASE_NAME, REFRESHED_SESSION_NAME, retiredTicket)) - .isFalse(); - assertThat(primer.getPrimeSession()).isNull(); - - // The replacement client registers itself as the new owner and registers normally. - long replacementTicket = OWNER_TICKETS.incrementAndGet(); - primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); - String replacementSession = DATABASE_NAME + "/sessions/replacement-session"; - assertThat(primer.registerPrimeSession(DATABASE_NAME, replacementSession, replacementTicket)) - .isTrue(); - assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); - - // An even later response of the retired client still does not replace the new session: the - // database is owned by the replacement. - assertThat(primer.registerPrimeSession(DATABASE_NAME, REFRESHED_SESSION_NAME, retiredTicket)) - .isFalse(); - assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); - assertThat(primer.getPrimeSessions()).hasSize(1); - assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); - } - - @Test - public void registrationWithoutOwnerIsDropped() { - DynamicChannelPoolPrimer primer = newPrimer(); - - // No owner at all, for example a session that was created without an owner ticket. - assertThat(primer.registerPrimeSession(DATABASE_NAME, SESSION_NAME, 42L)).isFalse(); - assertThat(primer.getPrimeSession()).isNull(); - - // An owner of another database does not own this database. - primer.registerPrimeOwner(OTHER_DATABASE_NAME, 42L); - assertThat(primer.registerPrimeSession(DATABASE_NAME, SESSION_NAME, 42L)).isFalse(); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.registerPrimeSession(OTHER_DATABASE_NAME, OTHER_SESSION_NAME, 42L)).isTrue(); - assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); - } - - @Test - public void replacementOwnerTakesOverTheDatabase() { - DynamicChannelPoolPrimer primer = newPrimer(); - long retiredTicket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - - // The replacement registers itself before the retired client has been closed, for example - // because the close of the retired client is still running. The retired client's session is - // dropped right away, because the retired client no longer maintains it. - long replacementTicket = OWNER_TICKETS.incrementAndGet(); - primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); - String replacementSession = DATABASE_NAME + "/sessions/replacement-session"; - assertThat(primer.registerPrimeSession(DATABASE_NAME, replacementSession, replacementTicket)) - .isTrue(); - - // The late close of the retired client leaves the replacement and its session untouched. - primer.unregisterPrimeOwner(DATABASE_NAME, retiredTicket); - assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); - assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, replacementTicket); - - // Registering the same owner again is a no-op that keeps its session. - primer.registerPrimeOwner(DATABASE_NAME, replacementTicket); - assertThat(primer.getPrimeSessionName()).isEqualTo(replacementSession); - } - - @Test - public void retirementOfOneDatabaseDoesNotAffectOtherDatabases() { - DynamicChannelPoolPrimer primer = newPrimer(); - long ticket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - long otherTicket = OWNER_TICKETS.incrementAndGet(); - primer.registerPrimeOwner(OTHER_DATABASE_NAME, otherTicket); - - primer.unregisterPrimeOwner(DATABASE_NAME, ticket); - - assertThat(primer.registerPrimeSession(OTHER_DATABASE_NAME, OTHER_SESSION_NAME, otherTicket)) - .isTrue(); - assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); - assertThat(primer.getPrimeOwners()).containsExactly(OTHER_DATABASE_NAME, otherTicket); - } - - @Test - public void churnOfManyDatabaseClientsLeavesNothingBehind() { - DynamicChannelPoolPrimer primer = newPrimer(); - // A long-lived client that churns through many database names, including repeated clients of - // the same database, must not retain anything for retired clients. - for (int i = 0; i < 10_000; i++) { - String databaseName = "projects/p/instances/i/databases/d" + (i % 100); - long ticket = OWNER_TICKETS.incrementAndGet(); - primer.registerPrimeOwner(databaseName, ticket); - assertThat( - primer.registerPrimeSession(databaseName, databaseName + "/sessions/s" + i, ticket)) - .isTrue(); - assertThat(primer.getPrimeOwners()).hasSize(1); - assertThat(primer.getPrimeSessions()).hasSize(1); - primer.unregisterPrimeOwner(databaseName, ticket); - // A late CreateSession response of the retired client finds no owner and is dropped. - assertThat( - primer.registerPrimeSession( - databaseName, databaseName + "/sessions/late" + i, ticket)) - .isFalse(); - assertThat(primer.getPrimeOwners()).isEmpty(); - assertThat(primer.getPrimeSessions()).isEmpty(); - } - } - - @Test - public void invalidSessionFailureClassification() { - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.NOT_FOUND.withDescription("Session not found").asRuntimeException())) - .isTrue(); - assertThat(DynamicChannelPoolPrimer.isCandidateSpecificFailure(Status.NOT_FOUND.asException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.FAILED_PRECONDITION - .withDescription("Session does not exist") - .asRuntimeException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.FAILED_PRECONDITION.withDescription("invalid session").asRuntimeException())) - .isTrue(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.FAILED_PRECONDITION.asRuntimeException())) - .isFalse(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.FAILED_PRECONDITION - .withDescription("Transaction was already committed") - .asRuntimeException())) - .isFalse(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.UNAVAILABLE.withDescription("Session not found").asRuntimeException())) - .isFalse(); - assertThat( - DynamicChannelPoolPrimer.isCandidateSpecificFailure( - Status.DEADLINE_EXCEEDED.asRuntimeException())) - .isFalse(); - assertThat(DynamicChannelPoolPrimer.isCandidateSpecificFailure(new IllegalStateException())) - .isFalse(); - } - @Test public void primeFailsWithDeadlineExceededWhenServerDoesNotRespond() throws Exception { service.holdResponses = true; DynamicChannelPoolPrimer primer = - newPrimer( - /* callCredentialsProvider= */ null, /* routeToLeader= */ true, Duration.ofMillis(200)); + newPrimer(() -> credentialsWithToken(DEFAULT_TOKEN), Duration.ofMillis(200)); registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); Throwable failure = failureOf(primer.prime(newChannel())); @@ -802,18 +525,16 @@ public void concurrentPrimesOnOneChannelAllSucceed() throws Exception { @Test public void latestRegisteredSessionIsUsedForPriming() throws Exception { DynamicChannelPoolPrimer primer = newPrimer(); - assertThat(primer.getPrimeSession()).isNull(); assertThat(primer.getPrimeSessionName()).isNull(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + MutableSessionSource source = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); - // A session that is created later, for example the refresh of the multiplexed session, - // replaces the previous prime session of that database. - registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); + // Refreshing changes what the source returns without changing its registration. + source.sessionName = REFRESHED_SESSION_NAME; assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); - assertThat(primer.getPrimeSessions()).hasSize(1); + assertThat(primer.getPrimeSessionSources()).containsExactly(source); getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); assertThat(service.requests).hasSize(2); @@ -822,79 +543,41 @@ public void latestRegisteredSessionIsUsedForPriming() throws Exception { } @Test - public void newestGenerationIsSelectedAcrossDatabases() { + public void primeAttemptsRotateAcrossAvailableSources() { DynamicChannelPoolPrimer primer = newPrimer(); + MutableSessionSource first = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + MutableSessionSource second = + registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); - + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); - assertThat(sessionNames(primer)).containsExactly(OTHER_SESSION_NAME, SESSION_NAME).inOrder(); - assertThat(primer.getPrimeSessions().get(0).getGeneration()) - .isGreaterThan(primer.getPrimeSessions().get(1).getGeneration()); - - // The refresh of the first database replaces its entry and makes it the newest again. - registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); - assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); - assertThat(sessionNames(primer)) - .containsExactly(REFRESHED_SESSION_NAME, OTHER_SESSION_NAME) - .inOrder(); - assertThat(primer.getPrimeSessions().get(0).getDatabaseName()).isEqualTo(DATABASE_NAME); - assertThat(primer.getPrimeSessions().get(0).getGeneration()) - .isGreaterThan(primer.getPrimeSessions().get(1).getGeneration()); + // An unavailable source is skipped without blocking, while the cursor still advances. + first.sessionName = null; + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + assertThat(primer.getPrimeSessionSources()).containsExactly(first, second).inOrder(); } @Test - public void evictionNeverClobbersNewerEntryOfSameDatabase() { + public void unregisterIsTargetedAndIdempotent() { DynamicChannelPoolPrimer primer = newPrimer(); - registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - PrimeSession stale = primer.getPrimeSession(); - assertThat(stale).isNotNull(); - - // A newer session of the same database was registered before the stale one is evicted. - registerPrimeSession(primer, DATABASE_NAME, REFRESHED_SESSION_NAME); - PrimeSession current = primer.getPrimeSession(); - assertThat(current).isNotEqualTo(stale); - - assertThat(primer.evictPrimeSession(stale)).isFalse(); - assertThat(primer.getPrimeSession()).isEqualTo(current); - assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); + EqualSessionSource first = new EqualSessionSource(SESSION_NAME); + EqualSessionSource second = new EqualSessionSource(OTHER_SESSION_NAME); + assertThat(first).isEqualTo(second); + primer.registerPrimeSessionSource(first); + primer.registerPrimeSessionSource(first); + primer.registerPrimeSessionSource(second); - assertThat(primer.evictPrimeSession(current)).isTrue(); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.evictPrimeSession(current)).isFalse(); - } + primer.unregisterPrimeSessionSource(first); + primer.unregisterPrimeSessionSource(first); - @Test - public void unregisterRemovesOnlyThatDatabaseClient() { - DynamicChannelPoolPrimer primer = newPrimer(); - long ticket = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); - long otherTicket = registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); - - primer.unregisterPrimeOwner(OTHER_DATABASE_NAME, otherTicket); - assertThat(sessionNames(primer)).containsExactly(SESSION_NAME); - assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, ticket); - - // Unknown databases and unknown tickets are ignored. - primer.unregisterPrimeOwner( - "projects/my-project/instances/my-instance/databases/unknown", ticket); - primer.unregisterPrimeOwner(DATABASE_NAME, ticket + 1000); - assertThat(sessionNames(primer)).containsExactly(SESSION_NAME); - assertThat(primer.getPrimeOwners()).containsExactly(DATABASE_NAME, ticket); - - primer.unregisterPrimeOwner(DATABASE_NAME, ticket); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.getPrimeSessions()).isEmpty(); - assertThat(primer.getPrimeOwners()).isEmpty(); - } - - private static List sessionNames(DynamicChannelPoolPrimer primer) { - List names = new ArrayList<>(); - for (PrimeSession entry : primer.getPrimeSessions()) { - names.add(entry.getSessionName()); - } - return names; + assertThat(primer.getPrimeSessionSources()).containsExactly(second); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + primer.unregisterPrimeSessionSource(second); + assertThat(primer.getPrimeSessionSources()).isEmpty(); + assertThat(primer.getPrimeSessionName()).isNull(); } @Test @@ -929,7 +612,6 @@ public void rpcDeadlineNormallyStaysBelowPoolPrimeTimeout() { assertThat( newPrimer( /* callCredentialsProvider= */ null, - /* routeToLeader= */ true, DynamicChannelPoolPrimer.rpcDeadlineFor(Duration.ofNanos(1))) .getRpcDeadline()) .isEqualTo(DynamicChannelPoolPrimer.MIN_RPC_DEADLINE); @@ -944,22 +626,11 @@ public void rpcDeadlineNormallyStaysBelowPoolPrimeTimeout() { } @Test - public void emptyNamesAreRejected() { + public void nullSourcesAreRejected() { DynamicChannelPoolPrimer primer = newPrimer(); - assertThrows( - IllegalArgumentException.class, () -> primer.registerPrimeSession(DATABASE_NAME, "", 1L)); - assertThrows( - IllegalArgumentException.class, () -> primer.registerPrimeSession(DATABASE_NAME, null, 1L)); - assertThrows( - IllegalArgumentException.class, () -> primer.registerPrimeSession("", SESSION_NAME, 1L)); - assertThrows( - IllegalArgumentException.class, () -> primer.registerPrimeSession(null, SESSION_NAME, 1L)); - assertThrows(IllegalArgumentException.class, () -> primer.registerPrimeOwner("", 1L)); - assertThrows(IllegalArgumentException.class, () -> primer.registerPrimeOwner(null, 1L)); - assertThrows(IllegalArgumentException.class, () -> primer.unregisterPrimeOwner("", 1L)); - assertThrows(IllegalArgumentException.class, () -> primer.unregisterPrimeOwner(null, 1L)); - assertThat(primer.getPrimeSession()).isNull(); - assertThat(primer.getPrimeOwners()).isEmpty(); + assertThrows(NullPointerException.class, () -> primer.registerPrimeSessionSource(null)); + assertThrows(NullPointerException.class, () -> primer.unregisterPrimeSessionSource(null)); + assertThat(primer.getPrimeSessionSources()).isEmpty(); } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java index 8c31cfc350ea..72c0f47243be 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java @@ -20,7 +20,6 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; @@ -71,7 +70,6 @@ import com.google.cloud.spanner.spi.v1.GapicSpannerRpc.AdminRequestsLimitExceededRetryAlgorithm; import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.Futures; import com.google.protobuf.ListValue; import com.google.rpc.ErrorInfo; @@ -99,7 +97,6 @@ import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder; import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; import io.grpc.protobuf.lite.ProtoLiteUtils; -import io.grpc.stub.StreamObserver; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; import io.opentelemetry.context.propagation.ContextPropagators; @@ -353,11 +350,6 @@ private static GapicSpannerRpc getRpc(Spanner spanner) throws Exception { return (GapicSpannerRpc) method.invoke(spanner.getOptions()); } - /** The CreateSession options of a database client with the given prime owner ticket. */ - private static Map primeOwner(long ownerTicket) { - return ImmutableMap.of(Option.CHANNEL_PRIME_OWNER, ownerTicket); - } - private static void registerPrimeStatementResult() { mockSpanner.putStatementResult( StatementResult.query( @@ -654,7 +646,7 @@ public void testDynamicChannelPoolPrimesScaledUpChannelsWithSelectOne() throws E .containsExactly("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE]"); assertThat(headers.get("x-goog-request-params")) .containsExactly("session=" + URLEncoder.encode(multiplexedSession, "UTF-8")); - assertThat(headers.get("x-goog-spanner-route-to-leader")).containsExactly("true"); + assertThat(headers).doesNotContainKey("x-goog-spanner-route-to-leader"); // The request id uses the client id of the rpc, the unknown channel 0 because the channel // is not part of the pool yet, and attempt 1 because the pool invokes prime() per attempt. assertThat(headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME)).hasSize(1); @@ -1386,433 +1378,6 @@ public void testCreateSession_whenMultiplexedSessionIsFalse_assertSessionProto() rpc.shutdown(); } - @Test - public void testMultiplexedCreateSessionRecordsLatestPrimeSession() { - SpannerOptions options = - createSpannerOptions().toBuilder() - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .build(); - GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); - try { - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - assertNull(primer.getPrimeSessionName()); - rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); - - // A regular session never becomes the prime session. - Session regular = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), false); - assertFalse(regular.getMultiplexed()); - assertNull(primer.getPrimeSessionName()); - - // A multiplexed session that is created without an owner ticket, for example by a batch - // client, never becomes the prime session either. - Session unowned = rpc.createSession("DATABASE_NAME", null, null, null, true); - assertTrue(unowned.getMultiplexed()); - assertNull(primer.getPrimeSessionName()); - - // The first successfully created multiplexed session of the owner becomes the prime session. - Session first = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); - assertTrue(first.getMultiplexed()); - assertEquals(first.getName(), primer.getPrimeSessionName()); - - // A failed multiplexed CreateSession leaves the prime session untouched. - mockSpanner.setCreateSessionExecutionTime( - SimulatedExecutionTime.ofException( - Status.PERMISSION_DENIED.withDescription("test").asRuntimeException())); - SpannerException exception = - assertThrows( - SpannerException.class, - () -> rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true)); - assertEquals(ErrorCode.PERMISSION_DENIED, exception.getErrorCode()); - assertEquals(first.getName(), primer.getPrimeSessionName()); - - // The periodic refresh of the multiplexed session carries the same ticket and replaces the - // prime session of its database. - Session refreshed = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); - assertTrue(refreshed.getMultiplexed()); - assertNotEquals(first.getName(), refreshed.getName()); - assertEquals(refreshed.getName(), primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(1); - - // The multiplexed session of another database client replaces the prime session as well. - rpc.registerChannelPrimeOwner("OTHER_DATABASE_NAME", 2L); - Session second = rpc.createSession("OTHER_DATABASE_NAME", null, null, primeOwner(2L), true); - assertTrue(second.getMultiplexed()); - assertNotEquals(refreshed.getName(), second.getName()); - assertEquals(second.getName(), primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(2); - } finally { - rpc.shutdown(); - } - } - - @Test - public void testChurnOfDatabaseClientsLeavesNoPrimeOwnersOrSessionsBehind() { - SpannerOptions options = - createSpannerOptions().toBuilder() - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .build(); - GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); - try { - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - for (int i = 0; i < 50; i++) { - String databaseName = "projects/p/instances/i/databases/d" + i; - long ticket = 100L + i; - rpc.registerChannelPrimeOwner(databaseName, ticket); - Session session = rpc.createSession(databaseName, null, null, primeOwner(ticket), true); - assertEquals(session.getName(), primer.getPrimeSessionName()); - assertThat(primer.getPrimeOwners()).containsExactly(databaseName, ticket); - rpc.unregisterChannelPrimeOwner(databaseName, ticket); - assertThat(primer.getPrimeOwners()).isEmpty(); - assertThat(primer.getPrimeSessions()).isEmpty(); - } - } finally { - rpc.shutdown(); - } - } - - @Test - public void testDynamicChannelPoolPrimingFallsBackToOlderSessionWhenNewestIsInvalid() - throws Exception { - registerPrimeStatementResult(); - // Keep the load queries open long enough to trigger a scale-up. - mockSpanner.setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(3000, 0)); - SpannerOptions options = - createSpannerOptions().toBuilder() - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .setGcpChannelPoolOptions( - GcpChannelPoolOptions.newBuilder() - .setInitSize(1) - .setMinSize(1) - .setMaxSize(4) - .setDynamicScaling(1, 2, Duration.ofMinutes(3)) - .build()) - .build(); - ExecutorService executor = Executors.newFixedThreadPool(8); - try (Spanner spanner = options.getService()) { - GapicSpannerRpc rpc = getRpc(spanner); - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - GcpManagedChannel pool = findGrpcGcpChannel(rpc); - assertNotNull(pool); - - // The older database client registers its multiplexed session first. - DatabaseClient older = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - awaitCondition(() -> primer.getPrimeSessionName() != null, Duration.ofSeconds(15)); - String olderSession = primer.getPrimeSessionName(); - assertThat(olderSession) - .startsWith("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE]/"); - // The newer database client registers the session that priming prefers. - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE2]")); - awaitCondition(() -> primer.getPrimeSessions().size() == 2, Duration.ofSeconds(15)); - String newestSession = primer.getPrimeSessionName(); - assertThat(newestSession) - .startsWith("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE2]/"); - // The newest session becomes invalid on the backend: the mock returns NOT_FOUND for it. - rpc.deleteSession(newestSession, null); - - // Six concurrent streams on one channel with maxRpcPerChannel=2 force a scale-up. - List> loads = new ArrayList<>(); - for (int i = 0; i < 6; i++) { - loads.add( - executor.submit( - () -> { - long rows = 0; - try (ResultSet resultSet = older.singleUse().executeQuery(SELECT1AND2)) { - while (resultSet.next()) { - rows++; - } - } - return rows; - })); - } - - // Priming recovers on the older session within the pool's own retries: the first attempt - // of every channel of the scale-up event uses the newest session, fails with NOT_FOUND and - // evicts it, and the retries after the pool's backoff all use the older session. - awaitCondition(() -> pool.getNumberOfChannels() > 1, Duration.ofSeconds(15)); - for (Future load : loads) { - assertEquals(2L, load.get().longValue()); - } - List primes = primeRequests(); - assertThat(primes.size()).isAtLeast(2); - assertEquals(newestSession, primes.get(0).getSession()); - int newestSessionPrimes = 0; - int olderSessionPrimes = 0; - for (ExecuteSqlRequest prime : primes) { - if (prime.getSession().equals(newestSession)) { - // The pool may add more than one channel per scale-up event, and their first attempts - // all start before the first failure evicts the newest session. No attempt uses the - // newest session once a retry has started. - assertEquals(0, olderSessionPrimes); - newestSessionPrimes++; - } else { - assertEquals(olderSession, prime.getSession()); - olderSessionPrimes++; - } - } - assertThat(newestSessionPrimes).isAtLeast(1); - // Every scaled-up channel was published after one successful retry on the older session. - assertEquals(pool.getNumberOfChannels() - 1, olderSessionPrimes); - assertEquals(olderSession, primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(1); - } finally { - executor.shutdownNow(); - } - } - - @Test - public void testLateMultiplexedCreateSessionAfterUnregisterNeverRecordsPrimeSession() - throws Exception { - SpannerOptions options = - createSpannerOptions().toBuilder() - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .build(); - GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); - ExecutorService executor = Executors.newSingleThreadExecutor(); - try { - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); - Session initial = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); - assertEquals(initial.getName(), primer.getPrimeSessionName()); - - // A refresh of the multiplexed session is still in flight when the database client that - // owns it is retired, which unregisters it as the owner of its database. - mockSpanner.freeze(); - Future refresh = - executor.submit( - () -> rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true)); - awaitCondition( - () -> mockSpanner.getRequestsOfType(CreateSessionRequest.class).size() == 2, - Duration.ofSeconds(15)); - rpc.unregisterChannelPrimeOwner("DATABASE_NAME", 1L); - assertNull(primer.getPrimeSessionName()); - assertThat(primer.getPrimeOwners()).isEmpty(); - - // The late response of the retired client is dropped. - mockSpanner.unfreeze(); - Session late = refresh.get(15, TimeUnit.SECONDS); - assertTrue(late.getMultiplexed()); - assertNull(primer.getPrimeSessionName()); - - // The replacement client registers its session normally, and it stays selected. - rpc.registerChannelPrimeOwner("DATABASE_NAME", 2L); - Session replacement = rpc.createSession("DATABASE_NAME", null, null, primeOwner(2L), true); - assertEquals(replacement.getName(), primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(1); - // An even later response of the retired client, for example a retried refresh, finds the - // replacement as the owner and is dropped as well. - Session evenLater = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); - assertTrue(evenLater.getMultiplexed()); - assertEquals(replacement.getName(), primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(1); - } finally { - mockSpanner.unfreeze(); - executor.shutdownNow(); - rpc.shutdown(); - } - } - - @Test - public void testDynamicChannelPoolPrimingRotatesToOlderSessionWhenNewestIsDenied() - throws Exception { - // The newest database rejects the priming query with PERMISSION_DENIED, for example because - // the caller lost access to it. The mock records the denied primes itself, because it does not - // pass them on to the regular handler. - String deniedDatabasePrefix = "projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE2]/"; - List deniedPrimes = new CopyOnWriteArrayList<>(); - MockSpannerServiceImpl denyingSpanner = - new MockSpannerServiceImpl() { - @Override - public void executeSql( - ExecuteSqlRequest request, - StreamObserver responseObserver) { - if (request.getSql().equals(DynamicChannelPoolPrimer.PRIME_SQL) - && request.getSession().startsWith(deniedDatabasePrefix)) { - deniedPrimes.add(request); - responseObserver.onError( - Status.PERMISSION_DENIED - .withDescription("Caller is missing IAM permission spanner.databases.select") - .asRuntimeException()); - return; - } - super.executeSql(request, responseObserver); - } - }; - denyingSpanner.setAbortProbability(0.0D); - denyingSpanner.putStatementResult(StatementResult.query(SELECT1AND2, SELECT1_RESULTSET)); - denyingSpanner.putStatementResult( - StatementResult.query( - Statement.of(DynamicChannelPoolPrimer.PRIME_SQL), - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues( - com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .setMetadata(SELECT1AND2_METADATA) - .build())); - // Keep the load queries open long enough to trigger a scale-up. - denyingSpanner.setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(3000, 0)); - Server denyingServer = - NettyServerBuilder.forAddress(new InetSocketAddress("localhost", 0)) - .addService(denyingSpanner) - .build() - .start(); - ExecutorService executor = Executors.newFixedThreadPool(8); - try { - SpannerOptions options = - createSpannerOptions().toBuilder() - .setHost("http://localhost:" + denyingServer.getPort()) - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .setGcpChannelPoolOptions( - GcpChannelPoolOptions.newBuilder() - .setInitSize(1) - .setMinSize(1) - .setMaxSize(4) - .setDynamicScaling(1, 2, Duration.ofMinutes(3)) - .build()) - .build(); - try (Spanner spanner = options.getService()) { - GapicSpannerRpc rpc = getRpc(spanner); - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - GcpManagedChannel pool = findGrpcGcpChannel(rpc); - assertNotNull(pool); - - DatabaseClient older = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - awaitCondition(() -> primer.getPrimeSessionName() != null, Duration.ofSeconds(15)); - String olderSession = primer.getPrimeSessionName(); - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE2]")); - awaitCondition(() -> primer.getPrimeSessions().size() == 2, Duration.ofSeconds(15)); - String newestSession = primer.getPrimeSessionName(); - assertThat(newestSession).startsWith(deniedDatabasePrefix); - - // Six concurrent streams on one channel with maxRpcPerChannel=2 force a scale-up. - List> loads = new ArrayList<>(); - for (int i = 0; i < 6; i++) { - loads.add( - executor.submit( - () -> { - long rows = 0; - try (ResultSet resultSet = older.singleUse().executeQuery(SELECT1AND2)) { - while (resultSet.next()) { - rows++; - } - } - return rows; - })); - } - - // The first attempt uses the newest session, is denied and evicts exactly that entry; - // the pool's retry succeeds on the older candidate. - awaitCondition(() -> pool.getNumberOfChannels() > 1, Duration.ofSeconds(15)); - for (Future load : loads) { - assertEquals(2L, load.get().longValue()); - } - assertThat(deniedPrimes).isNotEmpty(); - for (ExecuteSqlRequest denied : deniedPrimes) { - assertEquals(newestSession, denied.getSession()); - } - List succeededPrimes = new ArrayList<>(); - for (ExecuteSqlRequest request : - denyingSpanner.getRequestsOfType(ExecuteSqlRequest.class)) { - if (request.getSql().equals(DynamicChannelPoolPrimer.PRIME_SQL)) { - succeededPrimes.add(request); - } - } - assertEquals(pool.getNumberOfChannels() - 1, succeededPrimes.size()); - for (ExecuteSqlRequest prime : succeededPrimes) { - assertEquals(olderSession, prime.getSession()); - } - assertEquals(olderSession, primer.getPrimeSessionName()); - assertThat(primer.getPrimeSessions()).hasSize(1); - } - } finally { - executor.shutdownNow(); - denyingServer.shutdownNow(); - denyingServer.awaitTermination(); - } - } - - @Test - public void testMultiplexedCreateSessionWithoutEchoedFlagNeverRecordsPrimeSession() - throws Exception { - // A backend that ignores the multiplexed flag of the request creates a regular session, which - // is not safe for the concurrent use that priming implies and must never become the prime - // session. - MockSpannerServiceImpl nonEchoingSpanner = - new MockSpannerServiceImpl() { - @Override - public void createSession( - CreateSessionRequest request, StreamObserver responseObserver) { - super.createSession( - request, - new StreamObserver() { - @Override - public void onNext(Session session) { - responseObserver.onNext(session.toBuilder().setMultiplexed(false).build()); - } - - @Override - public void onError(Throwable t) { - responseObserver.onError(t); - } - - @Override - public void onCompleted() { - responseObserver.onCompleted(); - } - }); - } - }; - Server nonEchoingServer = - NettyServerBuilder.forAddress(new InetSocketAddress("localhost", 0)) - .addService(nonEchoingSpanner) - .build() - .start(); - try { - SpannerOptions options = - createSpannerOptions().toBuilder() - .setHost("http://localhost:" + nonEchoingServer.getPort()) - .enableGrpcGcpExtension() - .enableDynamicChannelPool() - .build(); - GapicSpannerRpc rpc = new GapicSpannerRpc(options, true); - try { - DynamicChannelPoolPrimer primer = rpc.getChannelPrimer(); - assertNotNull(primer); - rpc.registerChannelPrimeOwner("DATABASE_NAME", 1L); - - Session session = rpc.createSession("DATABASE_NAME", null, null, primeOwner(1L), true); - assertFalse(session.getMultiplexed()); - assertTrue( - nonEchoingSpanner - .getRequestsOfType(CreateSessionRequest.class) - .get(0) - .getSession() - .getMultiplexed()); - assertNull(primer.getPrimeSessionName()); - } finally { - rpc.shutdown(); - } - } finally { - nonEchoingServer.shutdownNow(); - nonEchoingServer.awaitTermination(); - } - } - @Test public void testChannelEndpointCacheFactoryUsedWhenLocationApiEnabled() { AtomicBoolean factoryCalled = new AtomicBoolean(false); From 51a6fd39d0214694edad2bb7c5f28c8893bdfc2a Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 17:46:26 +0530 Subject: [PATCH 5/7] chore(spanner): harden channel primer failure handling Ensure the primer API boundary converts synchronous failures into failed futures. Restore the built-in metrics test to the merged upstream emulator-host setup. --- .../spi/v1/DynamicChannelPoolPrimer.java | 28 ++++++++++++------- ...iltInOpenTelemetryMetricsProviderTest.java | 4 --- .../spi/v1/DynamicChannelPoolPrimerTest.java | 16 +++++++++++ 3 files changed, 34 insertions(+), 14 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java index 0ec2f9ca2958..c514ae1a851f 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -236,19 +236,27 @@ String getPrimeSessionName() { return null; } + /** + * Always returns a future and never throws, so every failure reaches the pool as a failed future. + */ @Override public ListenableFuture prime(ManagedChannel channel) { - String sessionName = getPrimeSessionName(); - if (sessionName == null) { - // The primer cannot gate the pool's scale-up decision, so the attempt fails fast. The pool's - // retry with backoff and its close-on-failure behaviour handle the unavailable session. - return Futures.immediateFailedFuture( - SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - "Cannot prime a dynamic channel pool channel before a multiplexed session is" - + " available")); + try { + String sessionName = getPrimeSessionName(); + if (sessionName == null) { + // The primer cannot gate the pool's scale-up decision, so the attempt fails fast. The + // pool's retry with backoff and its close-on-failure behaviour handle the unavailable + // session. + return Futures.immediateFailedFuture( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "Cannot prime a dynamic channel pool channel before a multiplexed session is" + + " available")); + } + return executePrimeStatement(channel, sessionName); + } catch (Throwable t) { + return Futures.immediateFailedFuture(t); } - return executePrimeStatement(channel, sessionName); } private ListenableFuture executePrimeStatement(ManagedChannel channel, String sessionName) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java index 8ebc51e578f6..9cbffdda1c0e 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProviderTest.java @@ -110,10 +110,6 @@ private void verifyHash(String hash) { private SpannerOptions newTestOptions() { return SpannerOptions.newBuilder() - // The builder picks up SPANNER_EMULATOR_HOST from the environment, and building with an - // emulator host replaces the credentials below with NoCredentials. Built-in metrics are - // disabled for a client without credentials, so leave the emulator out of these tests. - .setEmulatorHost(null) .setProjectId("host-project") // The builder picks up SPANNER_EMULATOR_HOST from the environment, and building with an // emulator host replaces the credentials below with NoCredentials. Built-in metrics are diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java index bf2547fdbb38..1467d0d7ba3a 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java @@ -415,6 +415,22 @@ public void primeOmitsCredentialsWithoutProviderOrDefaultCredentials() throws Ex assertThat(service.headers.get(0).containsKey(AUTHORIZATION_KEY)).isFalse(); } + @Test + public void primeReturnsFailedFutureWhenCallCredentialsProviderThrows() throws Exception { + RuntimeException expected = new RuntimeException("credentials unavailable"); + DynamicChannelPoolPrimer primer = + newPrimer( + () -> { + throw expected; + }, + Duration.ofSeconds(5)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + ListenableFuture future = primer.prime(newChannel()); + + assertThat(failureOf(future)).isSameInstanceAs(expected); + } + @Test public void primeFailsWhenRpcFails() throws Exception { service.failWith = Status.UNAVAILABLE.withDescription("backend unavailable"); From 1d01be3eb39aecb16dd77e4e02259877b71e6dbd Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 20:55:40 +0530 Subject: [PATCH 6/7] docs(spanner): trim comments documenting review discussion --- .../MultiplexedSessionDatabaseClient.java | 17 ++--------------- .../google/cloud/spanner/SpannerOptions.java | 4 +--- .../spi/v1/DynamicChannelPoolPrimer.java | 11 +++-------- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 6 ++---- 4 files changed, 8 insertions(+), 30 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 9d3d507a0355..59ad943230e0 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -258,13 +258,6 @@ private SharedChannelUsage(int numChannels) { final SettableApiFuture initialSessionReferenceFuture = SettableApiFuture.create(); this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); - // This escape is safe: resourceNotFoundException is initialized at declaration, - // multiplexedSessionReference above, and isClosed is volatile: the fields - // getChannelPrimeSessionName() reads. - // Registration adds to a CopyOnWriteArrayList under a lock, whose volatile array write happens - // before a primer reads that array, so this is safe publication, not a racy final-field escape. - // Keep registration before the first CreateSession and inside this constructor: clients are - // constructed at multiple call sites, and a missed registration would silently disable priming. spanner.getRpc().registerChannelPrimeSessionSource(this); try { @@ -277,11 +270,8 @@ private SharedChannelUsage(int numChannels) { sessionClient.getSpanner().getOptions().getSessionPoolOptions(), initialSessionReferenceFuture); } catch (Throwable t) { - // The constructor did not complete, so the caller never gets a reference to this client and - // will never close it. Undo everything that was registered above: mark the client closed so - // that a CreateSession that is still in flight neither starts the maintainer nor is handed - // to a waiter, stop the maintainer, unregister the prime-session source, and release the - // shared channel usage. + // The caller never receives this client, so it will never be closed; close() therefore undoes + // the registrations. close(); throw t; } @@ -417,9 +407,6 @@ AtomicLong getNumSessionsReleased() { @Override @Nullable public String getChannelPrimeSessionName() { - // This client already uses the returned session as multiplexed for all real traffic. If the - // backend failed to honor the multiplexed request flag, withholding it from the primer would - // not make the client safe. if (isClosed || !isValid()) { return null; } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 6e4261d21ed0..4d5b9d903508 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -2111,9 +2111,7 @@ public Builder setGrpcGcpOtelMetricsEnabled(boolean enableGrpcGcpOtelMetrics) { *

    Channels that the pool adds during scale-up are primed with {@code SELECT 1} on a * multiplexed session before they are published. A channel primer, prime timeout, or prime * attempt count that is set in the given options takes precedence over the Spanner primer and - * its defaults. The deadline of the Spanner primer's RPC is derived from the prime timeout and - * normally stays below it; only a prime timeout of two milliseconds or less yields the primer's - * minimum deadline of one millisecond, and the prime timeout then bounds the attempt. + * its defaults. * *

    Example usage: * diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java index c514ae1a851f..0d458ed1fb7d 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -68,7 +68,7 @@ * connection for the whole channel, which benefits every database that later uses it. It does not * warm database-specific server-side state, so a request served by a freshly scaled-up channel for * a database other than the one used to prime it can still see slightly higher latency than on a - * channel that database has already used. This trade-off is accepted. + * channel that database has already used. * *

    A database client unregisters itself when invalidated or closed. The primer therefore never * retains a retired client's session name, and a {@code CreateSession} that completes after its @@ -213,13 +213,8 @@ List getPrimeSessionSources() { @VisibleForTesting @Nullable String getPrimeSessionName() { - // The primer neither classifies failures nor evicts sessions. A dead session remains until its - // owner observes it and isValid() stops offering it. Client traffic normally reaches - // MultiplexedSessionTransaction.onError quickly; periodic refresh backs up an idle client. - // Attempts may thus use a dead session, and the shared cursor cannot ensure consecutive - // concurrent attempts on one channel use different sources. Deciding whether a session is still - // usable deliberately stays with the owning client, which already tracks it for its own - // traffic. + // The primer neither classifies failures nor evicts sessions; a source stops offering its + // session only when its owning client closes or becomes invalid. List sources = ImmutableList.copyOf(primeSessionSources); int size = sources.size(); if (size == 0) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index 05a1e0baa3f1..7d969d39fc59 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -804,10 +804,8 @@ private InstantiatingGrpcChannelProvider.Builder createBaseChannelProviderBuilde * as a normal Spanner call: the user-supplied {@link CallCredentialsProvider} takes precedence, * otherwise the scoped credentials that GAX attaches to every call are used, and the request ids * of priming RPCs come from the request id creator of this rpc. The deadline of the priming RPC - * is derived from the pool's prime timeout and normally stays below it; see {@link - * DynamicChannelPoolPrimer#rpcDeadlineFor(Duration)} for the minimum deadline that a prime - * timeout of two milliseconds or less yields. Live multiplexed-session database clients register - * themselves as session sources and unregister themselves when closed. + * is derived from the pool's prime timeout and normally stays below it. Live multiplexed-session + * database clients register themselves as session sources and unregister themselves when closed. */ @Nullable private DynamicChannelPoolPrimer createChannelPrimer( From 07cd5075ad63fee2fa271e9c587324e286657a6a Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 3 Sep 2026 20:57:18 +0530 Subject: [PATCH 7/7] perf(spanner): remove redundant closed-check synchronization --- .../google/cloud/spanner/MultiplexedSessionDatabaseClient.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 59ad943230e0..afe498219139 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -387,7 +387,7 @@ boolean isValid() { return resourceNotFoundException.get() == null; } - private synchronized boolean isClientClosed() { + private boolean isClientClosed() { return isClosed; }