Merge from main

Author: odeke-em

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

@@ -465,17 +465,22 @@ private void initTransactionInternal(BeginTransactionRequest request) {
                 request, reqId.withOptions(getTransactionChannelHint()), isRouteToLeader());
         if (!transaction.hasReadTimestamp()) {
           throw SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field");
+              ErrorCode.INTERNAL,
+              "Missing expected transaction.read_timestamp metadata field",
+              reqId);
         }
         if (transaction.getId().isEmpty()) {
           throw SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INTERNAL, "Missing expected transaction.id metadata field");
+              ErrorCode.INTERNAL, "Missing expected transaction.id metadata field", reqId);
         }
         try {
           timestamp = Timestamp.fromProto(transaction.getReadTimestamp());
         } catch (IllegalArgumentException e) {
           throw SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e);
+              ErrorCode.INTERNAL,
+              "Bad value in transaction.read_timestamp metadata field",
+              e,
+              reqId);
         }
         transactionId = transaction.getId();
         span.addAnnotation(
@@ -830,7 +835,7 @@ CloseableIterator<PartialResultSet> startStream(
             if (selector != null) {
               request.setTransaction(selector);
             }
-            this.incrementXGoogRequestIdAttempt();
+            this.ensureNonNullXGoogRequestId();
             SpannerRpc.StreamingCall call =
                 rpc.executeQuery(
                     request.build(),
@@ -1036,6 +1041,7 @@ CloseableIterator<PartialResultSet> startStream(
             }
             builder.setRequestOptions(buildRequestOptions(readOptions));
             this.incrementXGoogRequestIdAttempt();
+            this.xGoogRequestId.setChannelId(session.getChannel());
             SpannerRpc.StreamingCall call =
                 rpc.read(
                     builder.build(),

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

@@ -79,17 +79,17 @@ long executeStreamingPartitionedUpdate(
     boolean foundStats = false;
     long updateCount = 0L;
     Stopwatch stopwatch = Stopwatch.createStarted(ticker);
-    Options options = Options.fromUpdateOptions(updateOptions);
-    XGoogSpannerRequestId reqId = options.reqId();
-    if (reqId == null) {
-      reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 0);
-    }
+    XGoogSpannerRequestId reqId =
+        session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
+    UpdateOption[] allOptions = new UpdateOption[updateOptions.length + 1];
+    System.arraycopy(updateOptions, 0, allOptions, 0, updateOptions.length);
+    allOptions[allOptions.length - 1] = new Options.RequestIdOption(reqId);
+    Options options = Options.fromUpdateOptions(allOptions);
 
     try {
       ExecuteSqlRequest request = newTransactionRequestFrom(statement, options);
 
       while (true) {
-        reqId.incrementAttempt();
         final Duration remainingTimeout = tryUpdateTimeout(timeout, stopwatch);
 
         try {
@@ -110,6 +110,7 @@ long executeStreamingPartitionedUpdate(
         } catch (UnavailableException e) {
           LOGGER.log(
               Level.FINER, "Retrying PartitionedDml transaction after UnavailableException", e);
+          reqId.incrementAttempt();
           request = resumeOrRestartRequest(resumeToken, statement, request, options);
         } catch (InternalException e) {
           if (!isRetryableInternalErrorPredicate.apply(e)) {
@@ -118,6 +119,7 @@ long executeStreamingPartitionedUpdate(
 
           LOGGER.log(
               Level.FINER, "Retrying PartitionedDml transaction after InternalException - EOS", e);
+          reqId.incrementAttempt();
           request = resumeOrRestartRequest(resumeToken, statement, request, options);
         } catch (AbortedException e) {
           LOGGER.log(Level.FINER, "Retrying PartitionedDml transaction after AbortedException", e);
@@ -126,7 +128,7 @@ long executeStreamingPartitionedUpdate(
           updateCount = 0L;
           request = newTransactionRequestFrom(statement, options);
           // Create a new xGoogSpannerRequestId.
-          reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 0);
+          reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
         } catch (SpannerException e) {
           e.setRequestId(reqId);
           throw e;
@@ -141,7 +143,7 @@ long executeStreamingPartitionedUpdate(
       LOGGER.log(Level.FINER, "Finished PartitionedUpdate statement");
       return updateCount;
     } catch (Exception e) {
-      throw SpannerExceptionFactory.newSpannerException(e);
+      throw SpannerExceptionFactory.newSpannerException(e, reqId);
     }
   }
 
@@ -221,10 +223,8 @@ private ByteString initTransaction(final Options options) {
                     .setExcludeTxnFromChangeStreams(
                         options.withExcludeTxnFromChangeStreams() == Boolean.TRUE))
             .build();
-    XGoogSpannerRequestId reqId = options.reqId();
-    if (reqId == null) {
-      reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
-    }
+    XGoogSpannerRequestId reqId =
+        session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
     Transaction tx = rpc.beginTransaction(request, reqId.withOptions(session.getOptions()), true);
     if (tx.getId().isEmpty()) {
       throw SpannerExceptionFactory.newSpannerException(

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

@@ -196,11 +196,15 @@ private void backoffSleep(Context context, long backoffMillis) throws SpannerExc
     }
   }
 
-  public void incrementXGoogRequestIdAttempt() {
+  public void ensureNonNullXGoogRequestId() {
     if (this.xGoogRequestId == null) {
       this.xGoogRequestId =
-          this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: infer channelId*/, 0 /*attempt*/);
+          this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: infer channelId*/, 1 /*attempt*/);
     }
+  }
+
+  public void incrementXGoogRequestIdAttempt() {
+    this.ensureNonNullXGoogRequestId();
     this.xGoogRequestId.incrementAttempt();
   }
 
@@ -347,7 +351,8 @@ private void startGrpcStreaming() {
         stream.requestPrefetchChunks();
         if (this.xGoogRequestId == null) {
           this.xGoogRequestId =
-              this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: retrieve channelId*/, 0);
+              this.xGoogRequestIdCreator.nextRequestId(
+                  1 /* channelId shall be replaced by the instantiated class. */, 0);
         }
       }
     }

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

@@ -223,7 +223,7 @@ DatabaseId getDatabaseId() {
   @Override
   public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
     return XGoogSpannerRequestId.of(
-        this.nthId, channelId, this.nthRequest.incrementAndGet(), attempt);
+        this.nthId, this.nthRequest.incrementAndGet(), channelId, attempt);
   }
 
   /** Create a single session. */
@@ -423,7 +423,7 @@ private List<SessionImpl> internalBatchCreateSessions(
     span.addAnnotation(String.format("Requesting %d sessions", sessionCount));
     try (IScope s = spanner.getTracer().withSpan(span)) {
       XGoogSpannerRequestId reqId =
-          XGoogSpannerRequestId.of(this.nthId, channelHint, this.nthRequest.incrementAndGet(), 1);
+          XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelHint, 1);
       List<com.google.spanner.v1.Session> sessions =
           spanner
               .getRpc()

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

@@ -600,7 +600,14 @@ int getChannel() {
     if (getIsMultiplexed()) {
       return 0;
     }
-    Long channelHint = (Long) this.getOptions().get(SpannerRpc.Option.CHANNEL_HINT);
+    Map<SpannerRpc.Option, ?> options = this.getOptions();
+    if (options == null) {
+      return 0;
+    }
+    Long channelHint = (Long) options.get(SpannerRpc.Option.CHANNEL_HINT);
+    if (channelHint == null) {
+      return 0;
+    }
     return (int) (channelHint % this.spanner.getOptions().getNumChannels());
   }
 }