googleapis · arpan14 · Apr 9, 2024 · Mar 27, 2024 · Mar 28, 2024 · Apr 1, 2024
@@ -151,8 +151,8 @@ public void run() {
   }
 
   /**
-   * Callback interface to be used for BatchCreateSessions. When sessions become available or
-   * session creation fails, one of the callback methods will be called.
+   * Callback interface to be used for Sessions. When sessions become available or session creation
+   * fails, one of the callback methods will be called.
    */
   interface SessionConsumer {
     /** Called when a session has been created and is ready for use. */

@@ -104,6 +104,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -146,6 +147,7 @@ void maybeWaitOnMinSessions() {
   }
 
   private abstract static class CachedResultSetSupplier implements Supplier<ResultSet> {
+
     private ResultSet cached;
 
     abstract ResultSet load();
@@ -1155,6 +1157,46 @@ private PooledSessionFuture createPooledSessionFuture(
     return new PooledSessionFuture(future, span);
   }
 
+  /** Wrapper class for the {@link SessionFuture} implementations. */
+  interface SessionFutureWrapper<T extends SessionFuture> {
+
+    /** Method to resolve {@link SessionFuture} implementation for different use-cases. */
+    T get();
+  }
+
+  class PooledSessionFutureWrapper implements SessionFutureWrapper<PooledSessionFuture> {
+    PooledSessionFuture pooledSessionFuture;
+
+    public PooledSessionFutureWrapper(PooledSessionFuture pooledSessionFuture) {
+      this.pooledSessionFuture = pooledSessionFuture;
+    }
+
+    @Override
+    public PooledSessionFuture get() {
+      return this.pooledSessionFuture;
+    }
+  }
+
+  class MultiplexedSessionFutureWrapper implements SessionFutureWrapper<MultiplexedSessionFuture> {
+    SettableApiFuture<MultiplexedSessionFuture> multiplexedSessionSettableApiFuture;
+
+    public MultiplexedSessionFutureWrapper(
+        SettableApiFuture<MultiplexedSessionFuture> multiplexedSessionSettableApiFuture) {
+      this.multiplexedSessionSettableApiFuture = multiplexedSessionSettableApiFuture;
+    }
+
+    @Override
+    public MultiplexedSessionFuture get() {
+      try {
+        return this.multiplexedSessionSettableApiFuture.get();
+      } catch (InterruptedException interruptedException) {
+        throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
+      } catch (ExecutionException executionException) {
+        throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
+      }
+    }
+  }
+
   interface SessionFuture extends Session {
 
     /**
@@ -1435,12 +1477,9 @@ PooledSession get(final boolean eligibleForLongRunning) {
 
   class MultiplexedSessionFuture extends SimpleForwardingListenableFuture<MultiplexedSession>
       implements SessionFuture {
-    private final ISpan span;
-
     @VisibleForTesting
-    MultiplexedSessionFuture(ListenableFuture<MultiplexedSession> delegate, ISpan span) {
+    MultiplexedSessionFuture(ListenableFuture<MultiplexedSession> delegate) {
       super(delegate);
-      this.span = span;
     }
 
     @Override
@@ -1645,15 +1684,6 @@ private MultiplexedSession getOrNull() {
 
     @Override
     public MultiplexedSession get() {
-      MultiplexedSession res = null;
-      try {
-        res = super.get();
-      } catch (Throwable e) {
-        // ignore the exception as it will be handled by the call to super.get() below.
-      }
-      if (res != null) {
-        res.markBusy(span);
-      }
       try {
         return super.get();
       } catch (ExecutionException e) {
@@ -2231,6 +2261,9 @@ private PooledSession pollUninterruptiblyWithTimeout(
    */
   final class PoolMaintainer {
 
+    // Delay post which the maintainer will retry creating/replacing the current multiplexed session
+    private final Duration multiplexedSessionCreationRetryDelay = Duration.ofMinutes(10);
+
     // Length of the window in millis over which we keep track of maximum number of concurrent
     // sessions in use.
     private final Duration windowLength = Duration.ofMillis(TimeUnit.MINUTES.toMillis(10));
@@ -2254,6 +2287,8 @@ final class PoolMaintainer {
      */
     @VisibleForTesting Instant lastExecutionTime;
 
+    @VisibleForTesting Instant multiplexedSessionReplacementAttemptTime;
+
     /**
      * The previous numSessionsAcquired seen by the maintainer. This is used to calculate the
      * transactions per second, which again is used to determine whether to randomize the order of
@@ -2271,6 +2306,8 @@ final class PoolMaintainer {
 
     void init() {
       lastExecutionTime = clock.instant();
+      multiplexedSessionReplacementAttemptTime = clock.instant();
+
       // Scheduled pool maintenance worker.
       synchronized (lock) {
         scheduledFuture =
@@ -2312,6 +2349,7 @@ void maintainPool() {
         this.prevNumSessionsAcquired = SessionPool.this.numSessionsAcquired;
       }
       Instant currTime = clock.instant();
+      maintainMultiplexedSession(currTime);
       removeIdleSessions(currTime);
       // Now go over all the remaining sessions and see if they need to be kept alive explicitly.
       keepAliveSessions(currTime);
@@ -2480,6 +2518,46 @@ private void removeLongRunningSessions(
         }
       }
     }
+
+    void maintainMultiplexedSession(Instant currentTime) {
+      try {
+        if (options.getUseMultiplexedSession()) {
+          synchronized (lock) {
+            if (getMultiplexedSession().isDone()
+                && getMultiplexedSession().get() != null
+                && isMultiplexedSessionStale(currentTime)) {
+              final Instant minExecutionTime =
+                  multiplexedSessionReplacementAttemptTime.plus(
+                      multiplexedSessionCreationRetryDelay);
+              if (currentTime.isBefore(minExecutionTime)) {
+                return;
+              }
+
+              /*
+               This will attempt to create a new multiplexed session. if successfully created then
+               the existing session will be replaced. Note that there maybe active transactions
+               running on the stale session. Hence, it is important that we only replace the reference
+               and not invoke a DeleteSession RPC.
+              */
+              maybeCreateMultiplexedSession(multiplexedMaintainerConsumer);
+
+              // update this only after we have attempted to replace the multiplexed session
+              multiplexedSessionReplacementAttemptTime = currentTime;
+            }
+          }
+        }
+      } catch (final Throwable t) {
+        logger.log(Level.WARNING, "Failed to maintain multiplexed session", t);
+      }
+    }
+
+    boolean isMultiplexedSessionStale(Instant currentTime) {
+      final CachedSession session = getMultiplexedSession().get();
+      final Duration durationFromCreationTime =
+          Duration.between(session.getDelegate().getCreateTime(), currentTime);
+      return durationFromCreationTime.compareTo(options.getMultiplexedSessionMaintenanceDuration())
+          > 0;
+    }
   }
 
   enum Position {
@@ -2556,6 +2634,9 @@ enum Position {
   @GuardedBy("lock")
   private int numSessionsBeingCreated = 0;
 
+  @GuardedBy("lock")
+  private boolean multiplexedSessionBeingCreated = false;
+
   @GuardedBy("lock")
   private int numSessionsInUse = 0;
 
@@ -2585,6 +2666,10 @@ enum Position {
 
   private AtomicLong numWaiterTimeouts = new AtomicLong();
 
+  private final AtomicReference<SettableApiFuture<MultiplexedSessionFuture>>
+      currentMultiplexedSessionReference = new AtomicReference<>(SettableApiFuture.create());
+  MultiplexedSessionFutureWrapper wrappedMultiplexedSessionFuture = null;
+
   @GuardedBy("lock")
   private final Set<PooledSession> allSessions = new HashSet<>();
 
@@ -2593,9 +2678,16 @@ enum Position {
   final Set<PooledSessionFuture> checkedOutSessions = new HashSet<>();
 
   private final SessionConsumer sessionConsumer = new SessionConsumerImpl();
+
+  private final MultiplexedSessionInitializationConsumer multiplexedSessionInitializationConsumer =
+      new MultiplexedSessionInitializationConsumer();
+  private final MultiplexedSessionMaintainerConsumer multiplexedMaintainerConsumer =
+      new MultiplexedSessionMaintainerConsumer();
+
   @VisibleForTesting Function<PooledSession, Void> idleSessionRemovedListener;
 
   @VisibleForTesting Function<PooledSession, Void> longRunningSessionRemovedListener;
+  @VisibleForTesting Function<MultiplexedSession, Void> multiplexedSessionRemovedListener;
   private final CountDownLatch waitOnMinSessionsLatch;
   private final SessionReplacementHandler pooledSessionReplacementHandler =
       new PooledSessionReplacementHandler();
@@ -2839,6 +2931,9 @@ private void initPool() {
       if (options.getMinSessions() > 0) {
         createSessions(options.getMinSessions(), true);
       }
+      if (options.getUseMultiplexedSession()) {
+        maybeCreateMultiplexedSession(multiplexedSessionInitializationConsumer);
+      }
     }
   }
 
@@ -2900,6 +2995,38 @@ boolean isValid() {
     }
   }
 
+  /**
+   * Returns a multiplexed session. The method fallbacks to a regular session if {@link
+   * SessionPoolOptions#useMultiplexedSession} is not set.
+   */
+  SessionFutureWrapper getMultiplexedSessionWithFallback() throws SpannerException {
+    if (options.getUseMultiplexedSession()) {
+      try {
+        SessionFutureWrapper sessionFuture = getWrappedMultiplexedSessionFuture();
+        incrementNumSessionsInUse(true);
+        return sessionFuture;
+      } catch (Throwable t) {
+        ISpan span = tracer.getCurrentSpan();
+        span.addAnnotation("No multiplexed session available.");
+        throw SpannerExceptionFactory.asSpannerException(t.getCause());
+      }
+    } else {
+      return new PooledSessionFutureWrapper(getSession());
+    }
+  }
+
+  SessionFutureWrapper getWrappedMultiplexedSessionFuture() {
+    return wrappedMultiplexedSessionFuture;
+  }
+
+  /**
+   * This method is a blocking method. It will block until the underlying {@code
+   * SettableApiFuture<MultiplexedSessionFuture>} is resolved.
+   */
+  MultiplexedSessionFuture getMultiplexedSession() {
+    return (MultiplexedSessionFuture) getWrappedMultiplexedSessionFuture().get();
+  }
+
   /**
    * Returns a session to be used for requests to spanner. This method is always non-blocking and
    * returns a {@link PooledSessionFuture}. In case the pool is exhausted and {@link
@@ -3303,6 +3430,20 @@ private boolean canCreateSession() {
     }
   }
 
+  private void maybeCreateMultiplexedSession(SessionConsumer sessionConsumer) {
+    synchronized (lock) {
+      if (!multiplexedSessionBeingCreated) {
+        logger.log(Level.FINE, String.format("Creating multiplexed sessions"));
+        try {
+          multiplexedSessionBeingCreated = true;
+          sessionClient.createMultiplexedSession(sessionConsumer);
+        } catch (Throwable ignore) {
+          // such an exception will never be thrown. the exception will be passed onto the consumer.
+        }
+      }
+    }
+  }
+
   private void createSessions(final int sessionCount, boolean distributeOverChannels) {
     logger.log(Level.FINE, String.format("Creating %d sessions", sessionCount));
     synchronized (lock) {
@@ -3325,6 +3466,103 @@ private void createSessions(final int sessionCount, boolean distributeOverChanne
     }
   }
 
+  /**
+   * Callback interface which is invoked when a multiplexed session is being replaced by the
+   * background maintenance thread. When a multiplexed session creation fails during background
+   * thread, it would simply log the exception and retry the session creation in the next background
+   * thread invocation.
+   *
+   * <p>This consumer is not used when the multiplexed session is getting initialized for the first
+   * time during application startup. We instead use {@link
+   * MultiplexedSessionInitializationConsumer} for the first time when multiplexed session is
+   * getting created.
+   */
+  class MultiplexedSessionMaintainerConsumer implements SessionConsumer {
+    @Override
+    public void onSessionReady(SessionImpl sessionImpl) {
+      final SettableFuture<MultiplexedSession> settableFuture = SettableFuture.create();
+      final MultiplexedSession newSession = new MultiplexedSession(sessionImpl);
+      settableFuture.set(newSession);
+
+      synchronized (lock) {
+        MultiplexedSession oldSession = null;
+        if (currentMultiplexedSessionReference.get().isDone()) {
+          oldSession = getMultiplexedSession().get();
+        }
+        SettableApiFuture<MultiplexedSessionFuture> settableApiFuture = SettableApiFuture.create();
+        settableApiFuture.set(new MultiplexedSessionFuture(settableFuture));
+        currentMultiplexedSessionReference.set(settableApiFuture);
+        wrappedMultiplexedSessionFuture = new MultiplexedSessionFutureWrapper(settableApiFuture);
+        if (oldSession != null) {
+          logger.log(
+              Level.INFO,
+              String.format(
+                  "Removed Multiplexed Session => %s created at => %s and",
+                  oldSession.getName(), oldSession.getDelegate().getCreateTime()));
+          if (multiplexedSessionRemovedListener != null) {
+            multiplexedSessionRemovedListener.apply(oldSession);
+          }
+        }
+        multiplexedSessionBeingCreated = false;
+      }
+    }
+
+    /**
+     * Method which logs the exception so that session creation can be re-attempted in the next
+     * background thread invocation.
+     */
+    @Override
+    public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
+      synchronized (lock) {
+        multiplexedSessionBeingCreated = false;
+        wrappedMultiplexedSessionFuture =
+            new MultiplexedSessionFutureWrapper(currentMultiplexedSessionReference.get());
+      }
+      logger.log(
+          Level.WARNING,
+          String.format(
+              "Failed to create multiplexed session. "
+                  + "Pending replacing stale multiplexed session",
+              t));
+    }
+  }
+
+  /**
+   * Callback interface which is invoked when a multiplexed session is getting initialised for the
+   * first time when a session is getting created.
+   */
+  class MultiplexedSessionInitializationConsumer implements SessionConsumer {
+    @Override
+    public void onSessionReady(SessionImpl sessionImpl) {
+      final SettableFuture<MultiplexedSession> settableFuture = SettableFuture.create();
+      final MultiplexedSession newSession = new MultiplexedSession(sessionImpl);
+      settableFuture.set(newSession);
+
+      synchronized (lock) {
+        SettableApiFuture<MultiplexedSessionFuture> settableApiFuture =
+            currentMultiplexedSessionReference.get();
+        settableApiFuture.set(new MultiplexedSessionFuture(settableFuture));
+        wrappedMultiplexedSessionFuture = new MultiplexedSessionFutureWrapper(settableApiFuture);
+        multiplexedSessionBeingCreated = false;
+      }
+    }
+
+    /**
+     * When a multiplexed session fails during initialization we would like all pending threads to
+     * receive the exception and throw the error. This is done because at the time of start up there
+     * is no other multiplexed session which could have been assigned to the pending requests.
+     */
+    @Override
+    public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) {
+      synchronized (lock) {
+        multiplexedSessionBeingCreated = false;
+        wrappedMultiplexedSessionFuture =
+            new MultiplexedSessionFutureWrapper(currentMultiplexedSessionReference.get());
+        currentMultiplexedSessionReference.get().setException(newSpannerException(t));
+      }
+    }
+  }
+
   /**
    * {@link SessionConsumer} that receives the created sessions from a {@link SessionClient} and
    * releases these into the pool. The session pool only needs one instance of this, as all sessions