fix: avoid unbalanced session pool creation

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

@@ -54,6 +54,7 @@
 import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.cloud.spanner.SpannerImpl.ClosedException;
+import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
@@ -1372,6 +1373,13 @@ final class PooledSession implements Session {
     private volatile SpannerException lastException;
     private volatile boolean allowReplacing = true;
 
+    /**
+     * This ensures that the session is added at a random position in the pool the first time it is
+     * actually added to the pool.
+     */
+    @GuardedBy("lock")
+    private Position releaseToPosition = initialReleasePosition;
+
     /**
      * Property to mark if the session is eligible to be long-running. This can only be true if the
      * session is executing certain types of transactions (for ex - Partitioned DML) which can be
@@ -1403,6 +1411,13 @@ private PooledSession(SessionImpl delegate) {
       this.lastUseTime = clock.instant();
     }
 
+    int getChannel() {
+      Long channelHint = (Long) delegate.getOptions().get(SpannerRpc.Option.CHANNEL_HINT);
+      return channelHint == null
+          ? 0
+          : (int) (channelHint % sessionClient.getSpanner().getOptions().getNumChannels());
+    }
+
     @Override
     public String toString() {
       return getName();
@@ -1536,7 +1551,7 @@ public void close() {
         if (state != SessionState.CLOSING) {
           state = SessionState.AVAILABLE;
         }
-        releaseSession(this, Position.FIRST);
+        releaseSession(this, false);
       }
     }
 
@@ -1576,7 +1591,7 @@ private void determineDialectAsync(final SettableFuture<Dialect> dialect) {
               // in the database dialect, and there's nothing sensible that we can do with it here.
               dialect.setException(t);
             } finally {
-              releaseSession(this, Position.FIRST);
+              releaseSession(this, false);
             }
           });
     }
@@ -1830,7 +1845,7 @@ private void keepAliveSessions(Instant currTime) {
           logger.log(Level.FINE, "Keeping alive session " + sessionToKeepAlive.getName());
           numSessionsToKeepAlive--;
           sessionToKeepAlive.keepAlive();
-          releaseSession(sessionToKeepAlive, Position.FIRST);
+          releaseSession(sessionToKeepAlive, false);
         } catch (SpannerException e) {
           handleException(e, sessionToKeepAlive);
         }
@@ -1929,7 +1944,7 @@ private void removeLongRunningSessions(
     }
   }
 
-  private enum Position {
+  enum Position {
     FIRST,
     RANDOM
   }
@@ -1962,6 +1977,15 @@ private enum Position {
 
   final PoolMaintainer poolMaintainer;
   private final Clock clock;
+  /**
+   * initialReleasePosition determines where in the pool sessions are added when they are released
+   * into the pool the first time. This is always RANDOM in production, but some tests use FIRST to
+   * be able to verify the order of sessions in the pool. Using RANDOM ensures that we do not get an
+   * unbalanced session pool where all sessions belonging to one gRPC channel are added to the same
+   * region in the pool.
+   */
+  private final Position initialReleasePosition;
+
   private final Object lock = new Object();
   private final Random random = new Random();
 
@@ -2045,6 +2069,7 @@ static SessionPool createPool(
         ((GrpcTransportOptions) spannerOptions.getTransportOptions()).getExecutorFactory(),
         sessionClient,
         poolMaintainerClock == null ? new Clock() : poolMaintainerClock,
+        Position.RANDOM,
         Metrics.getMetricRegistry(),
         labelValues);
   }
@@ -2053,20 +2078,22 @@ static SessionPool createPool(
       SessionPoolOptions poolOptions,
       ExecutorFactory<ScheduledExecutorService> executorFactory,
       SessionClient sessionClient) {
-    return createPool(poolOptions, executorFactory, sessionClient, new Clock());
+    return createPool(poolOptions, executorFactory, sessionClient, new Clock(), Position.RANDOM);
   }
 
   static SessionPool createPool(
       SessionPoolOptions poolOptions,
       ExecutorFactory<ScheduledExecutorService> executorFactory,
       SessionClient sessionClient,
-      Clock clock) {
+      Clock clock,
+      Position initialReleasePosition) {
     return createPool(
         poolOptions,
         null,
         executorFactory,
         sessionClient,
         clock,
+        initialReleasePosition,
         Metrics.getMetricRegistry(),
         SPANNER_DEFAULT_LABEL_VALUES);
   }
