fix(spanner): prevent memory leak and thread blocking in transaction keep-alive

- Enable `setRemoveOnCancelPolicy(true)` on `KEEP_ALIVE_SERVICE` so canceled tasks are immediately purged from `DelayedWorkQueue`.
- Use a `WeakReference<ReadWriteTransaction>` in `KeepAliveRunnable` to prevent scheduled tasks from retaining strong references to transaction instances.
- Use `abortedLock.tryLock()` in `KeepAliveRunnable` so the shared executor thread does not block when a transaction is active or retrying.
- Remove duplicate `maybeScheduleKeepAlivePing` listener registration on keep-alive query completion.
- Add unit tests in `ReadWriteTransactionTest` verifying task removal on cancel, weak reference retention, non-blocking lock handling, and single ping scheduling on completion.

Author: olavloite

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java

@@ -65,15 +65,15 @@
 import io.grpc.Deadline;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.context.Scope;
+import java.lang.ref.WeakReference;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Callable;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
@@ -100,8 +100,20 @@ class ReadWriteTransaction extends AbstractMultiUseTransaction {
   private static final ThreadFactory KEEP_ALIVE_THREAD_FACTORY =
       ThreadFactoryUtil.createVirtualOrPlatformDaemonThreadFactory(
           "read-write-transaction-keep-alive", true);
-  private static final ScheduledExecutorService KEEP_ALIVE_SERVICE =
-      Executors.newSingleThreadScheduledExecutor(KEEP_ALIVE_THREAD_FACTORY);
+  private static final ScheduledThreadPoolExecutor KEEP_ALIVE_SERVICE = createKeepAliveService();
+
+  private static ScheduledThreadPoolExecutor createKeepAliveService() {
+    ScheduledThreadPoolExecutor executor =
+        new ScheduledThreadPoolExecutor(1, KEEP_ALIVE_THREAD_FACTORY);
+    executor.setRemoveOnCancelPolicy(true);
+    return executor;
+  }
+
+  @VisibleForTesting
+  static ScheduledThreadPoolExecutor getKeepAliveService() {
+    return KEEP_ALIVE_SERVICE;
+  }
+
   private static final ParsedStatement SELECT1_STATEMENT =
       AbstractStatementParser.getInstance(Dialect.GOOGLE_STANDARD_SQL)
           .parse(Statement.of("SELECT 1"));
@@ -146,7 +158,7 @@ class ReadWriteTransaction extends AbstractMultiUseTransaction {
   private Savepoint autoSavepoint;
 
   private final int maxInternalRetries;
-  private final ReentrantLock abortedLock = new ReentrantLock();
+  @VisibleForTesting final ReentrantLock abortedLock = new ReentrantLock();
   private final long transactionId;
   private final DatabaseClient dbClient;
   private final TransactionOption[] transactionOptions;
@@ -468,14 +480,15 @@ private boolean shouldPing() {
         && rolledBackToSavepointException == null;
   }
 
-  private void maybeScheduleKeepAlivePing() {
+  @VisibleForTesting
+  void maybeScheduleKeepAlivePing() {
     if (shouldPing()) {
       keepAliveLock.lock();
       try {
-        if (keepAliveFuture == null || keepAliveFuture.isDone()) {
+        if (shouldPing() && (keepAliveFuture == null || keepAliveFuture.isDone())) {
           keepAliveFuture =
               KEEP_ALIVE_SERVICE.schedule(
-                  new KeepAliveRunnable(),
+                  new KeepAliveRunnable(this),
                   keepAliveIntervalMillis > 0
                       ? keepAliveIntervalMillis
                       : DEFAULT_KEEP_ALIVE_INTERVAL_MILLIS,
@@ -487,36 +500,76 @@ private void maybeScheduleKeepAlivePing() {
     }
   }
 
+  @VisibleForTesting
+  ScheduledFuture<?> getKeepAliveFuture() {
+    return keepAliveFuture;
+  }
+
   private void cancelScheduledKeepAlivePing() {
     if (keepAliveLock != null) {
       keepAliveLock.lock();
       try {
         if (keepAliveFuture != null) {
           keepAliveFuture.cancel(false);
+          keepAliveFuture = null;
         }
       } finally {
         keepAliveLock.unlock();
       }
     }
   }
 
-  private class KeepAliveRunnable implements Runnable {
+  private void rescheduleKeepAlivePing() {
+    if (keepAliveLock != null) {
+      keepAliveLock.lock();
+      try {
+        keepAliveFuture = null;
+        maybeScheduleKeepAlivePing();
+      } finally {
+        keepAliveLock.unlock();
+      }
+    }
+  }
+
+  @VisibleForTesting
+  static class KeepAliveRunnable implements Runnable {
+    final WeakReference<ReadWriteTransaction> transactionRef;
+
+    KeepAliveRunnable(ReadWriteTransaction transaction) {
+      this.transactionRef = new WeakReference<>(transaction);
+    }
+
     @Override
     public void run() {
-      if (shouldPing()) {
-        // Do a shoot-and-forget ping and schedule a new ping over 8 seconds after this ping has
-        // finished.
-        ApiFuture<ResultSet> future =
-            executeQueryAsync(
+      ReadWriteTransaction transaction = transactionRef.get();
+      if (transaction != null && transaction.shouldPing()) {
+        if (transaction.abortedLock.tryLock()) {
+          boolean schedulePing = false;
+          try {
+            // Do a shoot-and-forget ping.
+            // Note: executeQueryAsync automatically adds StatementResultCallback,
+            // which calls maybeScheduleKeepAlivePing() upon completion.
+            transaction.executeQueryAsync(
                 CallType.SYNC,
                 SELECT1_STATEMENT,
                 AnalyzeMode.NONE,
                 Options.tag(
                     System.getProperty(
                         "spanner.connection.keep_alive_query_tag",
                         "connection.transaction-keep-alive")));
-        future.addListener(
-            ReadWriteTransaction.this::maybeScheduleKeepAlivePing, MoreExecutors.directExecutor());
+          } catch (Throwable t) {
+            schedulePing = true;
+          } finally {
+            transaction.abortedLock.unlock();
+          }
+          if (schedulePing) {
+            transaction.maybeScheduleKeepAlivePing();
+          }
+        } else {
+          // Transaction is currently busy (executing a statement or retrying).
+          // Reschedule keep-alive ping for later since it is active.
+          transaction.rescheduleKeepAlivePing();
+        }
       }
     }
   }

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java

@@ -24,8 +24,12 @@
 import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doThrow;
@@ -67,6 +71,8 @@
 import java.math.BigDecimal;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ScheduledFuture;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -158,12 +164,21 @@ private ReadWriteTransaction createSubject() {
     return createSubject(CommitBehavior.SUCCEED, false);
   }
 
+  private ReadWriteTransaction createSubject(boolean keepTransactionAlive) {
+    return createSubject(CommitBehavior.SUCCEED, false, keepTransactionAlive);
+  }
+
   private ReadWriteTransaction createSubject(CommitBehavior commitBehavior) {
-    return createSubject(commitBehavior, false);
+    return createSubject(commitBehavior, false, false);
   }
 
   private ReadWriteTransaction createSubject(
       final CommitBehavior commitBehavior, boolean withRetry) {
+    return createSubject(commitBehavior, withRetry, false);
+  }
+
+  private ReadWriteTransaction createSubject(
+      final CommitBehavior commitBehavior, boolean withRetry, boolean keepTransactionAlive) {
     DatabaseClient client = mock(DatabaseClient.class);
     when(client.transactionManager())
         .thenAnswer(
@@ -179,6 +194,7 @@ private ReadWriteTransaction createSubject(
             });
     return ReadWriteTransaction.newBuilder()
         .setDatabaseClient(client)
+        .setKeepTransactionAlive(keepTransactionAlive)
         .setRetryAbortsInternally(withRetry)
         .setIsolationLevel(IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED)
         .setSavepointSupport(SavepointSupport.FAIL_AFTER_ROLLBACK)
@@ -857,6 +873,172 @@ public void testGetCommitResponseAfterCommit() {
     assertNotNull(transaction.getCommitResponseOrNull());
   }
 
+  @Test
+  public void testKeepAliveTaskRemovedFromQueueOnCancel() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.UPDATE);
+    when(parsedStatement.isUpdate()).thenReturn(true);
+    Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
+    get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement));
+
+    ScheduledFuture<?> future = waitForKeepAliveFuture(transaction);
+    assertNotNull(future);
+    assertTrue(ReadWriteTransaction.getKeepAliveService().getQueue().contains(future));
+
+    get(transaction.commitAsync(CallType.SYNC, NoopEndTransactionCallback.INSTANCE));
+    assertFalse(ReadWriteTransaction.getKeepAliveService().getQueue().contains(future));
+  }
+
+  @Test
+  public void testKeepAliveWeakReference() {
+    ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
+    ReadWriteTransaction.KeepAliveRunnable runnable =
+        new ReadWriteTransaction.KeepAliveRunnable(transaction);
+
+    assertNotNull(runnable.transactionRef);
+    assertSame(transaction, runnable.transactionRef.get());
+  }
+
+  @Test
+  public void testKeepAliveRescheduledWhenLockBusy() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.UPDATE);
+    when(parsedStatement.isUpdate()).thenReturn(true);
+    Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
+    get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement));
+
+    ScheduledFuture<?> future1 = waitForKeepAliveFuture(transaction);
+    assertNotNull(future1);
+
+    CountDownLatch latch = new CountDownLatch(1);
+    CountDownLatch lockAcquired = new CountDownLatch(1);
+    Thread lockHoldingThread =
+        new Thread(
+            () -> {
+              transaction.abortedLock.lock();
+              try {
+                lockAcquired.countDown();
+                latch.await();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              } finally {
+                transaction.abortedLock.unlock();
+              }
+            });
+    lockHoldingThread.start();
+    try {
+      lockAcquired.await();
+      ReadWriteTransaction.KeepAliveRunnable runnable =
+          new ReadWriteTransaction.KeepAliveRunnable(transaction);
+      runnable.run();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      fail("Test interrupted");
+    } finally {
+      latch.countDown();
+      try {
+        lockHoldingThread.join();
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      }
+    }
+
+    ScheduledFuture<?> future2 = transaction.getKeepAliveFuture();
+    assertNotNull(future2);
+    assertNotSame(future1, future2);
+  }
+
+  @Test
+  public void testKeepAliveFutureNullifiedOnCancel() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.UPDATE);
+    when(parsedStatement.isUpdate()).thenReturn(true);
+    Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
+    get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement));
+
+    assertNotNull(waitForKeepAliveFuture(transaction));
+
+    get(transaction.commitAsync(CallType.SYNC, NoopEndTransactionCallback.INSTANCE));
+
+    assertNull(transaction.getKeepAliveFuture());
+  }
+
+  private static ScheduledFuture<?> waitForKeepAliveFuture(ReadWriteTransaction transaction) {
+    long deadline = System.currentTimeMillis() + 5000;
+    while (System.currentTimeMillis() < deadline) {
+      ScheduledFuture<?> future = transaction.getKeepAliveFuture();
+      if (future != null) {
+        return future;
+      }
+      try {
+        Thread.sleep(1);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+    }
+    fail("Keep-alive future was not populated within 5 seconds");
+    return null;
+  }
+
+  @Test
+  public void testKeepAliveRunnableHandlesSynchronousException() {
+    DatabaseClient client = mock(DatabaseClient.class);
+    when(client.transactionManager())
+        .thenAnswer(
+            invocation -> {
+              TransactionContext txContext = mock(TransactionContext.class);
+              when(txContext.executeQuery(any(Statement.class)))
+                  .thenThrow(new RuntimeException("Simulated synchronous execution error"));
+              return new SimpleTransactionManager(txContext, CommitBehavior.SUCCEED);
+            });
+
+    ReadWriteTransaction transaction =
+        ReadWriteTransaction.newBuilder()
+            .setDatabaseClient(client)
+            .setKeepTransactionAlive(true)
+            .setRetryAbortsInternally(false)
+            .setIsolationLevel(IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED)
+            .setSavepointSupport(SavepointSupport.FAIL_AFTER_ROLLBACK)
+            .setTransactionRetryListeners(Collections.emptyList())
+            .withStatementExecutor(new StatementExecutor())
+            .setSpan(Span.getInvalid())
+            .build();
+
+    ReadWriteTransaction.KeepAliveRunnable runnable =
+        new ReadWriteTransaction.KeepAliveRunnable(transaction);
+
+    runnable.run();
+
+    assertFalse(transaction.abortedLock.isLocked());
+  }
+
+  @Test
+  public void testKeepAliveNotScheduledIfTransactionClosed() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.UPDATE);
+    when(parsedStatement.isUpdate()).thenReturn(true);
+    Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
+    get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement));
+    get(transaction.commitAsync(CallType.SYNC, NoopEndTransactionCallback.INSTANCE));
+
+    transaction.maybeScheduleKeepAlivePing();
+
+    assertNull(transaction.getKeepAliveFuture());
+  }
+
   private static StatusRuntimeException createAbortedExceptionWithMinimalRetry() {
     Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
     Metadata trailers = new Metadata();