diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index ece6d862f87b..ebf1ee66b012 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -31,13 +31,14 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.spanner.v1.BatchWriteResponse; +import java.lang.ref.WeakReference; import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.util.BitSet; import java.util.EnumSet; -import java.util.HashMap; import java.util.Map; +import java.util.WeakHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -172,7 +173,13 @@ private SharedChannelUsage(int numChannels) { } } - private static final Map CHANNEL_USAGE = new HashMap<>(); + /** + * Keyed weakly on the {@link SpannerImpl} instance, so that an application that drops a {@link + * Spanner} instance without closing it does not permanently pin the entire client graph through + * this static map. The entry for a {@link SpannerImpl} that is still in use is kept alive by the + * strong reference that each {@link MultiplexedSessionDatabaseClient} holds to it. + */ + private static final Map CHANNEL_USAGE = new WeakHashMap<>(); private static final EnumSet RETRYABLE_ERROR_CODES = EnumSet.of(ErrorCode.DEADLINE_EXCEEDED, ErrorCode.RESOURCE_EXHAUSTED, ErrorCode.UNAVAILABLE); @@ -186,7 +193,12 @@ private SharedChannelUsage(int numChannels) { */ private final AtomicInteger numCurrentSingleUseTransactions = new AtomicInteger(); - private boolean isClosed; + /** + * Written under {@code synchronized(this)} in {@link #close()}, but read without any + * synchronization in {@link #createMultiplexedSessionTransaction(boolean)}. It must therefore be + * volatile for the latter to be guaranteed to observe the close. + */ + private volatile boolean isClosed; /** The duration before we try to replace the multiplexed session. The default is 7 days. */ private final Duration sessionExpirationDuration; @@ -623,16 +635,65 @@ public long executePartitionedUpdate(Statement stmt, UpdateOption... options) { ThreadFactoryUtil.createVirtualOrPlatformDaemonThreadFactory( "multiplexed-session-maintainer", /* tryVirtual= */ false)); + /** + * The task that is scheduled on the static {@link #MAINTAINER_SERVICE}. It only holds a {@link + * WeakReference} to the maintainer, so that an application that drops a {@link Spanner} instance + * without closing it is not kept alive by the static executor. The task cancels itself as soon as + * it notices that the maintainer has been garbage collected. + */ + @VisibleForTesting + static final class MaintainerTask implements Runnable { + private final WeakReference maintainerReference; + + private final AtomicReference> scheduledFuture = new AtomicReference<>(); + + MaintainerTask(MultiplexedSessionMaintainer maintainer) { + this.maintainerReference = new WeakReference<>(maintainer); + } + + void setScheduledFuture(ScheduledFuture scheduledFuture) { + this.scheduledFuture.set(scheduledFuture); + } + + @Override + public void run() { + MultiplexedSessionMaintainer maintainer = this.maintainerReference.get(); + if (maintainer == null) { + // The client was garbage collected without being closed. Cancel this task, as it would + // otherwise keep running for the lifetime of the JVM. + ScheduledFuture future = this.scheduledFuture.get(); + if (future != null) { + future.cancel(false); + } + return; + } + maintainer.maintain(); + } + } + final class MultiplexedSessionMaintainer { private final Clock clock; private ScheduledFuture scheduledFuture; + private boolean stopped; + MultiplexedSessionMaintainer(Clock clock) { this.clock = clock; } private synchronized void start() { + if (this.stopped) { + // The client was closed while the initial CreateSession RPC was still in flight. Scheduling + // a maintenance task at this point would leak it, as close() has already run and nothing + // would ever cancel it again. + return; + } + if (this.scheduledFuture != null) { + // Already started. Scheduling a second task would leak the first one, as only the last + // scheduled future is retained and can be cancelled by stop(). + return; + } // Schedule the maintainer to run once every ten minutes (by default). long loopFrequencyMillis = MultiplexedSessionDatabaseClient.this @@ -642,12 +703,15 @@ private synchronized void start() { .getSessionPoolOptions() .getMultiplexedSessionMaintenanceLoopFrequency() .toMillis(); + MaintainerTask task = new MaintainerTask(this); this.scheduledFuture = MAINTAINER_SERVICE.scheduleAtFixedRate( - this::maintain, loopFrequencyMillis, loopFrequencyMillis, TimeUnit.MILLISECONDS); + task, loopFrequencyMillis, loopFrequencyMillis, TimeUnit.MILLISECONDS); + task.setScheduledFuture(this.scheduledFuture); } private synchronized void stop() { + this.stopped = true; if (this.scheduledFuture != null) { this.scheduledFuture.cancel(false); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java index d398a78643aa..673fc8de6aa8 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientTest.java @@ -19,13 +19,19 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assume.assumeFalse; import static org.junit.Assume.assumeTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; @@ -34,13 +40,16 @@ import com.google.cloud.spanner.SessionClient.SessionConsumer; import java.io.PrintWriter; import java.io.StringWriter; +import java.lang.ref.WeakReference; import java.lang.reflect.Field; +import java.lang.reflect.Modifier; import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.util.Map; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -348,6 +357,159 @@ public void testCloseKeepsChannelUsageEntryWhileAnotherClientIsUsingSameSpanner( } } + @Test + public void testChannelUsageDoesNotPinSpannerThatIsNotClosed() throws Exception { + assertEquals(0, getChannelUsage().size()); + + SpannerImpl spanner = createTestSpanner(); + SessionClient sessionClient = createSessionClient(spanner); + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + assertEquals(1, getChannelUsage().size()); + + // Drop all references to the client and the Spanner instance *without* closing them. The + // CHANNEL_USAGE map is keyed weakly, so it must not keep the Spanner instance (and with it the + // entire client graph) alive. + WeakReference spannerReference = new WeakReference<>(spanner); + client = null; + sessionClient = null; + spanner = null; + + for (int attempt = 0; attempt < 100 && spannerReference.get() != null; attempt++) { + System.gc(); + Thread.sleep(10L); + } + + assertNull(spannerReference.get()); + assertEquals(0, getChannelUsage().size()); + } + + @Test + public void testSessionReadySchedulesMaintainer() throws Exception { + try (SpannerImpl spanner = createTestSpanner(); + DeferredMultiplexedSessionClient sessionClient = + new DeferredMultiplexedSessionClient(spanner)) { + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + // The CreateSession RPC is still in flight, so no maintenance task has been scheduled yet. + assertNull(getScheduledFuture(client)); + + sessionClient.completeSessionCreation(); + ScheduledFuture scheduledFuture = getScheduledFuture(client); + assertNotNull(scheduledFuture); + assertFalse(scheduledFuture.isCancelled()); + + client.close(); + assertTrue(scheduledFuture.isCancelled()); + } + } + + @Test + public void testMaintainerIsOnlyScheduledOnce() throws Exception { + try (SpannerImpl spanner = createTestSpanner(); + DeferredMultiplexedSessionClient sessionClient = + new DeferredMultiplexedSessionClient(spanner)) { + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + + sessionClient.completeSessionCreation(); + ScheduledFuture scheduledFuture = getScheduledFuture(client); + assertNotNull(scheduledFuture); + + // A second call must not schedule a second task, as only the last scheduled future is + // retained and can be cancelled by close(). + sessionClient.completeSessionCreation(); + assertSame(scheduledFuture, getScheduledFuture(client)); + + client.close(); + assertTrue(scheduledFuture.isCancelled()); + } + } + + @Test + public void testCloseBeforeSessionReadyDoesNotScheduleMaintainer() throws Exception { + try (SpannerImpl spanner = createTestSpanner(); + DeferredMultiplexedSessionClient sessionClient = + new DeferredMultiplexedSessionClient(spanner)) { + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + + // Close the client while the initial CreateSession RPC is still in flight. maintainer.stop() + // has nothing to cancel at this point. + client.close(); + + // The CreateSession RPC now completes. This must not schedule a maintenance task, as nothing + // would ever cancel it again. + sessionClient.completeSessionCreation(); + + assertNull(getScheduledFuture(client)); + } + } + + @Test + public void testMaintainerTaskCancelsItselfWhenMaintainerIsCollected() throws Exception { + try (SpannerImpl spanner = createTestSpanner(); + SessionClient sessionClient = createSessionClient(spanner)) { + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + + ScheduledFuture scheduledFuture = mock(ScheduledFuture.class); + MultiplexedSessionDatabaseClient.MaintainerTask task = + new MultiplexedSessionDatabaseClient.MaintainerTask(client.getMaintainer()); + task.setScheduledFuture(scheduledFuture); + + // As long as the maintainer is reachable, the task just runs the maintenance. The session has + // not expired, so this is a no-op. + task.run(); + verify(scheduledFuture, never()).cancel(anyBoolean()); + + // Clear the weak reference to simulate the client being garbage collected without being + // closed. The task must then cancel itself, so that the static executor stops running it. + clearMaintainerReference(task); + task.run(); + verify(scheduledFuture).cancel(false); + } + } + + @Test + public void testUseAfterCloseThrows() { + try (SpannerImpl spanner = createTestSpanner(); + SessionClient sessionClient = createSessionClient(spanner)) { + MultiplexedSessionDatabaseClient client = + new MultiplexedSessionDatabaseClient(sessionClient, Clock.systemUTC()); + client.close(); + + assertThrows(IllegalStateException.class, client::singleUse); + } + } + + @Test + public void testIsClosedIsVolatile() throws Exception { + // isClosed is written under synchronized(this) in close(), but read without any synchronization + // in createMultiplexedSessionTransaction(..). It must therefore be volatile for that read to be + // guaranteed to observe the close. + Field field = MultiplexedSessionDatabaseClient.class.getDeclaredField("isClosed"); + assertTrue(Modifier.isVolatile(field.getModifiers())); + } + + private ScheduledFuture getScheduledFuture(MultiplexedSessionDatabaseClient client) + throws Exception { + Field field = + MultiplexedSessionDatabaseClient.MultiplexedSessionMaintainer.class.getDeclaredField( + "scheduledFuture"); + field.setAccessible(true); + return (ScheduledFuture) field.get(client.getMaintainer()); + } + + private void clearMaintainerReference(MultiplexedSessionDatabaseClient.MaintainerTask task) + throws Exception { + Field field = + MultiplexedSessionDatabaseClient.MaintainerTask.class.getDeclaredField( + "maintainerReference"); + field.setAccessible(true); + ((WeakReference) field.get(task)).clear(); + } + private SessionClient createSessionClient(SpannerImpl spanner) { return new FailingMultiplexedSessionClient(spanner); } @@ -405,6 +567,33 @@ public void release(ScheduledExecutorService executor) { } } + /** + * {@link SessionClient} that does not complete the CreateSession RPC until {@link + * #completeSessionCreation()} is called. This allows tests to interleave {@link + * MultiplexedSessionDatabaseClient#close()} with the completion of the initial session creation. + */ + private static final class DeferredMultiplexedSessionClient extends SessionClient { + private static final DatabaseId TEST_DATABASE_ID = + DatabaseId.of("test-project", "test-instance", "test-database"); + + private SessionConsumer consumer; + + private DeferredMultiplexedSessionClient(SpannerImpl spanner) { + super(spanner, TEST_DATABASE_ID, new TestExecutorFactory()); + } + + @Override + void asyncCreateMultiplexedSession(SessionConsumer consumer) { + this.consumer = consumer; + } + + void completeSessionCreation() { + SessionImpl session = mock(SessionImpl.class); + when(session.getSessionReference()).thenReturn(mock(SessionReference.class)); + this.consumer.onSessionReady(session); + } + } + private static final class FailingMultiplexedSessionClient extends SessionClient { private static final DatabaseId TEST_DATABASE_ID = DatabaseId.of("test-project", "test-instance", "test-database");