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,25 +167,23 @@ 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 -> {
-          com.google.cloud.spanner.Statement statement =
-              Statement.newBuilder("SELECT 1").build();
+          Statement statement = Statement.newBuilder("SELECT 1").build();
           runSelectStatementAsFlux(this.dbClient.singleUse(), statement, sink);
-        })
-        .then(Mono.just(true))
-        .onErrorResume(error -> {
-          LOGGER.warn("Cloud Spanner healthcheck failed", error);
-          return Mono.just(false);
-        });
+        }).then(Mono.just(true))
+            .onErrorResume(error -> {
+              LOGGER.warn("Cloud Spanner healthcheck failed", error);
+              return Mono.just(false);
+            });
       }
     });
   }
 
   Mono<Boolean> localHealthcheck() {
-    return Mono.fromSupplier(() -> !this.executorService.isShutdown());
+    return Mono.fromSupplier(() -> this.active);
   }
 
   boolean isAutoCommit() {
@@ -206,7 +208,7 @@ Publisher<Void> setAutoCommit(boolean autoCommit) {
    *
    * @return reactive pipeline for running a DML statement
    */
-  Mono<Long> runDmlStatement(com.google.cloud.spanner.Statement statement) {
+  Mono<Long> runDmlStatement(Statement statement) {
     return runBatchDmlInternal(ctx -> ctx.executeUpdateAsync(statement));
   }
 
@@ -239,15 +241,14 @@ private <T> Mono<T> runBatchDmlInternal(
           ApiFuture<T> rowCountFuture =
               this.dbClient
                   .runAsync()
-                  .runAsync(txn -> asyncOperation.apply(txn), this.executorService);
+                  .runAsync(asyncOperation::apply, REACTOR_EXECUTOR);
           return rowCountFuture;
         }
       });
     });
   }
 
-  Flux<SpannerClientLibraryRow> runSelectStatement(
-      com.google.cloud.spanner.Statement statement) {
+  Flux<SpannerClientLibraryRow> runSelectStatement(Statement statement) {
     return Flux.create(
         sink -> {
           if (this.txnManager.isInReadWriteTransaction()) {
@@ -280,16 +281,12 @@ Mono<Void> runDdlStatement(String query) {
    * @return future suitable for transactional step chaining
    */
   private ApiFuture<Void> runSelectStatementAsFlux(
-      ReadContext readContext,
-      com.google.cloud.spanner.Statement statement,
-      FluxSink<SpannerClientLibraryRow> sink) {
+      ReadContext readContext, Statement statement, FluxSink<SpannerClientLibraryRow> sink) {
     AsyncResultSet ars = readContext.executeQueryAsync(statement);
     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 +336,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,21 @@ 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;
     }
+
+    if (isInReadonlyTransaction()) {
+      closeReadOnlyTransaction();
+    }
+
+    return returnFuture;
   }
 
   /**
@@ -174,9 +179,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;