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

Updates #3537

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/SessionClient.java

@@ -148,6 +148,7 @@ 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);

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

@@ -2904,6 +2904,21 @@ public void testPartitionedDmlDoesNotTimeout() {
                             return null;
                           }));
       assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode());
+
+      DatabaseClientImpl dbImpl = ((DatabaseClientImpl) client);
+      int channelId = dbImpl.getSession().getChannel();
+      int dbId = dbImpl.dbId;
+      XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = {
+        XGoogSpannerRequestIdTest.ofMethodAndRequestId(
+            "google.spanner.v1.Spanner/BatchCreateSessions",
+            new XGoogSpannerRequestId(1, dbId, channelId, 1)),
+        XGoogSpannerRequestIdTest.ofMethodAndRequestId(
+            "google.spanner.v1.Spanner/BatchCreateSessions",
+            new XGoogSpannerRequestId(1, dbId, channelId, 1)),
+      };
+      XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = {};
+      xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIds(wantUnaryValues);
+      xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues);
     }
   }
 
@@ -5138,6 +5153,7 @@ public <ReqT, RespT> ApiCallContext configure(
 
   @Test
   public void testRetryOnResourceExhausted() {
+    // MARK: Retries here.
     final RetrySettings retrySettings =
         RetrySettings.newBuilder()
             .setInitialRpcTimeoutDuration(Duration.ofSeconds(60L))
@@ -5172,9 +5188,9 @@ public void testRetryOnResourceExhausted() {
         .setRetryableCodes(StatusCode.Code.UNAVAILABLE, StatusCode.Code.RESOURCE_EXHAUSTED)
         .setRetrySettings(retrySettings);
 
+    DatabaseClient client;
     try (Spanner spanner = builder.build().getService()) {
-      DatabaseClient client =
-          spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+      client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
       final int expectedRowCount = 5;
       RandomResultSetGenerator generator = new RandomResultSetGenerator(expectedRowCount);
       Statement statement = Statement.of("select * from random_table");
@@ -5219,6 +5235,8 @@ public void testRetryOnResourceExhausted() {
           mockSpanner.clearRequests();
         }
       }
+
+      xGoogReqIdInterceptor.assertIntegrity();
     }
   }
 

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

@@ -157,40 +157,43 @@ private void assertMonotonicityOfIds(String prefix, List<XGoogSpannerRequestId>
               + String.join("\n\t", violations.toArray(new String[0])));
     }
 
-    public static class methodAndRequestId {
-      String method;
-      String requestId;
-
-      public methodAndRequestId(String method, String requestId) {
-        this.method = method;
-        this.requestId = requestId;
-      }
-
-      public String toString() {
-        return "{" + this.method + ":" + this.requestId + "}";
-      }
-    }
-
-    public methodAndRequestId[] accumulatedUnaryValues() {
-      List<methodAndRequestId> accumulated = new ArrayList();
+    public MethodAndRequestId[] accumulatedUnaryValues() {
+      List<MethodAndRequestId> accumulated = new ArrayList();
       this.unaryResults.forEach(
           (String method, CopyOnWriteArrayList<XGoogSpannerRequestId> values) -> {
             for (int i = 0; i < values.size(); i++) {
-              accumulated.add(new methodAndRequestId(method, values.get(i).toString()));
+              accumulated.add(new MethodAndRequestId(method, values.get(i).toString()));
             }
           });
-      return accumulated.toArray(new methodAndRequestId[0]);
+      return accumulated.toArray(new MethodAndRequestId[0]);
     }
 
-    public methodAndRequestId[] accumulatedStreamingValues() {
-      List<methodAndRequestId> accumulated = new ArrayList();
+    public MethodAndRequestId[] accumulatedStreamingValues() {
+      List<MethodAndRequestId> accumulated = new ArrayList();
       this.streamingResults.forEach(
           (String method, CopyOnWriteArrayList<XGoogSpannerRequestId> values) -> {
             for (int i = 0; i < values.size(); i++) {
-              accumulated.add(new methodAndRequestId(method, values.get(i).toString()));
+              accumulated.add(new MethodAndRequestId(method, values.get(i).toString()));
             }
           });
-      return accumulated.toArray(new methodAndRequestId[0]);
+      return accumulated.toArray(new MethodAndRequestId[0]);
+    }
+
+    public void checkExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryValues) {
+      MethodAndRequestId[] gotUnaryValues = this.accumulatedUnaryValues();
+      System.out.println("\033[34mUnary: " + gotUnaryValues + "\033[00m");
+      for (int i = 0; i < gotUnaryValues.length; i++) {
+        System.out.println("ith: " + i + ":: " + gotUnaryValues[i]);
+      }
+      assertEquals(wantUnaryValues, gotUnaryValues);
+    }
+
+    public void checkExpectedStreamingXGoogRequestIds(MethodAndRequestId... wantStreamingValues) {
+      MethodAndRequestId[] gotStreamingValues = this.accumulatedStreamingValues();
+      for (int i = 0; i < gotStreamingValues.length; i++) {
+        System.out.println("ith: " + i + ":: " + gotStreamingValues[i]);
+      }
+      assertEquals(wantStreamingValues, gotStreamingValues);
     }
 
     public void reset() {
@@ -199,4 +202,27 @@ public void reset() {
       this.streamingResults.clear();
     }
   }
+
+  public static class MethodAndRequestId {
+    String method;
+    String requestId;
+
+    public MethodAndRequestId(String method, String requestId) {
+      this.method = method;
+      this.requestId = requestId;
+    }
+
+    public String toString() {
+      return "{" + this.method + ":" + this.requestId + "}";
+    }
+  }
+
+  public static MethodAndRequestId ofMethodAndRequestId(String method, String reqId) {
+    return new MethodAndRequestId(method, reqId);
+  }
+
+  public static MethodAndRequestId ofMethodAndRequestId(
+      String method, XGoogSpannerRequestId reqId) {
+    return new MethodAndRequestId(method, reqId.toString());
+  }
 }