perf: close sessions async (#24)

* perf: close sessions async

Sessions should be closed using an async gRPC call in order
to speed up the closing of a large session pool. Instead of
using its own executor, which is limited to 8 worker threads,
to execute asynchronous delete session calls, the session pool
should use the asynchronous call option in gRPC. This allows a
larger number of asynchronous delete session calls to be executed
in parallel and speeds up closing a session pool with a large
number of sessions.

Testing against a real Cloud Spanner database with a session pool
containing 1,000 sessions shows the following performance for
closing the session pool:

Before (3 runs):
6603ms
8169ms
8367ms

After (3 runs):
1297ms
1710ms
1851ms

Fixes #19.

* fix: wait for test servers to terminate

* remove tracing for async call

* do not use directExecutor which could be a gRPC thread

* fix: return failed future instead of throwing exception

* fix: remove commented code

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Session.java

@@ -16,6 +16,9 @@
 
 package com.google.cloud.spanner;
 
+import com.google.api.core.ApiFuture;
+import com.google.protobuf.Empty;
+
 /**
  * A {@code Session} can be used to perform transactions that read and/or modify data in a Cloud
  * Spanner database.
@@ -54,4 +57,10 @@ public interface Session extends DatabaseClient, AutoCloseable {
 
   @Override
   void close();
+
+  /**
+   * Closes the session asynchronously and returns the {@link ApiFuture} that can be used to monitor
+   * the operation progress.
+   */
+  ApiFuture<Empty> asyncClose();
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -19,6 +19,7 @@
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.api.core.ApiFuture;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction;
 import com.google.cloud.spanner.AbstractReadContext.SingleReadContext;
@@ -27,6 +28,7 @@
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.common.collect.Lists;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Empty;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.CommitResponse;
@@ -196,6 +198,11 @@ public void prepareReadWriteTransaction() {
     readyTransactionId = beginTransaction();
   }
 
+  @Override
+  public ApiFuture<Empty> asyncClose() {
+    return spanner.getRpc().asyncDeleteSession(name, options);
+  }
+
   @Override
   public void close() {
     Span span = tracer.spanBuilder(SpannerImpl.DELETE_SESSION).startSpan();

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java

@@ -18,6 +18,8 @@
 
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
@@ -35,6 +37,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.common.util.concurrent.SettableFuture;
 import com.google.common.util.concurrent.Uninterruptibles;
+import com.google.protobuf.Empty;
 import io.opencensus.common.Scope;
 import io.opencensus.trace.Annotation;
 import io.opencensus.trace.AttributeValue;
@@ -767,6 +770,12 @@ public TransactionRunner readWriteTransaction() {
       return new SessionPoolTransactionRunner(SessionPool.this, this);
     }
 
+    @Override
+    public ApiFuture<Empty> asyncClose() {
+      close();
+      return ApiFutures.immediateFuture(Empty.getDefaultInstance());
+    }
+
     @Override
     public void close() {
       synchronized (lock) {
@@ -999,7 +1008,7 @@ private void closeIdleSessions(Instant currTime) {
       }
       for (PooledSession sess : sessionsToClose) {
         logger.log(Level.FINE, "Closing session {0}", sess.getName());
-        closeSession(sess);
+        closeSessionAsync(sess);
       }
     }
 
@@ -1612,37 +1621,27 @@ int totalSessions() {
     }
   }
 
-  private void closeSessionAsync(final PooledSession sess) {
-    executor.submit(
+  private ApiFuture<Empty> closeSessionAsync(final PooledSession sess) {
+    ApiFuture<Empty> res = sess.delegate.asyncClose();
+    res.addListener(
         new Runnable() {
           @Override
           public void run() {
-            closeSession(sess);
+            synchronized (lock) {
+              allSessions.remove(sess);
+              if (isClosed()) {
+                decrementPendingClosures(1);
+                return;
+              }
+              // Create a new session if needed to unblock some waiter.
+              if (numWaiters() > numSessionsBeingCreated) {
+                createSessions(getAllowedCreateSessions(numWaiters() - numSessionsBeingCreated));
+              }
+            }
           }
-        });
-  }
-
-  private void closeSession(PooledSession sess) {
-    try {
-      sess.delegate.close();
-    } catch (SpannerException e) {
-      // Backend will delete these sessions after a while even if we fail to close them.
-      if (logger.isLoggable(Level.FINE)) {
-        logger.log(Level.FINE, "Failed to close session: " + sess.getName(), e);
-      }
-    } finally {
-      synchronized (lock) {
-        allSessions.remove(sess);
-        if (isClosed()) {
-          decrementPendingClosures(1);
-          return;
-        }
-        // Create a new session if needed to unblock some waiter.
-        if (numWaiters() > numSessionsBeingCreated) {
-          createSessions(getAllowedCreateSessions(numWaiters() - numSessionsBeingCreated));
-        }
-      }
-    }
+        },
+        executor);
+    return res;
   }
 
   private void prepareSession(final PooledSession sess) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -18,6 +18,7 @@
 
 import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.NanoClock;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.ExecutorProvider;
@@ -523,9 +524,14 @@ public Session createSession(
   @Override
   public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
       throws SpannerException {
+    get(asyncDeleteSession(sessionName, options));
+  }
+
+  @Override
+  public ApiFuture<Empty> asyncDeleteSession(String sessionName, @Nullable Map<Option, ?> options) {
     DeleteSessionRequest request = DeleteSessionRequest.newBuilder().setName(sessionName).build();
     GrpcCallContext context = newCallContext(options, sessionName);
-    get(spannerStub.deleteSessionCallable().futureCall(request, context));
+    return spannerStub.deleteSessionCallable().futureCall(request, context);
   }
 
   @Override

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java

@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.spi.v1;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.ServiceRpc;
 import com.google.cloud.spanner.SpannerException;
@@ -219,6 +220,9 @@ Session createSession(
 
   void deleteSession(String sessionName, @Nullable Map<Option, ?> options) throws SpannerException;
 
+  ApiFuture<Empty> asyncDeleteSession(String sessionName, @Nullable Map<Option, ?> options)
+      throws SpannerException;
+
   StreamingCall read(
       ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options);
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java

@@ -23,8 +23,10 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
+import com.google.api.core.ApiFutures;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
 import com.google.cloud.spanner.SessionPool.Clock;
+import com.google.protobuf.Empty;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
@@ -61,6 +63,7 @@ SessionImpl mockSession() {
     when(session.getName())
         .thenReturn(
             "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex);
+    when(session.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
     sessionIndex++;
     return session;
   }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseAdminGaxTest.java

@@ -216,8 +216,9 @@ public static void startStaticServer() throws IOException {
   }
 
   @AfterClass
-  public static void stopServer() {
+  public static void stopServer() throws InterruptedException {
     server.shutdown();
+    server.awaitTermination();
   }
 
   @Before

google-cloud-spanner/src/test/java/com/google/cloud/spanner/InstanceAdminGaxTest.java

@@ -219,8 +219,9 @@ public static void startStaticServer() throws IOException {
   }
 
   @AfterClass
-  public static void stopServer() {
+  public static void stopServer() throws InterruptedException {
     server.shutdown();
+    server.awaitTermination();
   }
 
   @Before

google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java

@@ -172,9 +172,10 @@ public static void startStaticServer() throws IOException {
   }
 
   @AfterClass
-  public static void stopServer() {
+  public static void stopServer() throws InterruptedException {
     spannerClient.close();
     server.shutdown();
+    server.awaitTermination();
   }
 
   @Before

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java

@@ -29,6 +29,7 @@
 import io.grpc.StatusRuntimeException;
 import io.grpc.inprocess.InProcessServerBuilder;
 import java.io.IOException;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -55,13 +56,19 @@ public static void startStaticServer() throws IOException {
     mockSpanner = new MockSpannerServiceImpl();
     mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions.
     String uniqueName = InProcessServerBuilder.generateName();
-    server = InProcessServerBuilder.forName(uniqueName).addService(mockSpanner).build().start();
+    server =
+        InProcessServerBuilder.forName(uniqueName)
+            .scheduledExecutorService(new ScheduledThreadPoolExecutor(1))
+            .addService(mockSpanner)
+            .build()
+            .start();
     channelProvider = LocalChannelProvider.create(uniqueName);
   }
 
   @AfterClass
-  public static void stopServer() {
+  public static void stopServer() throws InterruptedException {
     server.shutdown();
+    server.awaitTermination();
   }
 
   @Before