Remove custom thread pools

cloud-spanner-r2dbc/src/main/java/com/google/cloud/spanner/r2dbc/v2/DatabaseClientReactiveAdapter.java

@@ -35,8 +35,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.Executor;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import org.reactivestreams.Publisher;
@@ -45,6 +44,7 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.FluxSink;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 /**
  * Converts gRPC/Cloud Spanner client library asyncronous abstractions into reactive ones.
@@ -54,6 +54,9 @@ class DatabaseClientReactiveAdapter {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DatabaseClientReactiveAdapter.class);
 
+  public static final Executor REACTOR_EXECUTOR =
+      runnable -> Schedulers.parallel().schedule(runnable);
+
   // used for DDL operations
   private final SpannerConnectionConfiguration config;
 
@@ -63,12 +66,12 @@ class DatabaseClientReactiveAdapter {
 
   private final DatabaseAdminClient dbAdminClient;
 
-  private final ExecutorService executorService;
-
   private DatabaseClientTransactionManager txnManager;
 
   private boolean autoCommit = true;
 
+  private boolean active = true;
+
   private QueryOptions queryOptions;
 
   /**
@@ -84,9 +87,8 @@ class DatabaseClientReactiveAdapter {
     this.dbClient = spannerClient.getDatabaseClient(
         DatabaseId.of(config.getProjectId(), config.getInstanceName(), config.getDatabaseName()));
     this.dbAdminClient = spannerClient.getDatabaseAdminClient();
-    this.executorService = Executors.newFixedThreadPool(config.getThreadPoolSize());
     this.config = config;
-    this.txnManager = new DatabaseClientTransactionManager(this.dbClient, this.executorService);
+    this.txnManager = new DatabaseClientTransactionManager(this.dbClient);
 
     QueryOptions.Builder builder =  QueryOptions.newBuilder();
     if (config.getOptimizerVersion() != null) {
@@ -148,12 +150,14 @@ Publisher<Void> rollback() {
    * @return reactive pipeline for closing the connection.
    */
   Mono<Void> close() {
-    // TODO: if txn is committed/rolled back and then connection closed, clearTransactionManager
-    // will run twice, causing trace span to be closed twice. Introduce `closed` field.
-    return Mono.fromRunnable(() -> {
-      this.txnManager.clearTransactionManager();
-      this.executorService.shutdown();
+    return Mono.defer(() -> {
+      if (!this.active) {
+        return Mono.empty();
+      }
+      this.active = false;
+      return convertFutureToMono(() -> this.txnManager.clearTransactionManager());
     });
+
   }
 
   /**
@@ -163,7 +167,7 @@ Mono<Void> close() {
    */
   Mono<Boolean> healthCheck() {
     return Mono.defer(() -> {
-      if (this.executorService.isShutdown() || this.spannerClient.isClosed()) {
+      if (!this.active || this.spannerClient.isClosed()) {
         return Mono.just(false);
       } else {
         return Flux.<SpannerClientLibraryRow>create(sink -> {
@@ -181,7 +185,7 @@ Mono<Boolean> healthCheck() {
   }
 
   Mono<Boolean> localHealthcheck() {
-    return Mono.fromSupplier(() -> !this.executorService.isShutdown());
+    return Mono.fromSupplier(() -> this.active);
   }
 
   boolean isAutoCommit() {
@@ -239,7 +243,7 @@ private <T> Mono<T> runBatchDmlInternal(
           ApiFuture<T> rowCountFuture =
               this.dbClient
                   .runAsync()
-                  .runAsync(txn -> asyncOperation.apply(txn), this.executorService);
+                  .runAsync(txn -> asyncOperation.apply(txn), REACTOR_EXECUTOR);
           return rowCountFuture;
         }
       });
@@ -287,9 +291,7 @@ private ApiFuture<Void> runSelectStatementAsFlux(
     sink.onCancel(ars::cancel);
     sink.onDispose(ars::close);
 
-    return ars.setCallback(
-        this.executorService,
-        new ResultSetReadyCallback(sink));
+    return ars.setCallback(REACTOR_EXECUTOR, new ResultSetReadyCallback(sink));
   }
 
   static class ResultSetReadyCallback implements ReadyCallback {
@@ -339,7 +341,7 @@ public void onSuccess(T result) {
                   sink.success(result);
                 }
               },
-              this.executorService);
+              REACTOR_EXECUTOR);
         });
   }
 

cloud-spanner-r2dbc/src/main/java/com/google/cloud/spanner/r2dbc/v2/DatabaseClientTransactionManager.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.spanner.r2dbc.v2;
 
+import static com.google.cloud.spanner.r2dbc.v2.DatabaseClientReactiveAdapter.REACTOR_EXECUTOR;
+
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
@@ -28,7 +30,6 @@
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.r2dbc.TransactionInProgressException;
-import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,12 +54,8 @@ class DatabaseClientTransactionManager {
 
   private AsyncTransactionStep<?, ? extends Object> lastStep;
 
-  private final ExecutorService executorService;
-
-  public DatabaseClientTransactionManager(
-      DatabaseClient dbClient, ExecutorService executorService) {
+  public DatabaseClientTransactionManager(DatabaseClient dbClient) {
     this.dbClient = dbClient;
-    this.executorService = executorService;
   }
 
   boolean isInReadWriteTransaction() {
@@ -105,13 +102,16 @@ void beginReadonlyTransaction(TimestampBound timestampBound) {
   /**
    * Closes the read/write transaction manager and clears its state.
    */
-  void clearTransactionManager() {
+  ApiFuture<Void> clearTransactionManager() {
     this.txnContextFuture = null;
     this.lastStep = null;
+    ApiFuture<Void> returnFuture = ApiFutures.immediateFuture(null);
+
     if (this.transactionManager != null) {
-      this.transactionManager.close();
+      returnFuture = this.transactionManager.closeAsync();
       this.transactionManager = null;
     }
+    return returnFuture;
   }
 
   /**
@@ -174,9 +174,9 @@ <T> ApiFuture<T> runInTransaction(Function<? super TransactionContext, ApiFuture
     AsyncTransactionStep<? extends Object, T> updateStatementFuture =
         this.lastStep == null
             ? this.txnContextFuture.then(
-                (ctx, unusedVoid) -> operation.apply(ctx), this.executorService)
+                (ctx, unusedVoid) -> operation.apply(ctx), REACTOR_EXECUTOR)
             : this.lastStep.then(
-                (ctx, unusedPreviousResult) -> operation.apply(ctx), this.executorService);
+                (ctx, unusedPreviousResult) -> operation.apply(ctx), REACTOR_EXECUTOR);
 
     this.lastStep = updateStatementFuture;
     return updateStatementFuture;

cloud-spanner-r2dbc/src/test/java/com/google/cloud/spanner/r2dbc/v2/DatabaseClientTransactionManagerTest.java

@@ -29,7 +29,6 @@
 import com.google.cloud.spanner.r2dbc.TransactionInProgressException;
 import java.io.ByteArrayOutputStream;
 import java.io.PrintStream;
-import java.util.concurrent.Executors;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -72,9 +71,7 @@ public void setUp() {
     when(this.mockDbClient.readOnlyTransaction(TimestampBound.strong()))
         .thenReturn(this.mockReadOnlyTransaction);
 
-    this.transactionManager =
-        new DatabaseClientTransactionManager(
-            this.mockDbClient, Executors.newSingleThreadExecutor());
+    this.transactionManager = new DatabaseClientTransactionManager(this.mockDbClient);
 
   }