chore(spanner): update precommit token logic in mock spanner

Author: harshachinta

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

@@ -198,12 +198,15 @@ private static class PartialResultSetsIterator implements Iterator<PartialResult
     private boolean hasNext;
     private boolean first = true;
     private int currentRow = 0;
-    private boolean isMultiplexedSession = false;
+    private final boolean setPrecommitToken;
+    private final ByteString transactionId;
 
-    private PartialResultSetsIterator(ResultSet resultSet, boolean isMultiplexedSession) {
+    private PartialResultSetsIterator(
+        ResultSet resultSet, boolean setPrecommitToken, ByteString transactionId) {
       this.resultSet = resultSet;
       this.hasNext = true;
-      this.isMultiplexedSession = isMultiplexedSession;
+      this.setPrecommitToken = setPrecommitToken;
+      this.transactionId = transactionId;
     }
 
     @Override
@@ -230,8 +233,8 @@ public PartialResultSet next() {
       }
       builder.setResumeToken(ByteString.copyFromUtf8(String.format("%09d", currentRow)));
       hasNext = currentRow < resultSet.getRowsCount();
-      if (this.isMultiplexedSession) {
-        builder.setPrecommitToken(getPartialResultSetPrecommitToken());
+      if (this.setPrecommitToken) {
+        builder.setPrecommitToken(getPartialResultSetPrecommitToken(this.transactionId));
       }
       return builder.build();
     }
@@ -606,6 +609,10 @@ private static void checkStreamException(
   private ConcurrentMap<String, AtomicLong> transactionCounters = new ConcurrentHashMap<>();
   private ConcurrentMap<String, List<ByteString>> partitionTokens = new ConcurrentHashMap<>();
   private ConcurrentMap<ByteString, Instant> transactionLastUsed = new ConcurrentHashMap<>();
+
+  // Stores latest sequence number required for precommit token.
+  private static ConcurrentMap<ByteString, AtomicInteger> transactionSequenceNo =
+      new ConcurrentHashMap<>();
   private int maxNumSessionsInOneBatch = 100;
   private int maxTotalSessions = Integer.MAX_VALUE;
   private Iterable<BatchWriteResponse> batchWriteResult = new ArrayList<>();
@@ -1056,8 +1063,8 @@ public void executeSql(ExecuteSqlRequest request, StreamObserver<ResultSet> resp
                                     ? Transaction.getDefaultInstance()
                                     : Transaction.newBuilder().setId(transactionId).build())
                             .build());
-            if (session.getMultiplexed()) {
-              resultSetBuilder.setPrecommitToken(getResultSetPrecommitToken());
+            if (session.getMultiplexed() && isReadWriteTransaction(transactionId)) {
+              resultSetBuilder.setPrecommitToken(getResultSetPrecommitToken(transactionId));
             }
             responseObserver.onNext(resultSetBuilder.build());
           }
@@ -1095,8 +1102,8 @@ private void returnResultSet(
     }
     ResultSet.Builder resultSetBuilder = resultSet.toBuilder();
     resultSetBuilder.setMetadata(metadata);
-    if (session.getMultiplexed()) {
-      resultSetBuilder.setPrecommitToken(getResultSetPrecommitToken());
+    if (session.getMultiplexed() && isReadWriteTransaction(transactionId)) {
+      resultSetBuilder.setPrecommitToken(getResultSetPrecommitToken(transactionId));
     }
     resultSet = resultSetBuilder.build();
     responseObserver.onNext(resultSet);
@@ -1193,8 +1200,8 @@ public void executeBatchDml(
                 .build());
       }
       builder.setStatus(status);
-      if (session.getMultiplexed()) {
-        builder.setPrecommitToken(getExecuteBatchDmlResponsePrecommitToken());
+      if (session.getMultiplexed() && isReadWriteTransaction(transactionId)) {
+        builder.setPrecommitToken(getExecuteBatchDmlResponsePrecommitToken(transactionId));
       }
       responseObserver.onNext(builder.build());
       responseObserver.onCompleted();
@@ -1726,7 +1733,10 @@ private void returnPartialResultSet(
     }
     resultSet = resultSet.toBuilder().setMetadata(metadata).build();
     PartialResultSetsIterator iterator =
-        new PartialResultSetsIterator(resultSet, isMultiplexedSession);
+        new PartialResultSetsIterator(
+            resultSet,
+            isMultiplexedSession && isReadWriteTransaction(transactionId),
+            transactionId);
     long index = 0L;
     while (iterator.hasNext()) {
       SimulatedExecutionTime.checkStreamException(
@@ -2060,6 +2070,7 @@ private void commitTransaction(ByteString transactionId) {
     transactions.remove(transactionId);
     isPartitionedDmlTransaction.remove(transactionId);
     transactionLastUsed.remove(transactionId);
+    transactionSequenceNo.remove(transactionId);
   }
 
   @Override
@@ -2091,13 +2102,15 @@ void rollbackTransaction(ByteString transactionId) {
     transactions.remove(transactionId);
     isPartitionedDmlTransaction.remove(transactionId);
     transactionLastUsed.remove(transactionId);
+    transactionSequenceNo.remove(transactionId);
   }
 
   void markAbortedTransaction(ByteString transactionId) {
     abortedTransactions.put(transactionId, Boolean.TRUE);
     transactions.remove(transactionId);
     isPartitionedDmlTransaction.remove(transactionId);
     transactionLastUsed.remove(transactionId);
+    transactionSequenceNo.remove(transactionId);
   }
 
   @Override
@@ -2302,6 +2315,7 @@ public void reset() {
     transactionCounters = new ConcurrentHashMap<>();
     partitionTokens = new ConcurrentHashMap<>();
     transactionLastUsed = new ConcurrentHashMap<>();
+    transactionSequenceNo = new ConcurrentHashMap<>();
 
     numSessionsCreated.set(0);
     stickyGlobalExceptions = false;
@@ -2474,22 +2488,31 @@ Session getSession(String name) {
     return null;
   }
 
-  static MultiplexedSessionPrecommitToken getResultSetPrecommitToken() {
-    return getPrecommitToken("ResultSetPrecommitToken", 1);
+  static MultiplexedSessionPrecommitToken getResultSetPrecommitToken(ByteString transactionId) {
+    return getPrecommitToken("ResultSetPrecommitToken", transactionId);
   }
 
-  static MultiplexedSessionPrecommitToken getPartialResultSetPrecommitToken() {
-    return getPrecommitToken("PartialResultSetPrecommitToken", 3);
+  static MultiplexedSessionPrecommitToken getPartialResultSetPrecommitToken(
+      ByteString transactionId) {
+    return getPrecommitToken("PartialResultSetPrecommitToken", transactionId);
   }
 
-  static MultiplexedSessionPrecommitToken getExecuteBatchDmlResponsePrecommitToken() {
-    return getPrecommitToken("ExecuteBatchDmlResponsePrecommitToken", 2);
+  static MultiplexedSessionPrecommitToken getExecuteBatchDmlResponsePrecommitToken(
+      ByteString transactionId) {
+    return getPrecommitToken("ExecuteBatchDmlResponsePrecommitToken", transactionId);
   }
 
-  static MultiplexedSessionPrecommitToken getPrecommitToken(String value, int seqNo) {
+  static MultiplexedSessionPrecommitToken getPrecommitToken(
+      String value, ByteString transactionId) {
+    if (!transactionSequenceNo.containsKey(transactionId)) {
+      transactionSequenceNo.put(transactionId, new AtomicInteger(0));
+    }
+
+    // Generates an incrementing sequence number
+    int seqNum = transactionSequenceNo.get(transactionId).addAndGet(1);
     return MultiplexedSessionPrecommitToken.newBuilder()
         .setPrecommitToken(ByteString.copyFromUtf8(value))
-        .setSeqNum(seqNo)
+        .setSeqNum(seqNum)
         .build();
   }
 }

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

@@ -858,42 +858,37 @@ public void testPrecommitTokenForPartialResultSet() {
   @Test
   public void testTxnTracksPrecommitTokenWithLatestSeqNo() {
     // This test ensures that the read-write transaction tracks the precommit token with the
-    // highest sequence number (in this case, PartialResultSetPrecommitToken) and sets it in the
-    // CommitRequest.
-    // ResultSetPrecommitToken                -> Seq no 1
-    // ExecuteBatchDmlResponsePrecommitToken  -> Seq no 2
-    // PartialResultSetPrecommitToken         -> Seq no 3
-
+    // highest sequence number and sets it in the CommitRequest.
     DatabaseClientImpl client =
         (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
 
     client
         .readWriteTransaction()
         .run(
             transaction -> {
-              // Returns a ResultSet containing the precommit token (ResultSetPrecommitToken) with
-              // Seq no 1
+              // Returns a ResultSet containing the precommit token (ResultSetPrecommitToken)
               transaction.executeUpdate(UPDATE_STATEMENT);
 
               // Returns a PartialResultSet containing the precommit token
-              // (PartialResultSetPrecommitToken) with Seq no 3
+              // (PartialResultSetPrecommitToken)
               ResultSet resultSet = transaction.executeQuery(STATEMENT);
               //noinspection StatementWithEmptyBody
               while (resultSet.next()) {
                 // ignore
               }
 
-              // Returns a ExecuteBatchDmlResponse containing the precommit token
-              // (ExecuteBatchDmlResponsePrecommitToken) with Seq no 2
+              // Returns an ExecuteBatchDmlResponse containing the precommit token
+              // (ExecuteBatchDmlResponsePrecommitToken).
+              // Since this is the last request received by the mock Spanner, it should be the most
+              // recent precommit token tracked by the transaction context.
               transaction.batchUpdate(Lists.newArrayList(UPDATE_STATEMENT));
 
               // Verify that the latest precommit token with highest sequence number is tracked in
-              // the
-              // transaction context.
+              // the transaction context.
               TransactionContextImpl impl = (TransactionContextImpl) transaction;
               assertNotNull(impl.getLatestPrecommitToken());
               assertEquals(
-                  ByteString.copyFromUtf8("PartialResultSetPrecommitToken"),
+                  ByteString.copyFromUtf8("ExecuteBatchDmlResponsePrecommitToken"),
                   impl.getLatestPrecommitToken().getPrecommitToken());
               return null;
             });
@@ -904,7 +899,7 @@ public void testTxnTracksPrecommitTokenWithLatestSeqNo() {
     assertTrue(mockSpanner.getSession(commitRequests.get(0).getSession()).getMultiplexed());
     assertNotNull(commitRequests.get(0).getPrecommitToken());
     assertEquals(
-        ByteString.copyFromUtf8("PartialResultSetPrecommitToken"),
+        ByteString.copyFromUtf8("ExecuteBatchDmlResponsePrecommitToken"),
         commitRequests.get(0).getPrecommitToken().getPrecommitToken());
   }