chore: wire up x-goog-spanner-request-id for BatchCreateSessions

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 #3537

Author: odeke-em

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

@@ -27,6 +27,9 @@
 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.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 
 class DatabaseClientImpl implements DatabaseClient {
@@ -40,6 +43,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 +91,15 @@ 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);
+    return Integer.parseInt(strWithValue);
   }
 
   @VisibleForTesting
@@ -179,8 +193,22 @@ 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<>();
+            allOptions.add(new Options.RequestIdOption(reqId));
+            allOptions.addAll(Arrays.asList(options));
+
+            return session.writeAtLeastOnceWithOptions(
+                mutations, allOptions.toArray(new TransactionOption[0]));
+          });
     } catch (RuntimeException e) {
       span.setStatus(e);
       throw e;
@@ -189,6 +217,10 @@ public CommitResponse writeAtLeastOnceWithOptions(
     }
   }
 
+  private int nextNthRequest() {
+    return this.nthRequest.incrementAndGet();
+  }
+
   @Override
   public ServerStream<BatchWriteResponse> batchWriteAtLeastOnce(
       final Iterable<MutationGroup> mutationGroups, final TransactionOption... options)

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

@@ -512,6 +512,7 @@ void appendToOptions(Options options) {
   private RpcOrderBy orderBy;
   private RpcLockHint lockHint;
   private Boolean lastStatement;
+  private XGoogSpannerRequestId reqId;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -568,6 +569,14 @@ String pageToken() {
     return pageToken;
   }
 
+  boolean hasReqId() {
+    return reqId != null;
+  }
+
+  XGoogSpannerRequestId reqId() {
+    return reqId;
+  }
+
   boolean hasFilter() {
     return filter != null;
   }
@@ -1018,4 +1027,29 @@ public boolean equals(Object o) {
       return o instanceof LastStatementUpdateOption;
     }
   }
+
+  static final class RequestIdOption extends InternalOption implements TransactionOption {
+    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/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. */
@@ -174,6 +175,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;
 
@@ -186,6 +193,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
@@ -206,19 +215,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);
@@ -387,6 +400,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)) {
+      XGoogSpannerRequestId reqId =
+          XGoogSpannerRequestId.of(this.nthId, this.nthRequest.incrementAndGet(), channelHint, 1);
       List<com.google.spanner.v1.Session> sessions =
           spanner
               .getRpc()
@@ -395,7 +410,7 @@ private List<SessionImpl> internalBatchCreateSessions(
                   sessionCount,
                   spanner.getOptions().getDatabaseRole(),
                   spanner.getOptions().getSessionLabels(),
-                  options);
+                  reqId.withOptions(options));
       span.addAnnotation(
           String.format(
               "Request for %d sessions returned %d sessions", sessionCount, sessions.size()));

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

@@ -17,9 +17,13 @@
 package com.google.cloud.spanner;
 
 import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.common.annotations.VisibleForTesting;
+import io.grpc.Metadata;
 import java.math.BigInteger;
 import java.security.SecureRandom;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
 @InternalApi
@@ -28,6 +32,9 @@ public class XGoogSpannerRequestId {
   @VisibleForTesting
   static final String RAND_PROCESS_ID = XGoogSpannerRequestId.generateRandProcessId();
 
+  public static final Metadata.Key<String> REQUEST_HEADER_KEY =
+      Metadata.Key.of("x-goog-spanner-request-id", Metadata.ASCII_STRING_MARSHALLER);
+
   @VisibleForTesting
   static final long VERSION = 1; // The version of the specification being implemented.
 
@@ -81,6 +88,18 @@ public boolean equals(Object other) {
         && Objects.equals(this.attempt, otherReqId.attempt);
   }
 
+  public void incrementAttempt() {
+    this.attempt++;
+  }
+
+  @SuppressWarnings("unchecked")
+  public Map withOptions(Map options) {
+    Map copyOptions = new HashMap<>();
+    copyOptions.putAll(options);
+    copyOptions.put(SpannerRpc.Option.REQUEST_ID, this.toString());
+    return copyOptions;
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(this.nthClientId, this.nthChannelId, this.nthRequest, this.attempt);

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

@@ -71,6 +71,7 @@
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
 import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider;
+import com.google.cloud.spanner.XGoogSpannerRequestId;
 import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
 import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
 import com.google.cloud.spanner.admin.database.v1.stub.GrpcDatabaseAdminCallableFactory;
@@ -88,6 +89,7 @@
 import com.google.common.base.Supplier;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.Resources;
 import com.google.common.util.concurrent.RateLimiter;
@@ -2023,6 +2025,8 @@ <ReqT, RespT> GrpcCallContext newCallContext(
         // Set channel affinity in GAX.
         context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
       }
+      String methodName = method.getFullMethodName();
+      context = withRequestId(context, options, methodName);
     }
     if (compressorName != null) {
       // This sets the compressor for Client -> Server.
@@ -2046,6 +2050,7 @@ <ReqT, RespT> GrpcCallContext newCallContext(
         context
             .withStreamWaitTimeoutDuration(waitTimeout)
             .withStreamIdleTimeoutDuration(idleTimeout);
+
     CallContextConfigurator configurator = SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY.get();
     ApiCallContext apiCallContextFromContext = null;
     if (configurator != null) {
@@ -2054,6 +2059,19 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     return (GrpcCallContext) context.merge(apiCallContextFromContext);
   }
 
+  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(
+            XGoogSpannerRequestId.REQUEST_HEADER_KEY.name(), Collections.singletonList(reqId));
+    return context.withExtraHeaders(withReqId);
+  }
+
   void registerResponseObserver(SpannerResponseObserver responseObserver) {
     responseObservers.add(responseObserver);
   }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java

@@ -78,7 +78,8 @@
 public interface SpannerRpc extends ServiceRpc {
   /** Options passed in {@link SpannerRpc} methods to control how an RPC is issued. */
   enum Option {
-    CHANNEL_HINT("Channel Hint");
+    CHANNEL_HINT("Channel Hint"),
+    REQUEST_ID("Request Id");
 
     private final String value;