chore(spanner): fix mutation only case in rw mux with aborted error (#3571)

* chore(Spanner): fix mutation only case ibn rw using mux with aborted errors

* chore(Spanner): fix mockspanner

* chore(spanner): update logic for mutations only

* chore(spanner): lint fix

* chore(spanner): update logic to handle begin txn retry only for aborted case

* chore(spanner): add span for abort

Author: harshachinta

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

@@ -326,6 +326,19 @@ private void createTxnAsync(
               }
               res.set(null);
             } catch (ExecutionException e) {
+              SpannerException spannerException = SpannerExceptionFactory.asSpannerException(e);
+              if (spannerException.getErrorCode() == ErrorCode.ABORTED
+                  && session.getIsMultiplexed()
+                  && mutation != null) {
+                // Begin transaction can return ABORTED errors. This can only happen if it included
+                // a mutation key, which again means that this is a mutation-only transaction on a
+                // multiplexed session.
+                span.addAnnotation(
+                    "Transaction Creation Failed with ABORT. Retrying",
+                    e.getCause() == null ? e : e.getCause());
+                createTxnAsync(res, mutation);
+                return;
+              }
               span.addAnnotation(
                   "Transaction Creation Failed", e.getCause() == null ? e : e.getCause());
               res.setException(e.getCause() == null ? e : e.getCause());

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

@@ -1919,7 +1919,8 @@ private Transaction beginTransaction(
     }
     if (session.getMultiplexed()
         && options.getModeCase() == ModeCase.READ_WRITE
-        && mutationKey != null) {
+        && mutationKey != null
+        && mutationKey != com.google.spanner.v1.Mutation.getDefaultInstance()) {
       // Mutation only case in a read-write transaction.
       builder.setPrecommitToken(getTransactionPrecommitToken(transactionId));
     }
@@ -2023,6 +2024,14 @@ public void commit(CommitRequest request, StreamObserver<CommitResponse> respons
       return;
     }
     sessionLastUsed.put(session.getName(), Instant.now());
+    if (session.getMultiplexed()
+        && !request.hasPrecommitToken()
+        && !request.hasSingleUseTransaction()) {
+      throw Status.INVALID_ARGUMENT
+          .withDescription(
+              "A Commit request for a read-write transaction on a multiplexed session must specify a precommit token.")
+          .asRuntimeException();
+    }
     try {
       commitExecutionTime.simulateExecutionTime(exceptions, stickyGlobalExceptions, freezeLock);
       // Find or start a transaction

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

@@ -1294,6 +1294,204 @@ public void testMutationOnlyUsingAsyncTransactionManager() {
         request.getPrecommitToken().getPrecommitToken());
   }
 
+  private Spanner setupSpannerForAbortedBeginTransactionTests() {
+    // Force the BeginTransaction RPC to return Aborted the first time it is called. The exception
+    // is cleared after the first call, so the retry should succeed.
+    mockSpanner.setBeginTransactionExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+
+    return SpannerOptions.newBuilder()
+        .setProjectId("test-project")
+        .setChannelProvider(channelProvider)
+        .setCredentials(NoCredentials.getInstance())
+        .setSessionPoolOption(
+            SessionPoolOptions.newBuilder()
+                .setUseMultiplexedSession(true)
+                .setUseMultiplexedSessionForRW(true)
+                .setSkipVerifyingBeginTransactionForMuxRW(true)
+                .build())
+        .build()
+        .getService();
+  }
+
+  private void verifyMutationKeySetInBeginTransactionRequests(
+      List<BeginTransactionRequest> beginTransactionRequests) {
+    assertEquals(2, beginTransactionRequests.size());
+    // Verify the requests are executed using multiplexed sessions
+    for (BeginTransactionRequest request : beginTransactionRequests) {
+      assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
+      assertTrue(request.hasMutationKey());
+      assertTrue(request.getMutationKey().hasInsert());
+    }
+  }
+
+  private void verifyPreCommitTokenSetInCommitRequest(List<CommitRequest> commitRequests) {
+    assertEquals(1L, commitRequests.size());
+    for (CommitRequest request : commitRequests) {
+      assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed());
+      assertNotNull(request.getPrecommitToken());
+      assertEquals(
+          ByteString.copyFromUtf8("TransactionPrecommitToken"),
+          request.getPrecommitToken().getPrecommitToken());
+    }
+  }
+
+  // The following 4 tests validate mutation-only cases where the BeginTransaction RPC fails with an
+  // ABORTED or retryable error
+  @Test
+  public void testMutationOnlyCaseAbortedDuringBeginTransaction() {
+    // This test ensures that when a transaction containing only mutations is retried after an
+    // ABORT error in the BeginTransaction RPC:
+    // 1. The mutation key is correctly included in the BeginTransaction request.
+    // 2. The precommit token is properly set in the Commit request.
+    Spanner spanner = setupSpannerForAbortedBeginTransactionTests();
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+    client
+        .readWriteTransaction()
+        .run(
+            transaction -> {
+              Mutation mutation =
+                  Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build();
+              transaction.buffer(mutation);
+              return null;
+            });
+
+    // Verify that for mutation only case, a mutation key is set in BeginTransactionRequest.
+    List<BeginTransactionRequest> beginTransactionRequests =
+        mockSpanner.getRequestsOfType(BeginTransactionRequest.class);
+    verifyMutationKeySetInBeginTransactionRequests(beginTransactionRequests);
+
+    // Verify that the latest precommit token is set in the CommitRequest
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    verifyPreCommitTokenSetInCommitRequest(commitRequests);
+
+    spanner.close();
+  }
+
+  @Test
+  public void testMutationOnlyUsingTransactionManagerAbortedDuringBeginTransaction() {
+    // This test ensures that when a transaction containing only mutations is retried after an
+    // ABORT error in the BeginTransaction RPC:
+    // 1. The mutation key is correctly included in the BeginTransaction request.
+    // 2. The precommit token is properly set in the Commit request.
+    Spanner spanner = setupSpannerForAbortedBeginTransactionTests();
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+    try (TransactionManager manager = client.transactionManager()) {
+      TransactionContext transaction = manager.begin();
+      while (true) {
+        try {
+          Mutation mutation =
+              Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build();
+          transaction.buffer(mutation);
+          manager.commit();
+          assertNotNull(manager.getCommitTimestamp());
+          break;
+        } catch (AbortedException e) {
+          transaction = manager.resetForRetry();
+        }
+      }
+    }
+
+    // Verify that for mutation only case, a mutation key is set in BeginTransactionRequest.
+    List<BeginTransactionRequest> beginTransactionRequests =
+        mockSpanner.getRequestsOfType(BeginTransactionRequest.class);
+    verifyMutationKeySetInBeginTransactionRequests(beginTransactionRequests);
+
+    // Verify that the latest precommit token is set in the CommitRequest
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    verifyPreCommitTokenSetInCommitRequest(commitRequests);
+
+    spanner.close();
+  }
+
+  @Test
+  public void testMutationOnlyUsingAsyncRunnerAbortedDuringBeginTransaction() {
+    // This test ensures that when a transaction containing only mutations is retried after an
+    // ABORT error in the BeginTransaction RPC:
+    // 1. The mutation key is correctly included in the BeginTransaction request.
+    // 2. The precommit token is properly set in the Commit request.
+
+    Spanner spanner = setupSpannerForAbortedBeginTransactionTests();
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+    AsyncRunner runner = client.runAsync();
+    get(
+        runner.runAsync(
+            txn -> {
+              txn.buffer(
+                  Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build());
+              return ApiFutures.immediateFuture(null);
+            },
+            MoreExecutors.directExecutor()));
+
+    // Verify that for mutation only case, a mutation key is set in BeginTransactionRequest.
+    List<BeginTransactionRequest> beginTransactionRequests =
+        mockSpanner.getRequestsOfType(BeginTransactionRequest.class);
+    verifyMutationKeySetInBeginTransactionRequests(beginTransactionRequests);
+
+    // Verify that the latest precommit token is set in the CommitRequest
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    verifyPreCommitTokenSetInCommitRequest(commitRequests);
+
+    spanner.close();
+  }
+
+  @Test
+  public void testMutationOnlyUsingTransactionManagerAsyncAbortedDuringBeginTransaction()
+      throws Exception {
+    // This test verifies that in the case of mutations-only, when a transaction is retried after an
+    // ABORT in BeginTransaction RPC, the mutation key is correctly set in the BeginTransaction
+    // request
+    // and precommit token is set in Commit request.
+    Spanner spanner = setupSpannerForAbortedBeginTransactionTests();
+    DatabaseClientImpl client =
+        (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+
+    try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+      TransactionContextFuture transaction = manager.beginAsync();
+      while (true) {
+        CommitTimestampFuture commitTimestamp =
+            transaction
+                .then(
+                    (txn, input) -> {
+                      txn.buffer(
+                          Mutation.newInsertBuilder("FOO")
+                              .set("ID")
+                              .to(1L)
+                              .set("NAME")
+                              .to("Bar")
+                              .build());
+                      return ApiFutures.immediateFuture(null);
+                    },
+                    MoreExecutors.directExecutor())
+                .commitAsync();
+        try {
+          assertThat(commitTimestamp.get()).isNotNull();
+          break;
+        } catch (AbortedException e) {
+          transaction = manager.resetForRetryAsync();
+        }
+      }
+    }
+
+    // Verify that for mutation only case, a mutation key is set in BeginTransactionRequest.
+    List<BeginTransactionRequest> beginTransactionRequests =
+        mockSpanner.getRequestsOfType(BeginTransactionRequest.class);
+    verifyMutationKeySetInBeginTransactionRequests(beginTransactionRequests);
+
+    // Verify that the latest precommit token is set in the CommitRequest
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    verifyPreCommitTokenSetInCommitRequest(commitRequests);
+
+    spanner.close();
+  }
+
   // Tests the behavior of the server-side kill switch for read-write multiplexed sessions..
   @Test
   public void testInitialBeginTransactionWithRW_receivesUnimplemented_fallsBackToRegularSession() {