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..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 @@ -28,8 +28,10 @@ import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; +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; @@ -38,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; @@ -48,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 @@ -186,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; @@ -253,15 +258,23 @@ private SharedChannelUsage(int numChannels) { final SettableApiFuture initialSessionReferenceFuture = SettableApiFuture.create(); this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); + spanner.getRpc().registerChannelPrimeSessionSource(this); - 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); + 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 caller never receives this client, so it will never be closed; close() therefore undoes + // the registrations. + close(); + throw t; + } } private void asyncCreateMultiplexedSession( @@ -270,10 +283,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,9 +318,11 @@ 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); @@ -363,6 +387,15 @@ boolean isValid() { return resourceNotFoundException.get() == null; } + private boolean isClientClosed() { + return isClosed; + } + + private static SpannerException newClosedException() { + return SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "This client has been closed"); + } + AtomicLong getNumSessionsAcquired() { return this.numSessionsAcquired; } @@ -371,6 +404,24 @@ AtomicLong getNumSessionsReleased() { return this.numSessionsReleased; } + @Override + @Nullable + public String getChannelPrimeSessionName() { + 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) { @@ -381,6 +432,9 @@ void close() { } } if (releaseChannelUsage) { + // 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) { @@ -659,12 +713,19 @@ void maintain() { new SessionConsumer() { @Override public void onSessionReady(SessionImpl session) { - multiplexedSessionReference.set( - ApiFutures.immediateFuture(session.getSessionReference())); - expirationDate.set( - clock - .instant() - .plus(MultiplexedSessionDatabaseClient.this.sessionExpirationDuration)); + 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)); + } } @Override 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..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,7 +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. + // 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); @@ -360,8 +361,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..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 @@ -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. 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); + + /** + * 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,16 @@ 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} * * + *

    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 + * 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 */ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() { @@ -208,13 +229,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) { @@ -251,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(); } @@ -2031,6 +2062,12 @@ 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. * + *

    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 * this method was called. @@ -2068,7 +2105,13 @@ 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. + * + *

    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. * *

    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..0d458ed1fb7d --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimer.java @@ -0,0 +1,346 @@ +/* + * 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.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; +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.stub.ClientCalls; +import java.io.UnsupportedEncodingException; +import java.net.URLEncoder; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +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. + * + *

    The channel pool is shared by all database clients of a {@link Spanner} instance, while + * 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. + * + *

    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. */ + 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); + + /** + * 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 CallCredentialsProvider callCredentialsProvider; + private final Duration rpcDeadline; + + /** + * Registered sources. Reference identity keeps overlapping client instances independent, and + * writers synchronize because the identity scan and mutation must be atomic. + */ + private final CopyOnWriteArrayList primeSessionSources = + new CopyOnWriteArrayList<>(); + + /** 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 + * 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 callCredentialsProvider provides the credentials of a normal Spanner call, or {@code + * null} when the client runs without credentials + * @param rpcDeadline the deadline of a single priming RPC + */ + DynamicChannelPoolPrimer( + SpannerMetadataProvider metadataProvider, + String projectName, + RequestIdCreator requestIdCreator, + @Nullable CallCredentialsProvider callCredentialsProvider, + Duration rpcDeadline) { + this.metadataProvider = Preconditions.checkNotNull(metadataProvider); + this.projectName = Preconditions.checkNotNull(projectName); + this.requestIdCreator = Preconditions.checkNotNull(requestIdCreator); + this.callCredentialsProvider = callCredentialsProvider; + 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 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; + } + } + primeSessionSources.add(source); + } + } + + /** 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; + } + } + } + } + + /** Returns a snapshot of registered sources in registration order. */ + @VisibleForTesting + List getPrimeSessionSources() { + return ImmutableList.copyOf(primeSessionSources); + } + + /** Returns a currently available session name without blocking, rotating the starting source. */ + @VisibleForTesting + @Nullable + String getPrimeSessionName() { + // 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) { + 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; + } + + /** + * Always returns a future and never throws, so every failure reaches the pool as a failed future. + */ + @Override + public ListenableFuture prime(ManagedChannel channel) { + 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); + } + } + + private ListenableFuture executePrimeStatement(ManagedChannel channel, String sessionName) { + ExecuteSqlRequest request = + ExecuteSqlRequest.newBuilder().setSession(sessionName).setSql(PRIME_SQL).build(); + // 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); + if (callCredentialsProvider != null) { + CallCredentials callCredentials = callCredentialsProvider.getCallCredentials(); + 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 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. + return Futures.transform( + ClientCalls.futureUnaryCall(call, request), + resultSet -> null, + MoreExecutors.directExecutor()); + } + + /** + * 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 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) { + 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( + 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 + // 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..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 @@ -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; @@ -90,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; @@ -309,6 +311,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 +376,7 @@ public GapicSpannerRpc(final SpannerOptions options) { if (initializeStubs) { CredentialsProvider credentialsProvider = GrpcTransportOptions.setUpCredentialsProvider(options); + this.channelPrimer = createChannelPrimer(options, credentialsProvider); InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder = createBaseChannelProviderBuilder( @@ -388,9 +392,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 +576,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 +652,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 +705,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 +739,7 @@ private InstantiatingGrpcChannelProvider.Builder createChannelProviderBuilder( InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder = createBaseChannelProviderBuilder( options, headerProviderWithUserAgent, isEnableDirectAccess); - maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options); + maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options, channelPrimer); return defaultChannelProviderBuilder; } @@ -790,8 +798,68 @@ 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. Live multiplexed-session + * database clients register themselves as session sources and unregister themselves when closed. + */ + @Nullable + private DynamicChannelPoolPrimer createChannelPrimer( + SpannerOptions options, CredentialsProvider credentialsProvider) { + if (!options.isGrpcGcpExtensionEnabled() || !options.isDynamicChannelPoolEnabled()) { + return null; + } + return new DynamicChannelPoolPrimer( + metadataProvider, + projectName, + requestIdCreator, + 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 + 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 +885,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 +893,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 +961,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(); @@ -1972,6 +2060,20 @@ public Session createSession( return get(spannerStub.createSessionCallable().futureCall(request, context)); } + @Override + public void registerChannelPrimeSessionSource(ChannelPrimeSessionSource source) { + if (channelPrimer != null) { + channelPrimer.registerPrimeSessionSource(source); + } + } + + @Override + public void unregisterChannelPrimeSessionSource(ChannelPrimeSessionSource source) { + if (channelPrimer != null) { + channelPrimer.unregisterPrimeSessionSource(source); + } + } + @Override public void deleteSession(String sessionName, @Nullable Map options) throws SpannerException { 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 282979b522b8..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 @@ -114,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. @@ -376,6 +383,12 @@ default Session createSession( void deleteSession(String sessionName, @Nullable Map options) throws SpannerException; + /** Registers a source of a multiplexed session for priming dynamic channel pool channels. */ + default void registerChannelPrimeSessionSource(ChannelPrimeSessionSource source) {} + + /** 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 new file mode 100644 index 000000000000..62c38511ee3b --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DynamicChannelPoolPrimeSessionTest.java @@ -0,0 +1,392 @@ +/* + * 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); + // 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 + // 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.getPrimeSessionSourceCount()).isEqualTo(1); + // 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.getPrimeSessionSourceCount()).isEqualTo(0); + 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.getPrimeSessionSourceCount()).isEqualTo(0); + // 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 churnOfDatabaseClientsLeavesNoPrimeSessionSourcesBehind() throws Exception { + // A long-lived Spanner instance that churns through many database names must not retain + // 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); + 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.getPrimeSessionSourceCount()).isEqualTo(liveDatabases.size()); + } + + spanner.close(); + + assertThat(rpc.getPrimeSessionNames()).isEmpty(); + assertThat(rpc.getPrimeSessionSourceCount()).isEqualTo(0); + } + + @Test + 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 session source 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.getPrimeSessionSourceCount()).isEqualTo(0); + 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 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.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 d398a78643aa..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,6 +19,7 @@ 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.assertNotNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeFalse; @@ -26,18 +27,26 @@ import static org.mockito.ArgumentMatchers.any; 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.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; +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; @@ -48,10 +57,14 @@ 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 +83,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()) @@ -126,6 +141,112 @@ public void testMaintainer() { assertEquals(client.getCurrentSessionReference(), session2.getSessionReference()); } + @Test + public void testClosedClientIgnoresInitialSessionThatArrivesAfterClose() { + assumeTrue(isJava8()); + 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)); + 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() { + assumeTrue(isJava8()); + 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)); + 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,91 @@ public void testCloseKeepsChannelUsageEntryWhileAnotherClientIsUsingSameSpanner( } } + @Test + public void testChannelPrimeSessionSourceLifecycleAndNonBlockingAccessor() throws Exception { + assumeTrue(isJava8()); + 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(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)); + List consumers = new ArrayList<>(); + doAnswer( + (Answer) + invocationOnMock -> { + consumers.add(invocationOnMock.getArgument(0)); + return null; + }) + .when(sessionClient) + .asyncCreateMultiplexedSession(any(SessionConsumer.class)); + + MultiplexedSessionDatabaseClient first = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + MultiplexedSessionDatabaseClient second = + new MultiplexedSessionDatabaseClient(sessionClient, clock); + + 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(); + 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()); + + first.close(); + first.close(); + verify(rpc).unregisterChannelPrimeSessionSource(first); + verify(rpc, never()).unregisterChannelPrimeSessionSource(second); + second.close(); + verify(rpc).unregisterChannelPrimeSessionSource(second); + assertThat(second.getChannelPrimeSessionName()).isNull(); + } + private SessionClient createSessionClient(SpannerImpl spanner) { return new FailingMultiplexedSessionClient(spanner); } @@ -406,9 +614,6 @@ 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()); } 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..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,6 +21,9 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +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 +42,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 +131,75 @@ public void getDbclientAgainGivesSame() { verify(spannerOptions).initializeBuiltInMetrics(db); } + @Test + public void invalidatedDatabaseClientUnregistersItsChannelPrimeSessionSource() { + 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 sources = + ArgumentCaptor.forClass(SpannerRpc.ChannelPrimeSessionSource.class); + verify(rpc).registerChannelPrimeSessionSource(sources.capture()); + SpannerRpc.ChannelPrimeSessionSource firstSource = sources.getValue(); + verify(rpc, never()).unregisterChannelPrimeSessionSource(firstSource); + + valid.set(false); + DatabaseClient second = spanner.getDatabaseClient(db); + + assertThat(second).isNotSameInstanceAs(first); + 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 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()) + .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 sources = + ArgumentCaptor.forClass(SpannerRpc.ChannelPrimeSessionSource.class); + verify(rpc, times(2)).registerChannelPrimeSessionSource(sources.capture()); + verify(rpc, never()).unregisterChannelPrimeSessionSource(any()); + + spanner.close(); + + verify(rpc).unregisterChannelPrimeSessionSource(sources.getAllValues().get(0)); + verify(rpc).unregisterChannelPrimeSessionSource(sources.getAllValues().get(1)); + verify(rpc, times(2)).unregisterChannelPrimeSessionSource(any()); + } + @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..7b46950746e8 --- /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.SpannerRpc.ChannelPrimeSessionSource; +import com.google.spanner.v1.Session; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; + +/** + * 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(); + private volatile Runnable afterUnregisterHook = () -> {}; + + public ChannelPrimerTestRpc(SpannerOptions options) { + super(options); + } + + /** Sets a hook that runs right after a database client's session source is unregistered. */ + 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 currently available session names in source-preference order. */ + public List getPrimeSessionNames() { + List names = new ArrayList<>(); + DynamicChannelPoolPrimer primer = getChannelPrimer(); + if (primer != null) { + for (ChannelPrimeSessionSource source : primer.getPrimeSessionSources()) { + String sessionName = source.getChannelPrimeSessionName(); + if (sessionName != null) { + names.add(sessionName); + } + } + } + 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 number of registered prime-session sources. */ + public int getPrimeSessionSourceCount() { + DynamicChannelPoolPrimer primer = getChannelPrimer(); + return primer == null ? 0 : primer.getPrimeSessionSources().size(); + } + + @Override + 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 new file mode 100644 index 000000000000..1467d0d7ba3a --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/DynamicChannelPoolPrimerTest.java @@ -0,0 +1,652 @@ +/* + * 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.SpannerRpc.ChannelPrimeSessionSource; +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.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +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; + + @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 = 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(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, 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, callCredentialsProvider, rpcDeadline); + } + + private DynamicChannelPoolPrimer newPrimer() { + return newPrimer(() -> credentialsWithToken(DEFAULT_TOKEN), Duration.ofSeconds(5)); + } + + private static class MutableSessionSource implements ChannelPrimeSessionSource { + @Nullable volatile String sessionName; + + private MutableSessionSource(@Nullable String sessionName) { + this.sessionName = sessionName; + } + + @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 { + 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)); + assertThat(headers.containsKey(ROUTE_TO_LEADER_KEY)).isFalse(); + 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); + // 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"); + } + + @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 primeOmitsRouteToLeaderHeader() throws Exception { + DynamicChannelPoolPrimer primer = newPrimer(); + 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 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)); + + assertThat(service.headers.get(0).getAll(AUTHORIZATION_KEY)) + .containsExactly("Bearer " + PROVIDER_TOKEN); + } + + @Test + public void primeFallsBackToDefaultCredentialsWhenProviderReturnsNull() throws Exception { + 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)); + + assertThat(service.headers.get(0).getAll(AUTHORIZATION_KEY)) + .containsExactly("Bearer " + DEFAULT_TOKEN); + } + + @Test + public void primeOmitsCredentialsWithoutProviderOrDefaultCredentials() throws Exception { + CallCredentialsProvider credentialsProvider = + GapicSpannerRpc.createChannelPrimeCallCredentialsProvider(() -> null, null); + DynamicChannelPoolPrimer primer = newPrimer(credentialsProvider, Duration.ofSeconds(5)); + registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + 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"); + 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 primeFailsWithDeadlineExceededWhenServerDoesNotRespond() throws Exception { + service.holdResponses = true; + DynamicChannelPoolPrimer primer = + newPrimer(() -> credentialsWithToken(DEFAULT_TOKEN), 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 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.getPrimeSessionName()).isNull(); + + MutableSessionSource source = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + getWithin(primer.prime(newChannel()), Duration.ofSeconds(10)); + + // Refreshing changes what the source returns without changing its registration. + source.sessionName = REFRESHED_SESSION_NAME; + assertThat(primer.getPrimeSessionName()).isEqualTo(REFRESHED_SESSION_NAME); + assertThat(primer.getPrimeSessionSources()).containsExactly(source); + 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 primeAttemptsRotateAcrossAvailableSources() { + DynamicChannelPoolPrimer primer = newPrimer(); + MutableSessionSource first = registerPrimeSession(primer, DATABASE_NAME, SESSION_NAME); + MutableSessionSource second = + registerPrimeSession(primer, OTHER_DATABASE_NAME, OTHER_SESSION_NAME); + + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + assertThat(primer.getPrimeSessionName()).isEqualTo(SESSION_NAME); + + // 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 unregisterIsTargetedAndIdempotent() { + DynamicChannelPoolPrimer primer = newPrimer(); + 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); + + primer.unregisterPrimeSessionSource(first); + primer.unregisterPrimeSessionSource(first); + + assertThat(primer.getPrimeSessionSources()).containsExactly(second); + assertThat(primer.getPrimeSessionName()).isEqualTo(OTHER_SESSION_NAME); + primer.unregisterPrimeSessionSource(second); + assertThat(primer.getPrimeSessionSources()).isEmpty(); + assertThat(primer.getPrimeSessionName()).isNull(); + } + + @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, + 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 nullSourcesAreRejected() { + DynamicChannelPoolPrimer primer = newPrimer(); + + 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 de21347e0609..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 @@ -22,6 +22,7 @@ import static org.junit.Assert.assertFalse; 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 +42,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 +66,14 @@ 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.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; @@ -105,16 +110,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 +131,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 +214,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 +261,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 +318,349 @@ 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()); + } + + 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).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); + 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 @@ -1221,9 +1580,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 +1943,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 +2019,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. }