Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
+208 −15
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 61 additions & 14 deletions ...cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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"));
Expand Down Expand Up @@ -146,7 +158,7 @@ class ReadWriteTransaction extends AbstractMultiUseTransaction {
private Savepoint autoSavepoint;

private final int maxInternalRetries;
private final ReentrantLock abortedLock = new ReentrantLock();
final ReentrantLock abortedLock = new ReentrantLock();
private final long transactionId;
private final DatabaseClient dbClient;
private final TransactionOption[] transactionOptions;
Expand Down Expand Up @@ -475,7 +487,7 @@ private void maybeScheduleKeepAlivePing() {
if (keepAliveFuture == null || keepAliveFuture.isDone()) {
keepAliveFuture =
KEEP_ALIVE_SERVICE.schedule(
new KeepAliveRunnable(),
new KeepAliveRunnable(this),
keepAliveIntervalMillis > 0
? keepAliveIntervalMillis
: DEFAULT_KEEP_ALIVE_INTERVAL_MILLIS,
Expand All @@ -487,36 +499,71 @@ 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();
}
}
}

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()) {
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) {
transaction.maybeScheduleKeepAlivePing();
} finally {
transaction.abortedLock.unlock();
}
} else {
// Transaction is currently busy (executing a statement or retrying).
// Reschedule keep-alive ping for later since it is active.
transaction.rescheduleKeepAlivePing();
}
}
}
}
Expand Down
148 changes: 147 additions & 1 deletion ...d-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,11 @@
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.fail;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.doThrow;
Expand Down Expand Up @@ -67,6 +70,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;
Expand Down Expand Up @@ -158,12 +163,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(
Expand All @@ -179,6 +193,7 @@ private ReadWriteTransaction createSubject(
});
return ReadWriteTransaction.newBuilder()
.setDatabaseClient(client)
.setKeepTransactionAlive(keepTransactionAlive)
.setRetryAbortsInternally(withRetry)
.setIsolationLevel(IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED)
.setSavepointSupport(SavepointSupport.FAIL_AFTER_ROLLBACK)
Expand Down Expand Up @@ -857,6 +872,137 @@ 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);

int initialQueueSize = ReadWriteTransaction.getKeepAliveService().getQueue().size();
ReadWriteTransaction transaction = createSubject(/* keepTransactionAlive= */ true);
get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement));

assertEquals(
initialQueueSize + 1, ReadWriteTransaction.getKeepAliveService().getQueue().size());

get(transaction.commitAsync(CallType.SYNC, NoopEndTransactionCallback.INSTANCE));
assertEquals(initialQueueSize, ReadWriteTransaction.getKeepAliveService().getQueue().size());
}

@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 = transaction.getKeepAliveFuture();
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(transaction.getKeepAliveFuture());

get(transaction.commitAsync(CallType.SYNC, NoopEndTransactionCallback.INSTANCE));

assertNull(transaction.getKeepAliveFuture());
}

@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());
}

private static StatusRuntimeException createAbortedExceptionWithMinimalRetry() {
Metadata.Key<RetryInfo> key = ProtoUtils.keyForProto(RetryInfo.getDefaultInstance());
Metadata trailers = new Metadata();
Expand Down
Loading