From a1f5720f3aee59a9c186d3b79b340e9dfc42a8f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 27 Jul 2026 11:07:27 +0200 Subject: [PATCH] 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` 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. --- .../connection/ReadWriteTransaction.java | 102 +++++++--- .../connection/ReadWriteTransactionTest.java | 185 +++++++++++++++++- 2 files changed, 263 insertions(+), 24 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java index ccb592e3f843..1b3830876d75 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java +++ b/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,12 +500,26 @@ private void maybeScheduleKeepAlivePing() { } } + @VisibleForTesting + ScheduledFuture getKeepAliveFuture() { + if (keepAliveLock != null) { + keepAliveLock.lock(); + try { + return keepAliveFuture; + } finally { + keepAliveLock.unlock(); + } + } + return null; + } + private void cancelScheduledKeepAlivePing() { if (keepAliveLock != null) { keepAliveLock.lock(); try { if (keepAliveFuture != null) { keepAliveFuture.cancel(false); + keepAliveFuture = null; } } finally { keepAliveLock.unlock(); @@ -500,23 +527,52 @@ private void cancelScheduledKeepAlivePing() { } } - private class KeepAliveRunnable implements Runnable { + @VisibleForTesting + static class KeepAliveRunnable implements Runnable { + final WeakReference transactionRef; + + KeepAliveRunnable(ReadWriteTransaction transaction) { + this.transactionRef = new WeakReference<>(transaction); + } + @Override + @SuppressWarnings("FutureReturnValueIgnored") 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 future = - 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()); + ReadWriteTransaction transaction = transactionRef.get(); + if (transaction != null && transaction.shouldPing()) { + transaction.keepAliveLock.lock(); + try { + transaction.keepAliveFuture = null; + } finally { + transaction.keepAliveLock.unlock(); + } + if (transaction.shouldPing()) { + boolean schedulePing = false; + if (transaction.abortedLock.tryLock()) { + 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"))); + } catch (Throwable t) { + schedulePing = true; + } finally { + transaction.abortedLock.unlock(); + } + } else { + schedulePing = true; + } + if (schedulePing) { + transaction.maybeScheduleKeepAlivePing(); + } + } } } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java index 7d0fa94c9b0b..41acd084da4d 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java +++ b/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,173 @@ 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(); + assertNotSame(future1, future2); + assertNotNull(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()); + assertNotNull(waitForKeepAliveFuture(transaction)); + } + + @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 key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance()); Metadata trailers = new Metadata();