chore: wire up x-goog-spanner-request-id to all

Wires up x-goog-spanner-request-id for piecemeal addition per
gRPC method, starting with BatchCreateSessions. This change
involves creating TransactionOption, UpdateOption variants that
allow holding the request id.

Updates googleapis#3537

Author: odeke-em

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

@@ -458,8 +458,11 @@ void initTransaction() {
 
     private void initTransactionInternal(BeginTransactionRequest request) {
       try {
+        XGoogSpannerRequestId reqId =
+            session.getRequestIdCreator().nextRequestId(1 /*TODO: retrieve channelId*/, 1);
         Transaction transaction =
-            rpc.beginTransaction(request, getTransactionChannelHint(), isRouteToLeader());
+            rpc.beginTransaction(
+                request, reqId.withOptions(getTransactionChannelHint()), isRouteToLeader());
         if (!transaction.hasReadTimestamp()) {
           throw SpannerExceptionFactory.newSpannerException(
               ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field");
@@ -803,7 +806,8 @@ ResultSet executeQueryInternalWithOptions(
             tracer.createStatementAttributes(statement, options),
             session.getErrorHandler(),
             rpc.getExecuteQueryRetrySettings(),
-            rpc.getExecuteQueryRetryableCodes()) {
+            rpc.getExecuteQueryRetryableCodes(),
+            session.getRequestIdCreator()) {
           @Override
           CloseableIterator<PartialResultSet> startStream(
               @Nullable ByteString resumeToken,
@@ -826,11 +830,13 @@ CloseableIterator<PartialResultSet> startStream(
             if (selector != null) {
               request.setTransaction(selector);
             }
+
+            this.incrementXGoogRequestIdAttempt();
             SpannerRpc.StreamingCall call =
                 rpc.executeQuery(
                     request.build(),
                     stream.consumer(),
-                    getTransactionChannelHint(),
+                    this.xGoogRequestId.withOptions(getTransactionChannelHint()),
                     isRouteToLeader());
             session.markUsed(clock.instant());
             stream.setCall(call, request.getTransaction().hasBegin());
@@ -1008,7 +1014,8 @@ ResultSet readInternalWithOptions(
             tracer.createTableAttributes(table, readOptions),
             session.getErrorHandler(),
             rpc.getReadRetrySettings(),
-            rpc.getReadRetryableCodes()) {
+            rpc.getReadRetryableCodes(),
+            session.getRequestIdCreator()) {
           @Override
           CloseableIterator<PartialResultSet> startStream(
               @Nullable ByteString resumeToken,
@@ -1029,11 +1036,12 @@ CloseableIterator<PartialResultSet> startStream(
               builder.setTransaction(selector);
             }
             builder.setRequestOptions(buildRequestOptions(readOptions));
+            this.incrementXGoogRequestIdAttempt();
             SpannerRpc.StreamingCall call =
                 rpc.read(
                     builder.build(),
                     stream.consumer(),
-                    getTransactionChannelHint(),
+                    this.xGoogRequestId.withOptions(getTransactionChannelHint()),
                     isRouteToLeader());
             session.markUsed(clock.instant());
             stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());

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

@@ -315,7 +315,9 @@ private List<Partition> partitionQuery(
 
       final PartitionQueryRequest request = builder.build();
       try {
-        PartitionResponse response = rpc.partitionQuery(request, options);
+        XGoogSpannerRequestId reqId =
+            this.session.requestIdCreator.nextRequestId(1 /* channelId */, 0);
+        PartitionResponse response = rpc.partitionQuery(request, reqId.withOptions(options));
         ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
         for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
           Partition partition =

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

@@ -27,6 +27,10 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.spanner.v1.BatchWriteResponse;
 import io.opentelemetry.api.common.Attributes;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 
 class DatabaseClientImpl implements DatabaseClient {
@@ -40,6 +44,8 @@ class DatabaseClientImpl implements DatabaseClient {
   @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
   @VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
   @VisibleForTesting final boolean useMultiplexedSessionForRW;
+  private final int dbId;
+  private final AtomicInteger nthRequest;
 
   final boolean useMultiplexedSessionBlindWrite;
 
@@ -86,6 +92,18 @@ class DatabaseClientImpl implements DatabaseClient {
     this.tracer = tracer;
     this.useMultiplexedSessionForRW = useMultiplexedSessionForRW;
     this.commonAttributes = commonAttributes;
+
+    this.dbId = this.dbIdFromClientId(this.clientId);
+    this.nthRequest = new AtomicInteger(0);
+  }
+
+  private int dbIdFromClientId(String clientId) {
+    int i = clientId.indexOf("-");
+    String strWithValue = clientId.substring(i + 1);
+    if (Objects.equals(strWithValue, "")) {
+      strWithValue = "0";
+    }
+    return Integer.parseInt(strWithValue);
   }
 
   @VisibleForTesting
@@ -183,8 +201,20 @@ public CommitResponse writeAtLeastOnceWithOptions(
         return getMultiplexedSessionDatabaseClient()
             .writeAtLeastOnceWithOptions(mutations, options);
       }
+
+      int nthRequest = this.nextNthRequest();
+      int channelId = 1; /* TODO: infer the channelId from the gRPC channel of the session */
+      XGoogSpannerRequestId reqId = XGoogSpannerRequestId.of(this.dbId, channelId, nthRequest, 0);
+
       return runWithSessionRetry(
-          session -> session.writeAtLeastOnceWithOptions(mutations, options));
+          (session) -> {
+            reqId.incrementAttempt();
+            // TODO: Update the channelId depending on the session that is inferred.
+            ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
+            allOptions.add(new Options.RequestIdOption(reqId));
+            return session.writeAtLeastOnceWithOptions(
+                mutations, allOptions.toArray(new TransactionOption[0]));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -193,16 +223,38 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  private int nextNthRequest() {
+    return this.nthRequest.incrementAndGet();
+  }
+
   @Override
   public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
       final Iterable<MutationGroup> mutationGroups, final TransactionOption... options)
       throws SpannerException {
     ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options);
     try (IScope s = tracer.withSpan(span)) {
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, 1 /*TODO:channelId*/, this.nextNthRequest(), 0);
+      System.out.println("\033[35mbatchWriteAtLeastOnceReq: " + reqId.toString() + "\033[00m");
       if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
-        return getMultiplexedSessionDatabaseClient().batchWriteAtLeastOnce(mutationGroups, options);
+        reqId.incrementAttempt();
+        ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
+        allOptions.add(new Options.RequestIdOption(reqId));
+        return getMultiplexedSessionDatabaseClient()
+            .batchWriteAtLeastOnce(mutationGroups, allOptions.toArray(new TransactionOption[0]));
       }
-      return runWithSessionRetry(session -> session.batchWriteAtLeastOnce(mutationGroups, options));
+      return runWithSessionRetry(
+          (session) -> {
+            reqId.incrementAttempt();
+            ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
+            System.out.println(
+                "\033[35mbatchWriteAtLeastOnce:: session.class: "
+                    + session.getClass()
+                    + "\033[00m");
+            allOptions.add(new Options.RequestIdOption(reqId));
+            return session.batchWriteAtLeastOnce(
+                mutationGroups, allOptions.toArray(new TransactionOption[0]));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -350,7 +402,15 @@ private long executePartitionedUpdateWithPooledSession(
       final Statement stmt, final UpdateOption... options) {
     ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION, commonAttributes);
     try (IScope s = tracer.withSpan(span)) {
-      return runWithSessionRetry(session -> session.executePartitionedUpdate(stmt, options));
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, 1 /*TODO: channelId*/, this.nextNthRequest(), 0);
+      return runWithSessionRetry(
+          (session) -> {
+            reqId.incrementAttempt();
+            ArrayList<TransactionOption> allOptions = new ArrayList(Arrays.asList(options));
+            allOptions.add(new Options.RequestIdOption(reqId));
+            return session.executePartitionedUpdate(stmt, allOptions.toArray(new UpdateOption[0]));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       span.end();

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

@@ -535,6 +535,7 @@ void appendToOptions(Options options) {
   private RpcLockHint lockHint;
   private Boolean lastStatement;
   private IsolationLevel isolationLevel;
+  private XGoogSpannerRequestId reqId;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -591,6 +592,14 @@ String pageToken() {
     return pageToken;
   }
 
+  boolean hasReqId() {
+    return reqId != null;
+  }
+
+  XGoogSpannerRequestId reqId() {
+    return reqId;
+  }
+
   boolean hasFilter() {
     return filter != null;
   }
@@ -1052,4 +1061,30 @@ public boolean equals(Object o) {
       return o instanceof LastStatementUpdateOption;
     }
   }
+
+  static final class RequestIdOption extends InternalOption
+      implements TransactionOption, UpdateOption {
+    private final XGoogSpannerRequestId reqId;
+
+    RequestIdOption(XGoogSpannerRequestId reqId) {
+      this.reqId = reqId;
+    }
+
+    @Override
+    void appendToOptions(Options options) {
+      options.reqId = this.reqId;
+    }
+
+    @Override
+    public int hashCode() {
+      return RequestIdOption.class.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      // TODO: Examine why the precedent for LastStatementUpdateOption
+      // does not check against the actual value.
+      return o instanceof RequestIdOption;
+    }
+  }
 }

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

@@ -81,15 +81,20 @@ long executeStreamingPartitionedUpdate(
     Stopwatch stopwatch = Stopwatch.createStarted(ticker);
     Options options = Options.fromUpdateOptions(updateOptions);
 
+    XGoogSpannerRequestId reqId =
+        session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 0);
+
     try {
       ExecuteSqlRequest request = newTransactionRequestFrom(statement, options);
 
       while (true) {
+        reqId.incrementAttempt();
         final Duration remainingTimeout = tryUpdateTimeout(timeout, stopwatch);
 
         try {
           ServerStream<PartialResultSet> stream =
-              rpc.executeStreamingPartitionedDml(request, session.getOptions(), remainingTimeout);
+              rpc.executeStreamingPartitionedDml(
+                  request, reqId.withOptions(session.getOptions()), remainingTimeout);
 
           for (PartialResultSet rs : stream) {
             if (rs.getResumeToken() != null && !rs.getResumeToken().isEmpty()) {
@@ -113,12 +118,17 @@ long executeStreamingPartitionedUpdate(
           LOGGER.log(
               Level.FINER, "Retrying PartitionedDml transaction after InternalException - EOS", e);
           request = resumeOrRestartRequest(resumeToken, statement, request, options);
+          if (resumeToken.isEmpty()) {
+            // Create a new xGoogSpannerRequestId.
+            reqId = session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 0);
+          }
         } catch (AbortedException e) {
           LOGGER.log(Level.FINER, "Retrying PartitionedDml transaction after AbortedException", e);
           resumeToken = ByteString.EMPTY;
           foundStats = false;
           updateCount = 0L;
           request = newTransactionRequestFrom(statement, options);
+          reqId = session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 0);
         }
       }
       if (!foundStats) {
@@ -209,7 +219,9 @@ private ByteString initTransaction(final Options options) {
                     .setExcludeTxnFromChangeStreams(
                         options.withExcludeTxnFromChangeStreams() == Boolean.TRUE))
             .build();
-    Transaction tx = rpc.beginTransaction(request, session.getOptions(), true);
+    XGoogSpannerRequestId reqId =
+        session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 1);
+    Transaction tx = rpc.beginTransaction(request, reqId.withOptions(session.getOptions()), true);
     if (tx.getId().isEmpty()) {
       throw SpannerExceptionFactory.newSpannerException(
           ErrorCode.INTERNAL,

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

@@ -71,6 +71,8 @@ abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet
   private CloseableIterator<PartialResultSet> stream;
   private ByteString resumeToken;
   private boolean finished;
+  public XGoogSpannerRequestId xGoogRequestId;
+  private XGoogSpannerRequestId.RequestIdCreator xGoogRequestIdCreator;
   /**
    * Indicates whether it is currently safe to retry RPCs. This will be {@code false} if we have
    * reached the maximum buffer size without seeing a restart token; in this case, we will drain the
@@ -85,7 +87,8 @@ protected ResumableStreamIterator(
       TraceWrapper tracer,
       ErrorHandler errorHandler,
       RetrySettings streamingRetrySettings,
-      Set<Code> retryableCodes) {
+      Set<Code> retryableCodes,
+      XGoogSpannerRequestId.RequestIdCreator xGoogRequestIdCreator) {
     this(
         maxBufferSize,
         streamName,
@@ -94,7 +97,8 @@ protected ResumableStreamIterator(
         Attributes.empty(),
         errorHandler,
         streamingRetrySettings,
-        retryableCodes);
+        retryableCodes,
+        xGoogRequestIdCreator);
   }
 
   protected ResumableStreamIterator(
@@ -105,14 +109,16 @@ protected ResumableStreamIterator(
       Attributes attributes,
       ErrorHandler errorHandler,
       RetrySettings streamingRetrySettings,
-      Set<Code> retryableCodes) {
+      Set<Code> retryableCodes,
+      XGoogSpannerRequestId.RequestIdCreator xGoogRequestIdCreator) {
     checkArgument(maxBufferSize >= 0);
     this.maxBufferSize = maxBufferSize;
     this.tracer = tracer;
     this.span = tracer.spanBuilderWithExplicitParent(streamName, parent, attributes);
     this.errorHandler = errorHandler;
     this.streamingRetrySettings = Preconditions.checkNotNull(streamingRetrySettings);
     this.retryableCodes = Preconditions.checkNotNull(retryableCodes);
+    this.xGoogRequestIdCreator = xGoogRequestIdCreator;
   }
 
   private ExponentialBackOff newBackOff() {
@@ -189,6 +195,14 @@ private void backoffSleep(Context context, long backoffMillis) throws SpannerExc
     }
   }
 
+  public void incrementXGoogRequestIdAttempt() {
+    if (this.xGoogRequestId == null) {
+      this.xGoogRequestId =
+          this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: infer channelId*/, 0 /*attempt*/);
+    }
+    this.xGoogRequestId.incrementAttempt();
+  }
+
   private enum DirectExecutor implements Executor {
     INSTANCE;
 
@@ -281,6 +295,7 @@ protected PartialResultSet computeNext() {
           assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
           stream = null;
           try (IScope s = tracer.withSpan(span)) {
+            incrementXGoogRequestIdAttempt();
             long delay = spannerException.getRetryDelayInMillis();
             if (delay != -1) {
               backoffSleep(context, delay);
@@ -301,6 +316,7 @@ protected PartialResultSet computeNext() {
             if (++numAttemptsOnOtherChannel < errorHandler.getMaxAttempts()
                 && prepareIteratorForRetryOnDifferentGrpcChannel()) {
               stream = null;
+              xGoogRequestId = null;
               continue;
             }
           }
@@ -326,6 +342,10 @@ private void startGrpcStreaming() {
         // When start a new stream set the Span as current to make the gRPC Span a child of
         // this Span.
         stream = checkNotNull(startStream(resumeToken, streamMessageListener));
+        if (xGoogRequestId == null) {
+          xGoogRequestId =
+              this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: retrieve channelId*/, 0);
+        }
         stream.requestPrefetchChunks();
       }
     }