fix(spanner): stop unclosed multiplexed clients leaking, and fix close() races (P1-7, P1-8, P1-9)

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<SpannerImpl, SharedChannelUsage> 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<SpannerImpl, SharedChannelUsage> CHANNEL_USAGE = new WeakHashMap<>();
 
   private static final EnumSet<ErrorCode> 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<MultiplexedSessionMaintainer> maintainerReference;
+
+    private final AtomicReference<ScheduledFuture<?>> 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);
       }

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<SpannerImpl> 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");