chore(x-goog-spanner-request-id): assert expectations in tests for retries + aborts

Updates #3537

chore(x-goog-spanner-request-id): add BeginTransaction+ResumableStreamIterator

Plumbs x-goog-spanner-request-id into BeginTransaction and
ResumableStreamIterator and for PartitionedDmlTransaction.

Updates #3537

Get more tests reveal needs for plumbing

Add requestId to PartitionQuery + PartitionRead

Propagate requestId into some more SpannerException values

Use session.getChannel() and assert for results

More debugging + fix up sorting comparator

Fix more channelId TODOs

Fix up withNthRequest for deterministic checks for BatchCreateSessions

More retrofits

More updates for tests

More validity checks

Add debugs to assert behavior of outgoing headers

Author: odeke-em

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

@@ -48,7 +48,7 @@ class DatabaseClientImpl implements DatabaseClient {
   @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient;
   @VisibleForTesting final boolean useMultiplexedSessionPartitionedOps;
   @VisibleForTesting final boolean useMultiplexedSessionForRW;
-  private final int dbId;
+  @VisibleForTesting final int dbId;
   private final AtomicInteger nthRequest;
   private final Map<String, Integer> clientIdToOrdinalMap;
 

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

@@ -93,6 +93,7 @@ long executeStreamingPartitionedUpdate(
         final Duration remainingTimeout = tryUpdateTimeout(timeout, stopwatch);
 
         try {
+          System.out.println("\033[31mreqIdPump: " + reqId + "\033[00m");
           ServerStream<PartialResultSet> stream =
               rpc.executeStreamingPartitionedDml(
                   request, reqId.withOptions(session.getOptions()), remainingTimeout);

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

@@ -198,6 +198,12 @@ private void backoffSleep(Context context, long backoffMillis) throws SpannerExc
 
   public void ensureNonNullXGoogRequestId() {
     if (this.xGoogRequestId == null) {
+      System.out.println(
+          "\033[34mXGoogRequestId.ensureNonNull: "
+              + this.xGoogRequestId
+              + " for:: "
+              + System.identityHashCode(this)
+              + "\033[00m");
       this.xGoogRequestId =
           this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: infer channelId*/, 1 /*attempt*/);
     }
@@ -206,6 +212,11 @@ public void ensureNonNullXGoogRequestId() {
   public void incrementXGoogRequestIdAttempt() {
     this.ensureNonNullXGoogRequestId();
     this.xGoogRequestId.incrementAttempt();
+    System.out.println(
+        "\033[35mincrementXGoogAttempt: "
+            + this.xGoogRequestId
+            + " :: "
+            + System.identityHashCode(this));
   }
 
   private enum DirectExecutor implements Executor {

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

@@ -190,7 +190,7 @@ interface SessionConsumer {
   // SessionClient is created long before a DatabaseClientImpl is created,
   // as batch sessions are firstly created then later attached to each Client.
   private static final AtomicInteger NTH_ID = new AtomicInteger(0);
-  private final int nthId = NTH_ID.incrementAndGet();
+  private final int nthId;
   private final AtomicInteger nthRequest = new AtomicInteger(0);
 
   @GuardedBy("this")
@@ -205,6 +205,7 @@ interface SessionConsumer {
     this.executorFactory = executorFactory;
     this.executor = executorFactory.get();
     this.commonAttributes = spanner.getTracer().createCommonAttributes(db);
+    this.nthId = NTH_ID.incrementAndGet();
   }
 
   @Override
@@ -223,7 +224,7 @@ DatabaseId getDatabaseId() {
   @Override
   public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
     return XGoogSpannerRequestId.of(
-        this.nthId, this.nthRequest.incrementAndGet(), channelId, attempt);
+        this.nthId, channelId, this.nthRequest.incrementAndGet(), attempt);
   }
 
   /** Create a single session. */
@@ -423,7 +424,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, this.nthRequest.incrementAndGet(), channelHint, 1);
+          XGoogSpannerRequestId.of(this.nthId, channelHint, this.nthRequest.incrementAndGet(), 1);
       List<com.google.spanner.v1.Session> sessions =
           spanner
               .getRpc()

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

@@ -2042,7 +2042,8 @@ <ReqT, RespT> GrpcCallContext newCallContext(
       }
     }
     if (options != null) {
-      context = withRequestId(context, options);
+      // TODO(@odeke-em): Infer the affinity if it doesn't match up with in the request-id.
+      context = withRequestId(context, options, method);
     }
     context = context.withExtraHeaders(metadataProvider.newExtraHeaders(resource, projectName));
     if (routeToLeader && leaderAwareRoutingEnabled) {
@@ -2060,15 +2061,40 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     if (configurator != null) {
       apiCallContextFromContext = configurator.configure(context, request, method);
     }
+
+    // Debug the call headers before this.
+    Map<String, List<String>> hdrs = context.getExtraHeaders();
+    if (method.getFullMethodName().compareTo("google.spanner.v1.Spanner/DeleteSession") != 0) {
+      System.out.println(
+          "\033[32mextraHeaders going out for " + method.getFullMethodName() + "\033[00m");
+      for (Map.Entry<String, List<String>> entry : hdrs.entrySet()) {
+        System.out.println(
+            "\t\033[36mcall.Key: " + entry.getKey() + ":: " + entry.getValue() + "\033[00m");
+      }
+    }
     return (GrpcCallContext) context.merge(apiCallContextFromContext);
   }
 
-  GrpcCallContext withRequestId(GrpcCallContext context, Map<SpannerRpc.Option, ?> options) {
+  <ReqT, RespT> GrpcCallContext withRequestId(
+      GrpcCallContext context,
+      Map<SpannerRpc.Option, ?> options,
+      MethodDescriptor<ReqT, RespT> method) {
     XGoogSpannerRequestId reqId = (XGoogSpannerRequestId) options.get(Option.REQUEST_ID);
     if (reqId == null) {
       return context;
     }
 
+    String methodName = method.getFullMethodName();
+    if (methodName.compareTo("google.spanner.v1.Spanner/ExecuteStreamingSql") == 0) {
+      System.out.println(
+          "\033[36mGapiSpannerRpc.withRequestId: "
+              + reqId
+              + " for: "
+              + method.getFullMethodName()
+              + " "
+              + System.identityHashCode(context)
+              + "\033[00m");
+    }
     Map<String, List<String>> withReqId =
         ImmutableMap.of(
             XGoogSpannerRequestId.REQUEST_HEADER_KEY.name(),