Wire up more x-goog-spanner-request-id propagations

Author: odeke-em

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -968,6 +968,7 @@ ResultSet readInternalWithOptions(
         builder.setLockHint(readOptions.lockHint());
       }
     }
+
     final int prefetchChunks =
         readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
     ResumableStreamIterator stream =

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java

@@ -169,7 +169,15 @@ public CommitResponse writeWithOptions(
       if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
         return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options);
       }
-      return runWithSessionRetry(session -> session.writeWithOptions(mutations, options));
+      int channelId = 1; /* TODO: infer the channelId from the gRPC channel of the session */
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, channelId, this.nextNthRequest(), 0);
+      return runWithSessionRetry(
+          session -> {
+            reqId.incrementAttempt();
+            // TODO: Update the channelId depending on the session that is inferred.
+            return session.writeWithOptions(mutations, appendReqIdToOptions(reqId, options));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -194,20 +202,15 @@ public CommitResponse writeAtLeastOnceWithOptions(
             .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);
-
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, channelId, this.nextNthRequest(), 0);
       return runWithSessionRetry(
           (session) -> {
             reqId.incrementAttempt();
             // TODO: Update the channelId depending on the session that is inferred.
-            ArrayList<TransactionOption> allOptions = new ArrayList<>();
-            allOptions.add(new Options.RequestIdOption(reqId));
-            allOptions.addAll(Arrays.asList(options));
-
             return session.writeAtLeastOnceWithOptions(
-                mutations, allOptions.toArray(new TransactionOption[0]));
+                mutations, appendReqIdToOptions(reqId, options));
           });
     } catch (RuntimeException e) {
       span.setStatus(e);
@@ -217,6 +220,24 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  private TransactionOption[] appendReqIdToOptions(
+      XGoogSpannerRequestId reqId, TransactionOption... options) {
+    ArrayList<TransactionOption> allOptions = new ArrayList<>();
+    allOptions.add(new Options.RequestIdOption(reqId));
+    allOptions.addAll(Arrays.asList(options));
+
+    return allOptions.toArray(new TransactionOption[0]);
+  }
+
+  private UpdateOption[] appendReqIdToOptions(
+      XGoogSpannerRequestId reqId, UpdateOption... options) {
+    ArrayList<UpdateOption> allOptions = new ArrayList<>();
+    allOptions.add(new Options.RequestIdOption(reqId));
+    allOptions.addAll(Arrays.asList(options));
+
+    return allOptions.toArray(new UpdateOption[0]);
+  }
+
   private int nextNthRequest() {
     return this.nthRequest.addAndGet(1);
   }
@@ -230,7 +251,17 @@ public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
       if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) {
         return getMultiplexedSessionDatabaseClient().batchWriteAtLeastOnce(mutationGroups, options);
       }