@@ -2077,6 +2104,7 @@ static SessionPool createPool(
       ExecutorFactory<ScheduledExecutorService> executorFactory,
       SessionClient sessionClient,
       Clock clock,
+      Position initialReleasePosition,
       MetricRegistry metricRegistry,
       List<LabelValue> labelValues) {
     SessionPool pool =
@@ -2087,6 +2115,7 @@ static SessionPool createPool(
             executorFactory.get(),
             sessionClient,
             clock,
+            initialReleasePosition,
             metricRegistry,
             labelValues);
     pool.initPool();
@@ -2100,6 +2129,7 @@ private SessionPool(
       ScheduledExecutorService executor,
       SessionClient sessionClient,
       Clock clock,
+      Position initialReleasePosition,
       MetricRegistry metricRegistry,
       List<LabelValue> labelValues) {
     this.options = options;
@@ -2108,6 +2138,7 @@ private SessionPool(
     this.executor = executor;
     this.sessionClient = sessionClient;
     this.clock = clock;
+    this.initialReleasePosition = initialReleasePosition;
     this.poolMaintainer = new PoolMaintainer();
     this.initMetricsCollection(metricRegistry, labelValues);
     this.waitOnMinSessionsLatch =
@@ -2233,7 +2264,7 @@ private void handleException(SpannerException e, PooledSession session) {
     if (isSessionNotFound(e)) {
       invalidateSession(session);
     } else {
-      releaseSession(session, Position.FIRST);
+      releaseSession(session, false);
     }
   }
 
@@ -2396,18 +2427,35 @@ private void maybeCreateSession() {
       }
     }
   }
+
   /** Releases a session back to the pool. This might cause one of the waiters to be unblocked. */
-  private void releaseSession(PooledSession session, Position position) {
+  private void releaseSession(PooledSession session, boolean isNewSession) {
     Preconditions.checkNotNull(session);
     synchronized (lock) {
       if (closureFuture != null) {
         return;
       }
       if (waiters.size() == 0) {
-        // No pending waiters
-        switch (position) {
+        // There are no pending waiters.
+        // Add to a random position if the head of the session pool already contains many sessions
+        // with the same channel as this one.
+        if (session.releaseToPosition == Position.FIRST && isUnbalanced(session)) {
+          session.releaseToPosition = Position.RANDOM;
+        } else if (session.releaseToPosition == Position.RANDOM
+            && !isNewSession
+            && checkedOutSessions.size() <= 2) {
+          // Do not randomize if there are few other sessions checked out and this session has been
+          // used. This ensures that this session will be re-used for the next transaction, which is
+          // more efficient.
+          session.releaseToPosition = Position.FIRST;
+        }
+        switch (session.releaseToPosition) {
           case RANDOM:
             if (!sessions.isEmpty()) {
+              // A session should only be added at a random position the first time it is added to
+              // the pool. All following releases into the pool should happen at the front of the
+              // pool.
+              session.releaseToPosition = Position.FIRST;
               int pos = random.nextInt(sessions.size() + 1);
               sessions.add(pos, session);
               break;
@@ -2423,6 +2471,48 @@ private void releaseSession(PooledSession session, Position position) {
     }
   }
 
+  private boolean isUnbalanced(PooledSession session) {
+    // Don't bother with any randomization if the number of checked out sessions is low, as it is
+    // better to re-use sessions as much as possible in a low-QPS scenario.
+    if (sessions.isEmpty() || checkedOutSessions.size() <= 2) {
+      return false;
+    }
+    int numChannels = sessionClient.getSpanner().getOptions().getNumChannels();
+    if (numChannels == 1) {
+      return false;
+    }
+
+    // Ideally, the first numChannels sessions in the pool should contain exactly one session for
+    // each channel.
+    // Check if the first numChannels sessions at the head of the pool already contain more than 2
+    // sessions that use the same channel as this one. If so, we randomize.
+    int channel = session.getChannel();
+    int count = 0;
+    for (int i = 0; i < Math.min(numChannels, sessions.size()); i++) {
+      PooledSession otherSession = sessions.get(i);
+      if (channel == otherSession.getChannel()) {
+        count++;
+        if (count > 2) {
+          return true;
+        }
+      }
+    }
+    // Ideally, the use of a channel in the checked out sessions is exactly
+    // numCheckedOut / numChannels
+    // We check whether we are more than a factor two away from that perfect distribution.
+    // If we are, then we randomize.
+    int checkedOutThreshold = Math.max(2, 2 * checkedOutSessions.size() / numChannels);
+    for (PooledSessionFuture otherSession : checkedOutSessions) {
+      if (otherSession.isDone() && channel == otherSession.get().getChannel()) {
+        count++;
+        if (count > checkedOutThreshold) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
   private void handleCreateSessionsFailure(SpannerException e, int count) {
     synchronized (lock) {
       for (int i = 0; i < count; i++) {
@@ -2622,7 +2712,7 @@ public void onSessionReady(SessionImpl session) {
             // Release the session to a random position in the pool to prevent the case that a batch
             // of sessions that are affiliated with the same channel are all placed sequentially in
             // the pool.
-            releaseSession(pooledSession, Position.RANDOM);
+            releaseSession(pooledSession, true);
           }
         }
       }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java

@@ -26,16 +26,21 @@
 import com.google.api.core.ApiFutures;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
 import com.google.cloud.spanner.SessionPool.Clock;
+import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.protobuf.Empty;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import org.threeten.bp.Instant;
 
 abstract class BaseSessionPoolTest {
   ScheduledExecutorService mockExecutor;
   int sessionIndex;
+  AtomicLong channelHint = new AtomicLong(0L);
 
   final class TestExecutorFactory implements ExecutorFactory<ScheduledExecutorService> {
 
@@ -64,6 +69,9 @@ public void release(ScheduledExecutorService executor) {
   @SuppressWarnings("unchecked")
   SessionImpl mockSession() {
     final SessionImpl session = mock(SessionImpl.class);
+    Map options = new HashMap<>();
+    options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement());
+    when(session.getOptions()).thenReturn(options);
     when(session.getName())
         .thenReturn(
             "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java

@@ -26,6 +26,7 @@
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
 import com.google.cloud.spanner.SessionPool.PooledSession;
 import com.google.cloud.spanner.SessionPool.PooledSessionFuture;
+import com.google.cloud.spanner.SessionPool.Position;
 import com.google.cloud.spanner.SessionPool.SessionConsumerImpl;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -58,6 +59,7 @@ public void setUp() {
     initMocks(this);
     when(client.getOptions()).thenReturn(spannerOptions);
     when(client.getSessionClient(db)).thenReturn(sessionClient);
+    when(sessionClient.getSpanner()).thenReturn(client);
     when(spannerOptions.getNumChannels()).thenReturn(4);
     when(spannerOptions.getDatabaseRole()).thenReturn("role");
     setupMockSessionCreation();
@@ -111,9 +113,11 @@ private SessionImpl setupMockSession(final SessionImpl session) {
   }
 
   private SessionPool createPool() throws Exception {
+    // Allow sessions to be added to the head of the pool in all cases in this test, as it is
+    // otherwise impossible to know which session exactly is getting pinged at what point in time.
     SessionPool pool =
         SessionPool.createPool(
-            options, new TestExecutorFactory(), client.getSessionClient(db), clock);
+            options, new TestExecutorFactory(), client.getSessionClient(db), clock, Position.FIRST);
     pool.idleSessionRemovedListener =
         input -> {
           idledSessions.add(input);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java

@@ -25,6 +25,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
 import com.google.cloud.spanner.SessionPool.PooledSessionFuture;
+import com.google.cloud.spanner.SessionPool.Position;
 import com.google.cloud.spanner.SessionPool.SessionConsumerImpl;
 import com.google.cloud.spanner.SessionPoolOptions.ActionOnInactiveTransaction;
 import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions;
@@ -92,6 +93,7 @@ private void setupSpanner(DatabaseId db) {
     when(spannerOptions.getNumChannels()).thenReturn(4);
     when(spannerOptions.getDatabaseRole()).thenReturn("role");
     SessionClient sessionClient = mock(SessionClient.class);
+    when(sessionClient.getSpanner()).thenReturn(mockSpanner);
     when(mockSpanner.getSessionClient(db)).thenReturn(sessionClient);
     when(mockSpanner.getOptions()).thenReturn(spannerOptions);
     doAnswer(
@@ -226,7 +228,11 @@ public void stressTest() throws Exception {
     }
     pool =
         SessionPool.createPool(
-            builder.build(), new TestExecutorFactory(), mockSpanner.getSessionClient(db), clock);
+            builder.build(),
+            new TestExecutorFactory(),
+            mockSpanner.getSessionClient(db),
+            clock,
+            Position.RANDOM);
     pool.idleSessionRemovedListener =
         pooled -> {
           String name = pooled.getName();