feat: long running transaction clean up background task. Adding configuration options for closing inactive transactions.

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

@@ -51,6 +51,7 @@
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.Options.UpdateOption;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
+import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
 import com.google.cloud.spanner.SpannerImpl.ClosedException;
 import com.google.common.annotations.VisibleForTesting;
@@ -1279,7 +1280,7 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti
     @Override
     public long executePartitionedUpdate(Statement stmt, UpdateOption... options) {
       try {
-        return get().executePartitionedUpdate(stmt, options);
+        return get(true).executePartitionedUpdate(stmt, options);
       } finally {
         close();
       }
@@ -1332,6 +1333,10 @@ private PooledSession getOrNull() {
 
     @Override
     public PooledSession get() {
+      return get(false);
+    }
+
+    PooledSession get(final boolean eligibleForLongRunning) {
       if (inUse.compareAndSet(false, true)) {
         PooledSession res = null;
         try {
@@ -1346,6 +1351,7 @@ public PooledSession get() {
             incrementNumSessionsInUse();
             checkedOutSessions.add(this);
           }
+          res.eligibleForLongRunning = eligibleForLongRunning;
         }
         initialized.countDown();
       }
@@ -1366,6 +1372,28 @@ final class PooledSession implements Session {
     private volatile SpannerException lastException;
     private volatile boolean allowReplacing = true;
 
+    /**
+     * 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
+     * long-running. By default, most transaction types are not expected to be long-running and
+     * hence this value is false.
+     */
+    private volatile boolean eligibleForLongRunning = false;
+
+    /**
+     * Property to mark if the session is no longer part of the session pool. For ex - A session
+     * which is long-running gets cleaned up and removed from the pool.
+     */
+    private volatile boolean isRemovedFromPool = false;
+
+    /**
+     * Property to mark if a leaked session exception is already logged. Given a session maintainer
+     * thread runs repeatedly at a defined interval, this property allows us to ensure that an
+     * exception is logged only once per leaked session. This is to avoid noisy repeated logs around
+     * session leaks for long-running sessions.
+     */
+    private volatile boolean isLeakedExceptionLogged = false;
+
     @GuardedBy("lock")
     private SessionState state;
 
@@ -1385,6 +1413,11 @@ void setAllowReplacing(boolean allowReplacing) {
       this.allowReplacing = allowReplacing;
     }
 
+    @VisibleForTesting
+    void setEligibleForLongRunning(boolean eligibleForLongRunning) {
+      this.eligibleForLongRunning = eligibleForLongRunning;
+    }
+
     @Override
     public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
       return writeWithOptions(mutations).getCommitTimestamp();
@@ -1485,7 +1518,7 @@ public void close() {
         numSessionsInUse--;
         numSessionsReleased++;
       }
-      if (lastException != null && isSessionNotFound(lastException)) {
+      if ((lastException != null && isSessionNotFound(lastException)) || isRemovedFromPool) {
         invalidateSession(this);
       } else {
         if (lastException != null && isDatabaseOrInstanceNotFound(lastException)) {
@@ -1499,6 +1532,7 @@ public void close() {
           }
         }
         lastException = null;
+        isRemovedFromPool = false;
         if (state != SessionState.CLOSING) {
           state = SessionState.AVAILABLE;
         }
@@ -1651,6 +1685,10 @@ private PooledSession pollUninterruptiblyWithTimeout(long timeoutMillis) {
    *   <li>Keeps alive sessions that have not been used for a user configured time in order to keep
    *       MinSessions sessions alive in the pool at any time. The keep-alive traffic is smeared out
    *       over a window of 10 minutes to avoid bursty traffic.
+   *   <li>Removes unexpected long running transactions from the pool. Only certain transaction
+   *       types (for ex - Partitioned DML / Batch Reads) can be long running. This tasks checks the
+   *       sessions which have been inactive for a longer than usual duration (for ex - 60 minutes)
+   *       and removes such sessions from the pool.
    * </ul>
    */
   final class PoolMaintainer {
@@ -1659,16 +1697,24 @@ final class PoolMaintainer {
     private final Duration windowLength = Duration.ofMillis(TimeUnit.MINUTES.toMillis(10));
     // Frequency of the timer loop.
     @VisibleForTesting final long loopFrequency = options.getLoopFrequency();
-    // Number of loop iterations in which we need to to close all the sessions waiting for closure.
+    // Number of loop iterations in which we need to close all the sessions waiting for closure.
     @VisibleForTesting final long numClosureCycles = windowLength.toMillis() / loopFrequency;
     private final Duration keepAliveMillis =
         Duration.ofMillis(TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes()));
     // Number of loop iterations in which we need to keep alive all the sessions
     @VisibleForTesting final long numKeepAliveCycles = keepAliveMillis.toMillis() / loopFrequency;
 
-    Instant lastResetTime = Instant.ofEpochMilli(0);
-    int numSessionsToClose = 0;
-    int sessionsToClosePerLoop = 0;
+    /**
+     * Variable maintaining the last execution time of the long-running transaction cleanup task.
+     *
+     * <p>The long-running transaction cleanup needs to be performed every X minutes. The X minutes
+     * recurs multiple times within the invocation of the pool maintainer thread. For ex - If the
+     * main thread runs every 10s and the long-running transaction clean-up needs to be performed
+     * every 2 minutes, then we need to keep a track of when was the last time that this task
+     * executed and makes sure we only execute it every 2 minutes and not every 10 seconds.
+     */
+    @VisibleForTesting Instant lastExecutionTime;
+
     boolean closed = false;
 
     @GuardedBy("lock")
@@ -1678,6 +1724,7 @@ final class PoolMaintainer {
     boolean running;
 
     void init() {
+      lastExecutionTime = clock.instant();
       // Scheduled pool maintenance worker.
       synchronized (lock) {
         scheduledFuture =
@@ -1723,6 +1770,7 @@ void maintainPool() {
           decrementPendingClosures(1);
         }
       }
+      removeLongRunningSessions(currTime);
     }
 
     private void removeIdleSessions(Instant currTime) {
@@ -1736,7 +1784,13 @@ private void removeIdleSessions(Instant currTime) {
           PooledSession session = iterator.next();
           if (session.lastUseTime.isBefore(minLastUseTime)) {
             if (session.state != SessionState.CLOSING) {
-              removeFromPool(session);
+              boolean isRemoved = removeFromPool(session);
+              if (isRemoved) {
+                numIdleSessionsRemoved++;
+                if (idleSessionRemovedListener != null) {
+                  idleSessionRemovedListener.apply(session);
+                }
+              }
               iterator.remove();
             }
           }
@@ -1792,6 +1846,87 @@ private void replenishPool() {
         }
       }
     }
+
+    // cleans up sessions which are unexpectedly long-running.
+    void removeLongRunningSessions(Instant currentTime) {
+      try {
+        if (SessionPool.this.isClosed()) {
+          return;
+        }
+        final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions =
+            options.getInactiveTransactionRemovalOptions();
+        final Instant minExecutionTime =
+            lastExecutionTime.plus(inactiveTransactionRemovalOptions.getExecutionFrequency());
+        if (currentTime.isBefore(minExecutionTime)) {
+          return;
+        }
+        lastExecutionTime = currentTime; // update this only after we have decided to execute task
+        if (options.closeInactiveTransactions()
+            || options.warnInactiveTransactions()
+            || options.warnAndCloseInactiveTransactions()) {
+          removeLongRunningSessions(currentTime, inactiveTransactionRemovalOptions);
+        }
+      } catch (final Throwable t) {
+        logger.log(Level.WARNING, "Failed removing long running transactions", t);
+      }
+    }
+
+    private void removeLongRunningSessions(
+        final Instant currentTime,
+        final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions) {
+      synchronized (lock) {
+        final double usedSessionsRatio = getRatioOfSessionsInUse();
+        if (usedSessionsRatio > inactiveTransactionRemovalOptions.getUsedSessionsRatioThreshold()) {
+          Iterator<PooledSessionFuture> iterator = checkedOutSessions.iterator();
+          while (iterator.hasNext()) {
+            final PooledSessionFuture sessionFuture = iterator.next();
+            // the below get() call on future object is non-blocking since checkedOutSessions
+            // collection is populated only when the get() method in {@code PooledSessionFuture} is
+            // called.
+            final PooledSession session = sessionFuture.get();
+            final Duration durationFromLastUse = Duration.between(session.lastUseTime, currentTime);
+            if (!session.eligibleForLongRunning
+                && durationFromLastUse.compareTo(
+                        inactiveTransactionRemovalOptions.getIdleTimeThreshold())
+                    > 0) {
+              if ((options.warnInactiveTransactions() || options.warnAndCloseInactiveTransactions())
+                  && !session.isLeakedExceptionLogged) {
+                if (options.warnAndCloseInactiveTransactions()) {
+                  logger.log(
+                      Level.WARNING,
+                      String.format("Removing long-running session => %s", session.getName()),
+                      sessionFuture.leakedException);
+                  session.isLeakedExceptionLogged = true;
+                } else if (options.warnInactiveTransactions()) {
+                  logger.log(
+                      Level.WARNING,
+                      String.format(
+                          "Detected long-running session => %s. To automatically remove "
+                              + "long-running sessions, set SessionOption ActionOnInactiveTransaction "
+                              + "to WARN_AND_CLOSE by invoking setWarnAndCloseIfInactiveTransactions() method.",
+                          session.getName()),
+                      sessionFuture.leakedException);
+                  session.isLeakedExceptionLogged = true;
+                }
+              }
+              if ((options.closeInactiveTransactions()
+                      || options.warnAndCloseInactiveTransactions())
+                  && session.state != SessionState.CLOSING) {
+                final boolean isRemoved = removeFromPool(session);
+                if (isRemoved) {
+                  session.isRemovedFromPool = true;
+                  numLeakedSessionsRemoved++;
+                  if (longRunningSessionRemovedListener != null) {
+                    longRunningSessionRemovedListener.apply(session);
+                  }
+                }
+                iterator.remove();
+              }
+            }
+          }
+        }
+      }
+    }
   }
 
   private enum Position {
@@ -1872,6 +2007,9 @@ private enum Position {
   @GuardedBy("lock")
   private long numIdleSessionsRemoved = 0;
 
+  @GuardedBy("lock")
+  private long numLeakedSessionsRemoved = 0;
+
   private AtomicLong numWaiterTimeouts = new AtomicLong();
 
   @GuardedBy("lock")
@@ -1885,6 +2023,8 @@ private enum Position {
 
   @VisibleForTesting Function<PooledSession, Void> idleSessionRemovedListener;
 
+  @VisibleForTesting Function<PooledSession, Void> longRunningSessionRemovedListener;
+
   private final CountDownLatch waitOnMinSessionsLatch;
 
   /**
@@ -1895,12 +2035,16 @@ private enum Position {
    */
   static SessionPool createPool(
       SpannerOptions spannerOptions, SessionClient sessionClient, List<LabelValue> labelValues) {
+    final SessionPoolOptions sessionPoolOptions = spannerOptions.getSessionPoolOptions();
+
+    // A clock instance is passed in {@code SessionPoolOptions} in order to allow mocking via tests.
+    final Clock poolMaintainerClock = sessionPoolOptions.getPoolMaintainerClock();
     return createPool(
-        spannerOptions.getSessionPoolOptions(),
+        sessionPoolOptions,
         spannerOptions.getDatabaseRole(),
         ((GrpcTransportOptions) spannerOptions.getTransportOptions()).getExecutorFactory(),
         sessionClient,
-        new Clock(),
+        poolMaintainerClock == null ? new Clock() : poolMaintainerClock,
         Metrics.getMetricRegistry(),
         labelValues);
   }
@@ -2015,18 +2159,26 @@ int getNumberOfSessionsInUse() {
     }
   }
 
-  void removeFromPool(PooledSession session) {
+  @VisibleForTesting
+  double getRatioOfSessionsInUse() {
+    synchronized (lock) {
+      final int maxSessions = options.getMaxSessions();
+      if (maxSessions == 0) {
+        return 0;
+      }
+      return (double) numSessionsInUse / maxSessions;
+    }
+  }
+
+  boolean removeFromPool(PooledSession session) {
     synchronized (lock) {
       if (isClosed()) {
         decrementPendingClosures(1);
-        return;
+        return false;
       }
       session.markClosing();
       allSessions.remove(session);
-      numIdleSessionsRemoved++;
-    }
-    if (idleSessionRemovedListener != null) {
-      idleSessionRemovedListener.apply(session);
+      return true;
     }
   }
 
@@ -2036,6 +2188,13 @@ long numIdleSessionsRemoved() {
     }
   }
 
+  @VisibleForTesting
+  long numLeakedSessionsRemoved() {
+    synchronized (lock) {
+      return numLeakedSessionsRemoved;
+    }
+  }
+
   @VisibleForTesting
   int getNumberOfSessionsInPool() {
     synchronized (lock) {