-      return runWithSessionRetry(session -> session.batchWriteAtLeastOnce(mutationGroups, options));
+
+      int channelId = 1; /* TODO: infer the channelId from the gRPC channel of the session */
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, channelId, this.nextNthRequest(), 0);
+
+      return runWithSessionRetry(
+          (session) -> {
+            reqId.incrementAttempt();
+            return session.batchWriteAtLeastOnce(
+                mutationGroups, appendReqIdToOptions(reqId, options));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -378,7 +409,14 @@ 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));
+      int channelId = 1; /* TODO: infer the channelId from the gRPC channel of the session */
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.dbId, channelId, this.nextNthRequest(), 0);
+      return runWithSessionRetry(
+          session -> {
+            reqId.incrementAttempt();
+            return session.executePartitionedUpdate(stmt, appendReqIdToOptions(reqId, options));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       span.end();

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java

@@ -1028,7 +1028,8 @@ public boolean equals(Object o) {
     }
   }
 
-  static final class RequestIdOption extends InternalOption implements TransactionOption {
+  static final class RequestIdOption extends InternalOption
+      implements UpdateOption, TransactionOption {
     private final XGoogSpannerRequestId reqId;
 
     RequestIdOption(XGoogSpannerRequestId reqId) {

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java

@@ -243,6 +243,7 @@ protected PartialResultSet computeNext() {
           && (finished || !safeToRetry || !buffer.getLast().getResumeToken().isEmpty())) {
         return buffer.pop();
       }
+
       try {
         if (stream.hasNext()) {
           PartialResultSet next = stream.next();

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java

@@ -31,6 +31,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.concurrent.GuardedBy;
 
 /** Client for creating single sessions and batches of sessions. */
@@ -175,6 +176,12 @@ interface SessionConsumer {
   private final DatabaseId db;
   private final Attributes commonAttributes;
 
+  // SessionClient is created long before a DatabaseClientImpl is created,
+  // as batch sessions are firstly created then later attached to each Client.
+  private static AtomicInteger NTH_ID = new AtomicInteger(0);
+  private final int nthId;
+  private final AtomicInteger nthRequest;
+
   @GuardedBy("this")
   private volatile long sessionChannelCounter;
 
@@ -187,6 +194,8 @@ interface SessionConsumer {
     this.executorFactory = executorFactory;
     this.executor = executorFactory.get();
     this.commonAttributes = spanner.getTracer().createCommonAttributes(db);
+    this.nthId = SessionClient.NTH_ID.incrementAndGet();
+    this.nthRequest = new AtomicInteger(0);
   }
 
   @Override
@@ -207,19 +216,23 @@ SessionImpl createSession() {
     // The sessionChannelCounter could overflow, but that will just flip it to Integer.MIN_VALUE,
     // which is also a valid channel hint.
     final Map<SpannerRpc.Option, ?> options;
+    final long channelId;
     synchronized (this) {
       options = optionMap(SessionOption.channelHint(sessionChannelCounter++));
+      channelId = sessionChannelCounter;
     }
     ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_SESSION, this.commonAttributes);
     try (IScope s = spanner.getTracer().withSpan(span)) {
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelId, 1);
       com.google.spanner.v1.Session session =
           spanner
               .getRpc()
               .createSession(
                   db.getName(),
                   spanner.getOptions().getDatabaseRole(),
                   spanner.getOptions().getSessionLabels(),
-                  options);
+                  reqId.withOptions(options));
       SessionReference sessionReference =
           new SessionReference(
               session.getName(), session.getCreateTime(), session.getMultiplexed(), options);
@@ -388,10 +401,8 @@ private List<SessionImpl> internalBatchCreateSessions(
             .spanBuilderWithExplicitParent(SpannerImpl.BATCH_CREATE_SESSIONS_REQUEST, parent);
     span.addAnnotation(String.format("Requesting %d sessions", sessionCount));
     try (IScope s = spanner.getTracer().withSpan(span)) {
-      // TODO: Infer the caller client if possible else create separate
-      // outside counter for such asynchronous operations and then also
-      // increment the operations for each asynchronous operation.
-      XGoogSpannerRequestId reqId = XGoogSpannerRequestId.of(1, 1, 1, 1);
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelHint, 1);
       List<com.google.spanner.v1.Session> sessions =
           spanner
               .getRpc()

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -423,7 +423,8 @@ public AsyncTransactionManagerImpl transactionManagerAsync(TransactionOption...
 
   @Override
   public ApiFuture<Empty> asyncClose() {
-    return spanner.getRpc().asyncDeleteSession(getName(), getOptions());
+    XGoogSpannerRequestId reqId = XGoogSpannerRequestId.of(1, 2, 1, 1);
+    return spanner.getRpc().asyncDeleteSession(getName(), reqId.withOptions(options));
   }
 
   @Override

Diff for: google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -2047,7 +2047,6 @@ <ReqT, RespT> GrpcCallContext newCallContext(
             .withStreamWaitTimeoutDuration(waitTimeout)
             .withStreamIdleTimeoutDuration(idleTimeout);
 
-    // TODO: Infer the x-goog-spanner-request-id header and inject it in accordingly.
     CallContextConfigurator configurator = SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY.get();
     ApiCallContext apiCallContextFromContext = null;
     if (configurator != null) {
@@ -2058,6 +2057,10 @@ <ReqT, RespT> GrpcCallContext newCallContext(
 
   GrpcCallContext withRequestId(GrpcCallContext context, Map options, String methodName) {
     String reqId = (String) options.get(Option.REQUEST_ID);
+    if (reqId == null) {
+      return context;
+    }
+
     System.out.println("\033[32moptions.reqId: " + reqId + "\033[00m " + methodName);
     Map<String, List<String>> withReqId =
         ImmutableMap.of(

Diff for: google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -120,6 +120,7 @@
 import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Random;
 import java.util.Set;
@@ -222,7 +223,9 @@ public static void startStaticServer() throws IOException {
         StatementResult.query(SELECT1_FROM_TABLE, MockSpannerTestUtil.SELECT1_RESULTSET));
     mockSpanner.setBatchWriteResult(BATCH_WRITE_RESPONSES);
 
-    xGoogReqIdInterceptor = new XGoogSpannerRequestIdTest.ServerHeaderEnforcer();
+    Set<String> checkMethods =
+        new HashSet(Arrays.asList("google.spanner.v1.Spanner/BatchCreateSessions"));
+    xGoogReqIdInterceptor = new XGoogSpannerRequestIdTest.ServerHeaderEnforcer(checkMethods);
     executor = Executors.newSingleThreadExecutor();
     String uniqueName = InProcessServerBuilder.generateName();
     server =

Diff for: google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java

@@ -27,6 +27,7 @@
 import io.grpc.ServerInterceptor;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import org.junit.Test;
@@ -61,27 +62,35 @@ public void testEnsureHexadecimalFormatForRandProcessID() {
 
   public static class ServerHeaderEnforcer implements ServerInterceptor {
     private List<String> gotValues;
+    private Set<String> checkMethods;
 
-    ServerHeaderEnforcer() {
+    ServerHeaderEnforcer(Set<String> checkMethods) {
       this.gotValues = new ArrayList<String>();
+      this.checkMethods = checkMethods;
     }
 
     @Override
     public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
         ServerCall<ReqT, RespT> call,
         final Metadata requestHeaders,
         ServerCallHandler<ReqT, RespT> next) {
+      String methodName = call.getMethodDescriptor().getFullMethodName();
+      if (!this.checkMethods.contains(methodName)) {
+        return next.startCall(call, requestHeaders);
+      }
+
       // Firstly assert and validate that at least we've got a requestId.
       String gotReqId = requestHeaders.get(XGoogSpannerRequestId.REQUEST_HEADER_KEY);
-      assertNotNull(gotReqId);
       Matcher m = XGoogSpannerRequestIdTest.REGEX_RAND_PROCESS_ID.matcher(gotReqId);
-      String message =
-          String.format(
-              "%s lacks %s in %s",
-              call.getMethodDescriptor().getFullMethodName(),
-              XGoogSpannerRequestId.REQUEST_HEADER_KEY.name(),
-              gotReqId);
-      System.out.println("\033[32mMessage: " + message + "\033[00m");
+      if (!m.matches()) {
+        String message =
+            String.format(
+                "%s lacks %s", methodName, XGoogSpannerRequestId.REQUEST_HEADER_KEY.name());
+        System.out.println("\033[31mMessage: " + message + "\033[00m");
+      } else {
+        System.out.println("\033[32mMessage: " + methodName + " has " + gotReqId + "\033[00m");
+      }
+      assertNotNull(gotReqId);
       assertTrue(m.matches());
 
       this.gotValues.add(gotReqId);