fix: rollback could return session to pool twice (#888)

* fix: rollback could return session to pool twice

Transaction rollbacks could cause a session to be added to the session pool twice.
This again could cause two transactions to try to use the same session at the same
time, which again could lead to unpredictable errors.

* chore: add explanation + test for aborted tx

Author: olavloite

cloud-spanner-r2dbc/pom.xml

@@ -83,6 +83,20 @@
       <scope>provided</scope>
     </dependency>
 
+    <!-- Test dependencies for a mock Spanner server -->
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-spanner</artifactId>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.api</groupId>
+      <artifactId>gax-grpc</artifactId>
+      <classifier>testlib</classifier>
+      <scope>test</scope>
+    </dependency>
+
   </dependencies>
 
   <build>

cloud-spanner-r2dbc/src/main/java/com/google/cloud/spanner/r2dbc/SpannerConnectionConfiguration.java

@@ -20,6 +20,7 @@
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.r2dbc.util.Assert;
+import io.grpc.ManagedChannelBuilder;
 import io.r2dbc.spi.ConnectionFactoryOptions;
 import io.r2dbc.spi.R2dbcNonTransientResourceException;
 import java.io.IOException;
@@ -51,11 +52,13 @@ public class SpannerConnectionConfiguration {
   // TODO: check how to handle full URL (it gets parsed by SPI, we only get pieces)
   private final String fullyQualifiedDbName;
 
-  private String projectId;
+  private final ConnectionFactoryOptions options;
 
-  private String instanceName;
+  private final String projectId;
 
-  private String databaseName;
+  private final String instanceName;
+
+  private final String databaseName;
 
   private final OAuth2Credentials credentials;
 
@@ -82,6 +85,7 @@ public class SpannerConnectionConfiguration {
    * @param credentials GCP credentials to authenticate service calls with.
    */
   private SpannerConnectionConfiguration(
+      ConnectionFactoryOptions options,
       String projectId,
       String instanceName,
       String databaseName,
@@ -91,6 +95,7 @@ private SpannerConnectionConfiguration(
     Assert.requireNonNull(instanceName, "instanceName must not be null");
     Assert.requireNonNull(databaseName, "databaseName must not be null");
 
+    this.options = Assert.requireNonNull(options, "options must not be null");
     this.projectId = projectId;
     this.instanceName = instanceName;
     this.databaseName = databaseName;
@@ -189,13 +194,24 @@ public int hashCode() {
   public SpannerOptions buildSpannerOptions() {
     SpannerOptions.Builder optionsBuilder = SpannerOptions.newBuilder();
 
+    if (this.options.hasOption(ConnectionFactoryOptions.HOST)) {
+      int port = 443;
+      if (this.options.hasOption(ConnectionFactoryOptions.PORT)) {
+        port = (Integer) Objects.requireNonNull(
+            this.options.getValue(ConnectionFactoryOptions.PORT));
+      }
+      String host = String.format("http://%s:%d", this.options.getValue(ConnectionFactoryOptions.HOST), port);
+      optionsBuilder.setHost(host);
+    }
     if (this.projectId != null) {
       optionsBuilder.setProjectId(this.projectId);
     }
-
     if (this.credentials != null) {
       optionsBuilder.setCredentials(this.credentials);
     }
+    if (this.usePlainText) {
+      optionsBuilder.setChannelConfigurator(ManagedChannelBuilder::usePlaintext);
+    }
 
     optionsBuilder.setHeaderProvider(() ->
         Collections.singletonMap(USER_AGENT_KEY, USER_AGENT_LIBRARY_NAME + "/" + PACKAGE_VERSION));
@@ -210,6 +226,7 @@ public SpannerOptions buildSpannerOptions() {
    * Builder for the enclosing class.
    */
   public static class Builder {
+    private final ConnectionFactoryOptions options;
 
     private String fullyQualifiedDatabaseName;
 
@@ -235,6 +252,10 @@ public static class Builder {
 
     private boolean autocommit = true;
 
+    public Builder(ConnectionFactoryOptions options) {
+      this.options = Assert.requireNonNull(options, "options must not be null");
+    }
+
     /**
      * R2DBC SPI does not provide the full URL to drivers after parsing the connection string.
      * Therefore, this usecase is only possible if the client application provides a URL property
@@ -359,7 +380,7 @@ public SpannerConnectionConfiguration build() {
       }
 
       SpannerConnectionConfiguration configuration = new SpannerConnectionConfiguration(
-            this.projectId, this.instanceName, this.databaseName, this.credentials);
+            this.options, this.projectId, this.instanceName, this.databaseName, this.credentials);
 
       configuration.partialResultSetFetchSize = this.partialResultSetFetchSize;
       configuration.ddlOperationTimeout = this.ddlOperationTimeout;

cloud-spanner-r2dbc/src/main/java/com/google/cloud/spanner/r2dbc/SpannerConnectionFactoryProvider.java

@@ -22,7 +22,6 @@
 import static io.r2dbc.spi.ConnectionFactoryOptions.DATABASE;
 import static io.r2dbc.spi.ConnectionFactoryOptions.DRIVER;
 
-import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.r2dbc.util.Assert;
@@ -68,7 +67,7 @@ public class SpannerConnectionFactoryProvider implements ConnectionFactoryProvid
   /**
    * Option specifying the already-instantiated credentials object.
    */
-  public static final Option<GoogleCredentials> GOOGLE_CREDENTIALS =
+  public static final Option<OAuth2Credentials> GOOGLE_CREDENTIALS =
       Option.valueOf("google_credentials");
 
   public static final Option<Boolean> AUTOCOMMIT = Option.valueOf(AUTOCOMMIT_PROPERTY_NAME);
@@ -117,7 +116,8 @@ public String getDriver() {
   SpannerConnectionConfiguration createConfiguration(
       ConnectionFactoryOptions options) {
 
-    SpannerConnectionConfiguration.Builder config = new SpannerConnectionConfiguration.Builder();
+    SpannerConnectionConfiguration.Builder config =
+        new SpannerConnectionConfiguration.Builder(options);
 
     // Directly passed URL is supported for backwards compatibility. R2DBC SPI does not provide
     // the original URL when creating connection through ConnectionFactories.get(String).

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

@@ -29,6 +29,7 @@
 import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.TransactionContext;
+import com.google.cloud.spanner.TransactionManager.TransactionState;
 import com.google.cloud.spanner.r2dbc.TransactionInProgressException;
 import java.util.function.Function;
 import org.slf4j.Logger;
@@ -108,7 +109,13 @@ ApiFuture<Void> clearTransactionManager() {
     ApiFuture<Void> returnFuture = ApiFutures.immediateFuture(null);
 
     if (this.transactionManager != null) {
-      returnFuture = this.transactionManager.closeAsync();
+      // We don't close the transaction manager here if it has already been rolled back.
+      // Rolling back a transaction manager automatically also closes it, and returns the underlying
+      // session to the pool. Closing it a second time here, would cause it to be added to the pool
+      // a second time.
+      if (this.transactionManager.getState() != TransactionState.ROLLED_BACK) {
+        returnFuture = this.transactionManager.closeAsync();
+      }
       this.transactionManager = null;
     }
 

cloud-spanner-r2dbc/src/test/java/com/google/cloud/spanner/r2dbc/SpannerConnectionConfigurationTest.java

@@ -18,21 +18,22 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.mock;
 
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.r2dbc.SpannerConnectionConfiguration.Builder;
+import io.r2dbc.spi.ConnectionFactoryOptions;
 import java.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
 
 /**
  * Test for {@link SpannerConnectionConfiguration}.
  */
 class SpannerConnectionConfigurationTest {
 
-  GoogleCredentials mockCredentials = Mockito.mock(GoogleCredentials.class);
+  GoogleCredentials mockCredentials = mock(GoogleCredentials.class);
 
   SpannerConnectionConfiguration.Builder configurationBuilder;
 
@@ -41,13 +42,15 @@ class SpannerConnectionConfigurationTest {
    */
   @BeforeEach
   public void setUpMockCredentials() {
-    this.configurationBuilder = new SpannerConnectionConfiguration.Builder()
+    this.configurationBuilder = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setCredentials(this.mockCredentials);
   }
 
   @Test
   void missingInstanceNameTriggersException() {
-    Builder builder = new SpannerConnectionConfiguration.Builder()
+    Builder builder = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setProjectId("project1")
         .setDatabaseName("db")
         .setCredentials(NoCredentials.getInstance());
@@ -59,7 +62,8 @@ void missingInstanceNameTriggersException() {
 
   @Test
   void missingDatabaseNameTriggersException() {
-    Builder builder = new SpannerConnectionConfiguration.Builder()
+    Builder builder = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setProjectId("project1")
         .setInstanceName("an-instance")
         .setCredentials(NoCredentials.getInstance());
@@ -71,7 +75,8 @@ void missingDatabaseNameTriggersException() {
 
   @Test
   void missingProjectIdTriggersException() {
-    Builder builder = new SpannerConnectionConfiguration.Builder()
+    Builder builder = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setInstanceName("an-instance")
         .setDatabaseName("db")
         .setCredentials(NoCredentials.getInstance());

cloud-spanner-r2dbc/src/test/java/com/google/cloud/spanner/r2dbc/SpannerConnectionFactoryProviderTest.java

@@ -48,11 +48,13 @@
 import com.google.spanner.v1.StructType.Field;
 import com.google.spanner.v1.Type;
 import com.google.spanner.v1.TypeCode;
+import io.r2dbc.spi.Closeable;
 import io.r2dbc.spi.ConnectionFactories;
 import io.r2dbc.spi.ConnectionFactory;
 import io.r2dbc.spi.ConnectionFactoryOptions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
 
 /**
  * Unit test for {@link SpannerConnectionFactoryProvider}.
@@ -101,6 +103,8 @@ void testCreate() {
         this.spannerConnectionFactoryProvider.create(SPANNER_OPTIONS);
     assertThat(spannerConnectionFactory).isNotNull();
     assertThat(spannerConnectionFactory).isInstanceOf(SpannerClientLibraryConnectionFactory.class);
+
+    Flux.from(((Closeable) spannerConnectionFactory).close()).blockLast();
   }
 
   @Test
@@ -111,6 +115,8 @@ void testCreateFactoryWithOldDriverNameWillReturnCorrectFactory() {
     assertThat(spannerConnectionFactory)
         .isNotNull()
         .isInstanceOf(SpannerClientLibraryConnectionFactory.class);
+
+    Flux.from(((Closeable) spannerConnectionFactory).close()).blockLast();
   }
 
   @Test
@@ -121,6 +127,8 @@ void testCreateFactoryWithDriverNameWillReturnCorrectFactory() {
     assertThat(spannerConnectionFactory)
         .isNotNull()
         .isInstanceOf(SpannerClientLibraryConnectionFactory.class);
+
+    Flux.from(((Closeable) spannerConnectionFactory).close()).blockLast();
   }
 
   @Test
@@ -138,6 +146,8 @@ void testCreateFactoryWithUrl() {
     assertThat(spannerConnectionFactory)
         .isNotNull()
         .isInstanceOf(SpannerClientLibraryConnectionFactory.class);
+
+    Flux.from(((Closeable) spannerConnectionFactory).close()).blockLast();
   }
 
   @Test

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

@@ -45,6 +45,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import io.r2dbc.spi.ConnectionFactoryOptions;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import org.junit.jupiter.api.AfterEach;
@@ -74,7 +75,7 @@ class DatabaseClientReactiveAdapterTest {
   @BeforeEach
   void setup() throws Exception {
     this.config =
-        new SpannerConnectionConfiguration.Builder()
+        new SpannerConnectionConfiguration.Builder(mock(ConnectionFactoryOptions.class))
             .setFullyQualifiedDatabaseName("projects/p/instances/i/databases/d")
             .setCredentials(mock(GoogleCredentials.class))
             .build();
@@ -148,7 +149,8 @@ void testSameAutocommitNoop() {
 
   @Test
   void unsetQueryOptimizerResultsInDefaultQueryOptions() {
-    SpannerConnectionConfiguration config = new SpannerConnectionConfiguration.Builder()
+    SpannerConnectionConfiguration config = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setFullyQualifiedDatabaseName("projects/p/instances/i/databases/d")
         .setCredentials(mock(GoogleCredentials.class))
         .build();
@@ -160,7 +162,8 @@ void unsetQueryOptimizerResultsInDefaultQueryOptions() {
 
   @Test
   void queryOptimizerPropagatesToQueryOptions() {
-    SpannerConnectionConfiguration config = new SpannerConnectionConfiguration.Builder()
+    SpannerConnectionConfiguration config = new SpannerConnectionConfiguration.Builder(
+        mock(ConnectionFactoryOptions.class))
         .setFullyQualifiedDatabaseName("projects/p/instances/i/databases/d")
         .setCredentials(mock(GoogleCredentials.class))
         .setOptimizerVersion("2")