chore(spanner): track precommit token for R/W multiplexed session

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

@@ -48,6 +48,7 @@
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
+import com.google.spanner.v1.MultiplexedSessionPrecommitToken;
 import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.ReadRequest;
 import com.google.spanner.v1.RequestOptions;
@@ -893,6 +894,13 @@ public void onDone(boolean withBeginTransaction) {
     this.session.onReadDone();
   }
 
+  /**
+   * For transactions other than read-write, the MultiplexedSessionPrecommitToken will not be
+   * present in the RPC response. In such cases, this method will be a no-op.
+   */
+  @Override
+  public void onPrecommitToken(MultiplexedSessionPrecommitToken token) {}
+
   private ResultSet readInternal(
       String table,
       @Nullable String index,

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

@@ -27,6 +27,7 @@
 import com.google.protobuf.ListValue;
 import com.google.protobuf.ProtocolMessageEnum;
 import com.google.protobuf.Value.KindCase;
+import com.google.spanner.v1.MultiplexedSessionPrecommitToken;
 import com.google.spanner.v1.Transaction;
 import java.io.IOException;
 import java.io.Serializable;
@@ -57,6 +58,12 @@ void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId)
 
     /** Called when the read finishes normally. */
     void onDone(boolean withBeginTransaction);
+
+    /**
+     * Called when the RPC response contains a MultiplexedSessionPrecommitToken. A precommit token
+     * will be included if the read-write transaction is executed on a multiplexed session.
+     */
+    void onPrecommitToken(MultiplexedSessionPrecommitToken token);
   }
 
   static final class LazyByteArray implements Serializable {

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

@@ -47,7 +47,7 @@ class GrpcResultSet extends AbstractResultSet<List<Object>> implements ProtobufR
 
   GrpcResultSet(
       CloseableIterator<PartialResultSet> iterator, Listener listener, DecodeMode decodeMode) {
-    this.iterator = new GrpcValueIterator(iterator);
+    this.iterator = new GrpcValueIterator(iterator, listener);
     this.listener = listener;
     this.decodeMode = decodeMode;
   }

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

@@ -20,6 +20,7 @@
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.cloud.spanner.AbstractResultSet.CloseableIterator;
+import com.google.cloud.spanner.AbstractResultSet.Listener;
 import com.google.common.collect.AbstractIterator;
 import com.google.protobuf.ListValue;
 import com.google.protobuf.Value.KindCase;
@@ -44,9 +45,11 @@ private enum StreamValue {
   private PartialResultSet current;
   private int pos;
   private ResultSetStats statistics;
+  private final Listener listener;
 
-  GrpcValueIterator(CloseableIterator<PartialResultSet> stream) {
+  GrpcValueIterator(CloseableIterator<PartialResultSet> stream, Listener listener) {
     this.stream = stream;
+    this.listener = listener;
   }
 
   @SuppressWarnings("unchecked")
@@ -154,6 +157,10 @@ private boolean ensureReady(StreamValue requiredValue) throws SpannerException {
               ErrorCode.INTERNAL, "Invalid type metadata: " + e.getMessage(), e);
         }
       }
+      // collect the precommit token from each PartialResultSet
+      if (current.hasPrecommitToken()) {
+        listener.onPrecommitToken(current.getPrecommitToken());
+      }
       if (current.hasStats()) {
         statistics = current.getStats();
       }

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

@@ -46,6 +46,7 @@
 import com.google.spanner.v1.ExecuteBatchDmlResponse;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
+import com.google.spanner.v1.MultiplexedSessionPrecommitToken;
 import com.google.spanner.v1.RequestOptions;
 import com.google.spanner.v1.ResultSet;
 import com.google.spanner.v1.ResultSetStats;
@@ -179,6 +180,11 @@ public void removeListener(Runnable listener) {
     @GuardedBy("committingLock")
     private volatile boolean committing;
 
+    private final Object precommitTokenLock = new Object();
+
+    @GuardedBy("precommitTokenLock")
+    private MultiplexedSessionPrecommitToken latestPrecommitToken;
+
     @GuardedBy("lock")
     private volatile SettableApiFuture<Void> finishedAsyncOperations = SettableApiFuture.create();
 
@@ -439,6 +445,10 @@ public void run() {
             }
             requestBuilder.setRequestOptions(requestOptionsBuilder.build());
           }
+          if (session.getIsMultiplexed() && getLatestPrecommitToken() != null) {
+            // Set the precommit token in the CommitRequest for multiplexed sessions.
+            requestBuilder.setPrecommitToken(getLatestPrecommitToken());
+          }
           final CommitRequest commitRequest = requestBuilder.build();
           span.addAnnotation("Starting Commit");
           final ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
@@ -643,6 +653,25 @@ public void onTransactionMetadata(Transaction transaction, boolean shouldInclude
       }
     }
 
