feat(spanner): prime scaled-up dynamic channel pool channels with SELECT 1

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<SessionReference> 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<SessionReference> 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<SessionReference> 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

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<DatabaseId, DatabaseClientImpl> dbClient : dbClients.entrySet()) {
+        closureFutures.add(dbClient.getValue().closeAsync(closedException));
       }
       dbClients.clear();
       Futures.successfulAsList(closureFutures).get(timeout, unit);

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java

@@ -178,6 +178,19 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
    */
   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<Spanner, SpannerOptions> {
    *   <li>Scale down interval: 3 minutes
    *   <li>Affinity key lifetime: 10 minutes
    *   <li>Cleanup interval: 1 minute
+   *   <li>Channel prime timeout: 10 seconds
+   *   <li>Channel prime max attempts: {@value #DEFAULT_DYNAMIC_POOL_CHANNEL_PRIME_MAX_ATTEMPTS}
    * </ul>
    *
+   * <p>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.
      *
+     * <p>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.
+     *
      * <p>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.
      *
      * <p>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.
+     *
+     * <p>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.
      *
      * <p>Example usage:
      *