chore(spanner): tracking precommit token from response

Author: harshachinta

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;
@@ -171,6 +172,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();
 
@@ -625,6 +631,24 @@ 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;
+          System.out.println("Updating precommit token to " + this.latestPrecommitToken);
+          txnLogger.log(Level.ALL, "Updating precommit token to " + this.latestPrecommitToken);
+        }
+      }
+    }
+
     @Nullable
     String getTransactionTag() {
       if (this.options.hasTag()) {
@@ -633,6 +657,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);
@@ -811,6 +842,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(
@@ -885,6 +919,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.
@@ -940,6 +977,11 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... update
             }
           }
 
+          // TODO(sriharshach): check if we need to get precommit_token from response.getResultSets
+          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) {
@@ -1004,6 +1046,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) {