+    /**
+     * In read-write transactions, the precommit token with the highest sequence number from this
+     * transaction attempt will be tracked and included in the
+     * [Commit][google.spanner.v1.Spanner.Commit] request for the transaction.
+     */
+    @Override
+    public void onPrecommitToken(MultiplexedSessionPrecommitToken token) {
+      if (token == null) {
+        return;
+      }
+      synchronized (precommitTokenLock) {
+        if (this.latestPrecommitToken == null
+            || token.getSeqNum() > this.latestPrecommitToken.getSeqNum()) {
+          this.latestPrecommitToken = token;
+          txnLogger.log(Level.FINE, "Updating precommit token to " + this.latestPrecommitToken);
+        }
+      }
+    }
+
     @Nullable
     String getTransactionTag() {
       if (this.options.hasTag()) {
@@ -651,6 +680,13 @@ String getTransactionTag() {
       return null;
     }
 
+    @Nullable
+    MultiplexedSessionPrecommitToken getLatestPrecommitToken() {
+      synchronized (precommitTokenLock) {
+        return this.latestPrecommitToken;
+      }
+    }
+
     @Override
     public SpannerException onError(SpannerException e, boolean withBeginTransaction) {
       e = super.onError(e, withBeginTransaction);
@@ -829,6 +865,9 @@ private ResultSet internalExecuteUpdate(
           throw new IllegalArgumentException(
               "DML response missing stats possibly due to non-DML statement as input");
         }
+        if (resultSet.hasPrecommitToken()) {
+          onPrecommitToken(resultSet.getPrecommitToken());
+        }
         return resultSet;
       } catch (Throwable t) {
         throw onError(
@@ -903,6 +942,9 @@ public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... u
                       resultSet.get().getMetadata().getTransaction(),
                       builder.getTransaction().hasBegin());
                 }
+                if (resultSet.get().hasPrecommitToken()) {
+                  onPrecommitToken(resultSet.get().getPrecommitToken());
+                }
               } catch (Throwable e) {
                 // Ignore this error here as it is handled by the future that is returned by the
                 // executeUpdateAsync method.
@@ -958,6 +1000,10 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... update
             }
           }
 
+          if (response.hasPrecommitToken()) {
+            onPrecommitToken(response.getPrecommitToken());
+          }
+
           // If one of the DML statements was aborted, we should throw an aborted exception.
           // In all other cases, we should throw a BatchUpdateException.
           if (response.getStatus().getCode() == Code.ABORTED_VALUE) {
@@ -1022,6 +1068,9 @@ public ApiFuture<long[]> batchUpdateAsync(
                           builder.getTransaction().hasBegin());
                     }
                   }
+                  if (batchDmlResponse.hasPrecommitToken()) {
+                    onPrecommitToken(batchDmlResponse.getPrecommitToken());
+                  }
                   // If one of the DML statements was aborted, we should throw an aborted exception.
                   // In all other cases, we should throw a BatchUpdateException.
                   if (batchDmlResponse.getStatus().getCode() == Code.ABORTED_VALUE) {

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

@@ -36,6 +36,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.ByteString;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
+import com.google.spanner.v1.MultiplexedSessionPrecommitToken;
 import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.QueryPlan;
 import com.google.spanner.v1.ResultSetMetadata;
@@ -77,6 +78,9 @@ public SpannerException onError(SpannerException e, boolean withBeginTransaction
 
     @Override
     public void onDone(boolean withBeginTransaction) {}
+
+    @Override
+    public void onPrecommitToken(MultiplexedSessionPrecommitToken token) {}
   }
 
   @Before