Fix up withNthRequest for deterministic checks for BatchCreateSessions

Author: odeke-em

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

@@ -41,6 +41,7 @@
 import io.grpc.Status;
 import java.time.Duration;
 import java.time.temporal.ChronoUnit;
+import java.util.Arrays;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
@@ -79,11 +80,11 @@ 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(), 1);
-    }
+    XGoogSpannerRequestId reqId =
+        session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
+    UpdateOption[] allOptions = Arrays.copyOf(updateOptions, updateOptions.length + 1);
+    allOptions[updateOptions.length] = new Options.RequestIdOption(reqId);
+    Options options = Options.fromUpdateOptions(allOptions);
 
     try {
       ExecuteSqlRequest request = newTransactionRequestFrom(statement, options);
@@ -222,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/SessionClient.java

@@ -148,7 +148,6 @@ public void run() {
             .addAnnotation(String.format("Creating %d sessions", sessionCount));
         while (remainingSessionsToCreate > 0) {
           try {
-            System.out.println("\033[35mchannelHint: " + channelHint + "\033[00m");
             sessions = internalBatchCreateSessions(remainingSessionsToCreate, channelHint);
           } catch (Throwable t) {
             spanner.getTracer().getCurrentSpan().setStatus(t);
@@ -224,10 +223,8 @@ DatabaseId getDatabaseId() {
 
   @Override
   public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
-    long nthReq = this.nthRequest.incrementAndGet();
-    // System.out.println("\033[36mnthRequest.addr: " + System.identityHashCode(this.nthRequest) + "
-    // value: " + nthReq + "\033[00m");
-    return XGoogSpannerRequestId.of(this.nthId, nthReq, channelId, attempt);
+    return XGoogSpannerRequestId.of(
+        this.nthId, channelId, this.nthRequest.incrementAndGet(), attempt);
   }
 
   /** Create a single session. */

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

@@ -176,4 +176,10 @@ public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) {
       return XGoogSpannerRequestId.of(1, 1, 1, 0);
     }
   }
+
+  @VisibleForTesting
+  XGoogSpannerRequestId withNthRequest(long replacementNthRequest) {
+    return XGoogSpannerRequestId.of(
+        this.nthClientId, this.nthChannelId, replacementNthRequest, this.attempt);
+  }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -2913,28 +2913,29 @@ public void testPartitionedDmlDoesNotTimeout() {
       XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = {
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/ExecuteStreamingSql",
-            new XGoogSpannerRequestId(dbId, channelId, 1, 1)),
+            new XGoogSpannerRequestId(dbId, channelId, 5, 1)),
       };
       xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues);
 
       XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = {
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/BatchCreateSessions",
-            new XGoogSpannerRequestId(dbId, 0, 4, 1)),
+            new XGoogSpannerRequestId(dbId, 0, XGoogSpannerRequestIdTest.UNNECESSARY, 1)),
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/BatchCreateSessions",
-            new XGoogSpannerRequestId(dbId, 1, 2, 1)),
+            new XGoogSpannerRequestId(dbId, 1, XGoogSpannerRequestIdTest.UNNECESSARY, 1)),
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/BatchCreateSessions",
-            new XGoogSpannerRequestId(dbId, 2, 3, 1)),
+            new XGoogSpannerRequestId(dbId, 2, XGoogSpannerRequestIdTest.UNNECESSARY, 1)),
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/BatchCreateSessions",
-            new XGoogSpannerRequestId(dbId, 3, 1, 1)),
+            new XGoogSpannerRequestId(dbId, 3, XGoogSpannerRequestIdTest.UNNECESSARY, 1)),
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
-            "google.spanner.v1.Spanner/BeginTransaction", new XGoogSpannerRequestId(dbId, 3, 1, 1)),
+            "google.spanner.v1.Spanner/BeginTransaction",
+            new XGoogSpannerRequestId(dbId, channelId, 6, 1)),
         XGoogSpannerRequestIdTest.ofMethodAndRequestId(
             "google.spanner.v1.Spanner/ExecuteSql",
-            new XGoogSpannerRequestId(dbId, channelId, 5, 1)),
+            new XGoogSpannerRequestId(dbId, channelId, 7, 1)),
       };
       xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIds(wantUnaryValues);
     }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java

@@ -43,6 +43,7 @@
 
 @RunWith(JUnit4.class)
 public class XGoogSpannerRequestIdTest {
+  public static long UNNECESSARY = -1;
 
   @Test
   public void testEquals() {
@@ -184,20 +185,21 @@ public MethodAndRequestId[] accumulatedStreamingValues() {
     public void checkExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryValues) {
       MethodAndRequestId[] gotUnaryValues = this.accumulatedUnaryValues();
       sortValues(gotUnaryValues);
-      for (int i = 0; i < gotUnaryValues.length; i++) {
+      for (int i = 0; i < gotUnaryValues.length && false; i++) {
         System.out.println("\033[33misUnary: #" + i + ":: " + gotUnaryValues[i] + "\033[00m");
       }
       assertEquals(wantUnaryValues, gotUnaryValues);
     }
 
     private void sortValues(MethodAndRequestId[] values) {
+      massageValues(values);
       Arrays.sort(values, new MethodAndRequestIdComparator());
     }
 
     public void checkExpectedStreamingXGoogRequestIds(MethodAndRequestId... wantStreamingValues) {
       MethodAndRequestId[] gotStreamingValues = this.accumulatedStreamingValues();
       sortValues(gotStreamingValues);
-      for (int i = 0; i < gotStreamingValues.length; i++) {
+      for (int i = 0; i < gotStreamingValues.length && false; i++) {
         System.out.println(
             "\033[32misStreaming: #" + i + ":: " + gotStreamingValues[i] + "\033[00m");
       }
@@ -242,6 +244,7 @@ public int compare(MethodAndRequestId mr1, MethodAndRequestId mr2) {
       if (cmpMethod != 0) {
         return cmpMethod;
       }
+
       if (Objects.equals(mr1.requestId, mr2.requestId)) {
         return 0;
       }
@@ -252,6 +255,18 @@ public int compare(MethodAndRequestId mr1, MethodAndRequestId mr2) {
     }
   }
 
+  static void massageValues(MethodAndRequestId[] mreqs) {
+    for (int i = 0; i < mreqs.length; i++) {
+      MethodAndRequestId mreq = mreqs[i];
+      // BatchCreateSessions is so hard to control as the round-robin doling out
+      // hence we might need to be able to scrub the nth_request that won't match
+      // nth_req in consecutive order of nth_client.
+      if (mreq.method.compareTo("google.spanner.v1.Spanner/BatchCreateSessions") == 0) {
+        mreqs[i] = new MethodAndRequestId(mreq.method, mreq.requestId.withNthRequest(UNNECESSARY));
+      }
+    }
+  }
+
   public static MethodAndRequestId ofMethodAndRequestId(String method, String reqId) {
     return new MethodAndRequestId(method, XGoogSpannerRequestId.of(reqId));
   }