chore(spanner): add mock spanner tests for R/W txn with multiplexed sessions

Author: harshachinta

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

@@ -40,17 +40,6 @@ public String getDatabaseRole() {
     throw new UnsupportedOperationException();
   }
 
-  @Override
-  public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public CommitResponse writeWithOptions(Iterable<Mutation> mutations, TransactionOption... options)
-      throws SpannerException {
-    throw new UnsupportedOperationException();
-  }
-
   @Override
   public Timestamp writeAtLeastOnce(Iterable<Mutation> mutations) throws SpannerException {
     return writeAtLeastOnceWithOptions(mutations).getCommitTimestamp();
@@ -63,16 +52,6 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
     throw new UnsupportedOperationException();
   }
 
-  @Override
-  public TransactionRunner readWriteTransaction(TransactionOption... options) {
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
-  public TransactionManager transactionManager(TransactionOption... options) {
-    throw new UnsupportedOperationException();
-  }
-
   @Override
   public AsyncRunner runAsync(TransactionOption... options) {
     throw new UnsupportedOperationException();

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

@@ -20,6 +20,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction;
 import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction;
 import com.google.cloud.spanner.Options.TransactionOption;
@@ -136,6 +137,19 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  /**
+   * This is a blocking method, as the interface that it implements is also defined as a blocking
+   * method.
+   */
+  @Override
+  public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
+    SessionReference sessionReference = getSessionReference();
+    try (MultiplexedSessionTransaction transaction =
+        new MultiplexedSessionTransaction(client, span, sessionReference, NO_CHANNEL_HINT, false)) {
+      return transaction.write(mutations);
+    }
+  }
+
   /**
    * This is a blocking method, as the interface that it implements is also defined as a blocking
    * method.

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

@@ -21,6 +21,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
+import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.SessionClient.SessionConsumer;
 import com.google.cloud.spanner.SpannerException.ResourceNotFoundException;
@@ -367,6 +368,11 @@ private int getSingleUseChannelHint() {
     }
   }
 
+  @Override
+  public Timestamp write(Iterable<Mutation> mutations) throws SpannerException {
+    return createMultiplexedSessionTransaction(false).write(mutations);
+  }
+
   @Override
   public CommitResponse writeWithOptions(
       final Iterable<Mutation> mutations, final TransactionOption... options)

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

@@ -822,6 +822,8 @@ Builder setUseMultiplexedSessionBlindWrite(boolean useMultiplexedSessionBlindWri
      * Sets whether the client should use multiplexed session for R/W operations or not. This method
      * is intentionally package-private and intended for internal use.
      */
+    @InternalApi
+    @VisibleForTesting
     Builder setUseMultiplexedSessionForRW(boolean useMultiplexedSessionForRW) {
       this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
       return this;

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

@@ -71,6 +71,7 @@ public void createSpannerInstance() {
                 SessionPoolOptions.newBuilder()
                     .setUseMultiplexedSession(true)
                     .setUseMultiplexedSessionBlindWrite(true)
+                    .setUseMultiplexedSessionForRW(true)
                     // Set the maintainer to loop once every 1ms
                     .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMillis(1L))
                     // Set multiplexed sessions to be replaced once every 1ms
@@ -467,6 +468,148 @@ public void testWriteAtLeastOnceWithExcludeTxnFromChangeStreams() {
     assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
   }
 
+  @Test
+  public void testReadWriteTransactionUsingTransactionRunner() {
+    // Queries executed within a R/W transaction via TransactionRunner should use a multiplexed
+    // session.
+    // During a retry (due to an ABORTED error), the transaction should use the same multiplexed
+    // session as before, assuming the maintainer hasn't run in the meantime.
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+    // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+    // after the first call, so the retry should succeed.
+    mockSpanner.setCommitExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+
+    client
+        .readWriteTransaction()
+        .run(
+            transaction -> {
+              try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+                //noinspection StatementWithEmptyBody
+                while (resultSet.next()) {
+                  // ignore
+                }
+              }
+              return null;
+            });
+
+    List<ExecuteSqlRequest> executeSqlRequests =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, executeSqlRequests.size());
+    assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession());
+
+    // 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());
+    // TODO: fix this
+    // assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+  }
+
+  @Test
+  public void testReadWriteTransactionUsingTransactionManager() {
+    // Queries executed within a R/W transaction via TransactionManager should use a multiplexed
+    // session.
+    // During a retry (due to an ABORTED error), the transaction should use the same multiplexed
+    // session as before, assuming the maintainer hasn't run in the meantime.
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+    // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+    // after the first call, so the retry should succeed.
+    mockSpanner.setCommitExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+
+    try (TransactionManager manager = client.transactionManager()) {
+      TransactionContext transaction = manager.begin();
+      while (true) {
+        try {
+          try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) {
+            //noinspection StatementWithEmptyBody
+            while (resultSet.next()) {
+              // ignore
+            }
+          }
+          manager.commit();
+          assertNotNull(manager.getCommitTimestamp());
+          break;
+        } catch (AbortedException e) {
+          transaction = manager.resetForRetry();
+        }
+      }
+    }
+
+    List<ExecuteSqlRequest> executeSqlRequests =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(2, executeSqlRequests.size());
+    assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession());
+
+    // 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());
+    // TODO: fix this
+    // assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+  }
+
+  @Test
+  public void testMutationUsingWrite() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+    // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+    // after the first call, so the retry should succeed.
+    mockSpanner.setCommitExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+    Timestamp timestamp =
+        client.write(
+            Collections.singletonList(
+                Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()));
+    assertNotNull(timestamp);
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(2, commitRequests.size());
+    for (CommitRequest request : commitRequests) {
+      assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
+    }
+
+    assertNotNull(client.multiplexedSessionDatabaseClient);
+    assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get());
+    // assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get());
+  }
+
+  @Test
+  public void testMutationUsingWriteWithOptions() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+    CommitResponse response =
+        client.writeWithOptions(
+            Collections.singletonList(
+                Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()),
+            Options.tag("app=spanner,env=test"));
+    assertNotNull(response);
+    assertNotNull(response.getCommitTimestamp());
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertThat(commitRequests).hasSize(1);
+    CommitRequest commit = commitRequests.get(0);
+    assertNotNull(commit.getRequestOptions());
+    assertThat(commit.getRequestOptions().getTransactionTag()).isEqualTo("app=spanner,env=test");
+    assertTrue(mockSpanner.getSession(commit.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 =