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;
@@ -1365,6 +1366,7 @@ final class PooledSession implements Session {
     private volatile Instant lastUseTime;
     private volatile SpannerException lastException;
     private volatile boolean allowReplacing = true;
+    private volatile boolean isLongRunning = false;
 
     @GuardedBy("lock")
     private SessionState state;
@@ -1385,6 +1387,11 @@ void setAllowReplacing(boolean allowReplacing) {
       this.allowReplacing = allowReplacing;
     }
 
+    @VisibleForTesting
+    void setIsLongRunning(boolean isLongRunning) {
+      this.isLongRunning = isLongRunning;
+    }
+
     @Override
     public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
       return writeWithOptions(mutations).getCommitTimestamp();
@@ -1422,6 +1429,7 @@ public long executePartitionedUpdate(Statement stmt, UpdateOption... options)
         throws SpannerException {
       try {
         markUsed();
+        markLongRunning();
         return delegate.executePartitionedUpdate(stmt, options);
       } catch (SpannerException e) {
         throw lastException = e;
@@ -1572,6 +1580,14 @@ void markUsed() {
       lastUseTime = clock.instant();
     }
 
+    /**
+     * Method to mark a session occupied by a long-running transaction. Any transaction that is
+     * expected to be long-running (for ex - Partitioned DML, Batch Read) must use this method.
+     */
+    void markLongRunning() {
+      isLongRunning = true;
+    }
+
     @Override
     public TransactionManager transactionManager(TransactionOption... options) {
       return delegate.transactionManager(options);
@@ -1651,6 +1667,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 executing for a longer than usual duration (60 minutes) and
+   *       returns such sessions back to the pool.
    * </ul>
    */
   final class PoolMaintainer {
@@ -1659,16 +1679,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 public volatile Instant lastExecutionTime;
+
     boolean closed = false;
 
     @GuardedBy("lock")
@@ -1678,6 +1706,7 @@ final class PoolMaintainer {
     boolean running;
 
     void init() {
+      lastExecutionTime = clock.instant();
       // Scheduled pool maintenance worker.
       synchronized (lock) {
         scheduledFuture =
@@ -1723,6 +1752,7 @@ void maintainPool() {
           decrementPendingClosures(1);
         }
       }
+      closeLongRunningTransactions(currTime);
     }
 
     private void removeIdleSessions(Instant currTime) {
@@ -1736,7 +1766,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 +1828,60 @@ private void replenishPool() {
         }
       }
     }
+
+    // cleans up transactions which are unexpectedly long-running.
+    void closeLongRunningTransactions(Instant currentTime) {
+      try {
+        synchronized (lock) {
+          if (SessionPool.this.isClosed()) {
+            return;
+          }
+          final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions =
+              options.getInactiveTransactionRemovalOptions();
+          final Instant minExecutionTime =
+              lastExecutionTime.plus(inactiveTransactionRemovalOptions.getInterval());
+          if (currentTime.isBefore(minExecutionTime)) {
+            return;
+          }
+          lastExecutionTime = currentTime; // update this only after we have decided to execute task
+          if (options.closeInactiveTransactions() || options.warnInactiveTransactions()) {
+            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.isLongRunning
+                && durationFromLastUse.toMillis()
+                    > inactiveTransactionRemovalOptions.getExecutionTimeThreshold().toMillis()) {
+              logger.log(
+                  Level.WARNING, "Removing long running session", sessionFuture.leakedException);
+              numInactiveSessionsRemoved++;
+              if (options.closeInactiveTransactions() && session.state != SessionState.CLOSING) {
+                removeFromPool(session);
+                iterator.remove();
+              }
+            }
+          }
+        }
+      }
+    }
   }
 
   private enum Position {
@@ -1872,6 +1962,9 @@ private enum Position {
   @GuardedBy("lock")
   private long numIdleSessionsRemoved = 0;
 
+  @GuardedBy("lock")
+  private long numInactiveSessionsRemoved = 0;
+
   private AtomicLong numWaiterTimeouts = new AtomicLong();
 
   @GuardedBy("lock")
@@ -2015,18 +2108,24 @@ 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 +2135,13 @@ long numIdleSessionsRemoved() {
     }
   }
 
+  @VisibleForTesting
+  long numInactiveSessionsRemoved() {
+    synchronized (lock) {
+      return numInactiveSessionsRemoved;
+    }
+  }
+
   @VisibleForTesting
   int getNumberOfSessionsInPool() {
     synchronized (lock) {