chore(spanner): make AsyncRunner non blocking with multiplexed sessions

Author: harshachinta

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

@@ -17,8 +17,9 @@
 package com.google.cloud.spanner;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
-import java.util.concurrent.ExecutionException;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.Executor;
 
 /**
@@ -34,33 +35,36 @@ public DelayedAsyncRunner(ApiFuture<AsyncRunner> asyncRunnerFuture) {
     this.asyncRunnerFuture = asyncRunnerFuture;
   }
 
-  AsyncRunner getAsyncRunner() {
-    try {
-      return this.asyncRunnerFuture.get();
-    } catch (ExecutionException executionException) {
-      // Propagate the underlying exception as a RuntimeException (SpannerException is also a
-      // RuntimeException).
-      if (executionException.getCause() instanceof RuntimeException) {
-        throw (RuntimeException) executionException.getCause();
-      }
-      throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
-    } catch (InterruptedException interruptedException) {
-      throw SpannerExceptionFactory.propagateInterrupt(interruptedException);
-    }
+  ApiFuture<AsyncRunner> getAsyncRunner() {
+    return ApiFutures.catchingAsync(
+        asyncRunnerFuture,
+        Exception.class,
+        exception -> {
+          if (exception instanceof InterruptedException) {
+            throw SpannerExceptionFactory.propagateInterrupt((InterruptedException) exception);
+          }
+          throw SpannerExceptionFactory.asSpannerException(exception.getCause());
+        },
+        MoreExecutors.directExecutor());
   }
 
   @Override
   public <R> ApiFuture<R> runAsync(AsyncWork<R> work, Executor executor) {
-    return getAsyncRunner().runAsync(work, executor);
+    return ApiFutures.transformAsync(
+        getAsyncRunner(),
+        asyncRunner -> asyncRunner.runAsync(work, executor),
+        MoreExecutors.directExecutor());
   }
 
   @Override
   public ApiFuture<Timestamp> getCommitTimestamp() {
-    return getAsyncRunner().getCommitTimestamp();
+    return ApiFutures.transformAsync(
+        getAsyncRunner(), AsyncRunner::getCommitTimestamp, MoreExecutors.directExecutor());
   }
 
   @Override
   public ApiFuture<CommitResponse> getCommitResponse() {
-    return getAsyncRunner().getCommitResponse();
+    return ApiFutures.transformAsync(
+        getAsyncRunner(), AsyncRunner::getCommitResponse, MoreExecutors.directExecutor());
   }
 }

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

@@ -713,6 +713,38 @@ public void testReadWriteTransactionUsingAsyncRunner() throws Exception {
     assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
   }
 
+  @Test
+  public void testAsyncRunnerIsNonBlockingWithMultiplexedSession() throws Exception {
+    mockSpanner.freeze();
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+    AsyncRunner runner = client.runAsync();
+    ApiFuture<Void> res =
+        runner.runAsync(
+            txn -> {
+              txn.executeUpdateAsync(UPDATE_STATEMENT);
+              return ApiFutures.immediateFuture(null);
+            },
+            MoreExecutors.directExecutor());
+    ApiFuture<Timestamp> ts = runner.getCommitTimestamp();
+    mockSpanner.unfreeze();
+    assertThat(res.get()).isNull();
+    assertThat(ts.get()).isNotNull();
+
+    List<ExecuteSqlRequest> executeSqlRequests =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(1L, executeSqlRequests.size());
+
+    // Verify the requests are executed using multiplexed sessions
+    for (ExecuteSqlRequest request : executeSqlRequests) {
+      assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
+    }
+
+    assertNotNull(client.multiplexedSessionDatabaseClient);
+    assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get());
+    assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+  }
+
   private void waitForSessionToBeReplaced(DatabaseClientImpl client) {
     assertNotNull(client.multiplexedSessionDatabaseClient);
     SessionReference sessionReference =