Use session.getChannel() and assert for results

Author: odeke-em

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

@@ -251,7 +251,7 @@ private List<Partition> partitionReadUsingIndex(
       builder.setPartitionOptions(pbuilder.build());
 
       XGoogSpannerRequestId reqId =
-          session.getRequestIdCreator().nextRequestId(1 /*TODO: retrieve channelId*/, 1);
+          session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
       final PartitionReadRequest request = builder.build();
       try {
         PartitionResponse response = rpc.partitionRead(request, reqId.withOptions(options));
@@ -317,7 +317,7 @@ private List<Partition> partitionQuery(
       builder.setPartitionOptions(pbuilder.build());
 
       XGoogSpannerRequestId reqId =
-          session.getRequestIdCreator().nextRequestId(1 /*TODO: retrieve channelId*/, 1);
+          session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
       final PartitionQueryRequest request = builder.build();
       try {
         PartitionResponse response = rpc.partitionQuery(request, reqId.withOptions(options));

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

@@ -82,7 +82,7 @@ long executeStreamingPartitionedUpdate(
     Options options = Options.fromUpdateOptions(updateOptions);
     XGoogSpannerRequestId reqId = options.reqId();
     if (reqId == null) {
-      reqId = session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 0);
+      reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 0);
     }
 
     try {
@@ -223,7 +223,7 @@ private ByteString initTransaction(final Options options) {
             .build();
     XGoogSpannerRequestId reqId = options.reqId();
     if (reqId == null) {
-      reqId = session.getRequestIdCreator().nextRequestId(1 /*TODO: infer channelId*/, 1);
+      reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 1);
     }
     Transaction tx = rpc.beginTransaction(request, reqId.withOptions(session.getOptions()), true);
     if (tx.getId().isEmpty()) {

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

@@ -320,7 +320,7 @@ public CommitResponse writeAtLeastOnceWithOptions(
   private XGoogSpannerRequestId reqIdOrFresh(Options options) {
     XGoogSpannerRequestId reqId = options.reqId();
     if (reqId == null) {
-      reqId = this.getRequestIdCreator().nextRequestId(1 /* TODO: channelId */, 1);
+      reqId = this.getRequestIdCreator().nextRequestId(this.getChannel(), 1);
     }
     return reqId;
   }
@@ -466,7 +466,7 @@ public AsyncTransactionManagerImpl transactionManagerAsync(TransactionOption...
   @Override
   public ApiFuture<Empty> asyncClose() {
     XGoogSpannerRequestId reqId =
-        this.getRequestIdCreator().nextRequestId(1 /* TODO: channelId */, 0);
+        this.getRequestIdCreator().nextRequestId(this.getChannel(), 1);
     return spanner.getRpc().asyncDeleteSession(getName(), reqId.withOptions(getOptions()));
   }
 
@@ -475,7 +475,7 @@ public void close() {
     ISpan span = tracer.spanBuilder(SpannerImpl.DELETE_SESSION);
     try (IScope s = tracer.withSpan(span)) {
       XGoogSpannerRequestId reqId =
-          this.getRequestIdCreator().nextRequestId(1 /* TODO: channelId */, 0);
+          this.getRequestIdCreator().nextRequestId(this.getChannel(), 1);
       spanner.getRpc().deleteSession(getName(), reqId.withOptions(getOptions()));
     } catch (RuntimeException e) {
       span.setStatus(e);
@@ -507,7 +507,7 @@ ApiFuture<Transaction> beginTransactionAsync(
     final BeginTransactionRequest request = requestBuilder.build();
     final ApiFuture<Transaction> requestFuture;
     XGoogSpannerRequestId reqId =
-        this.getRequestIdCreator().nextRequestId(1 /* TODO: channelId */, 1);
+        this.getRequestIdCreator().nextRequestId(this.getChannel(), 1);
     try (IScope ignore = tracer.withSpan(span)) {
       requestFuture =
           spanner
@@ -598,4 +598,9 @@ public void setRequestIdCreator(XGoogSpannerRequestId.RequestIdCreator creator)
   public XGoogSpannerRequestId.RequestIdCreator getRequestIdCreator() {
     return this.requestIdCreator;
   }
+
+  int getChannel() {
+    Long channelHint = (Long) this.getOptions().get(SpannerRpc.Option.CHANNEL_HINT);
+    return (int) (channelHint % this.spanner.getOptions().getNumChannels());
+  }
 }

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

@@ -2918,9 +2918,13 @@ public void testPartitionedDmlDoesNotTimeout() {
             "google.spanner.v1.Spanner/BatchCreateSessions",
             new XGoogSpannerRequestId(1, dbId, channelId, 1)),
       };
-      XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = {};
-      xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIds(wantUnaryValues);
+      XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = {
+        XGoogSpannerRequestIdTest.ofMethodAndRequestId(
+            "google.spanner.v1.Spanner/ExecuteStreamingSql",
+            new XGoogSpannerRequestId(1, channelId, dbId, 2)),
+      };
       xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues);
+      xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIds(wantUnaryValues);
     }
   }
 

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

@@ -28,6 +28,8 @@
 import io.grpc.ServerInterceptor;
 import io.grpc.Status;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -162,7 +164,7 @@ public MethodAndRequestId[] accumulatedUnaryValues() {
       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)));
             }
           });
       return accumulated.toArray(new MethodAndRequestId[0]);
@@ -173,25 +175,33 @@ public MethodAndRequestId[] accumulatedStreamingValues() {
       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)));
             }
           });
       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]);
+        System.out.println("\033[34misUnary: #" + i + ":: " + gotUnaryValues[i] + "\033[00m");
+      }
+      sortValues(gotUnaryValues);
+      for (int i = 0; i < gotUnaryValues.length; i++) {
+        System.out.println("\033[33misUnary: #" + i + ":: " + gotUnaryValues[i] + "\033[00m");
       }
       assertEquals(wantUnaryValues, gotUnaryValues);
     }
 
+    private void sortValues(MethodAndRequestId[] 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++) {
-        System.out.println("ith: " + i + ":: " + gotStreamingValues[i]);
+        System.out.println("\033[32misStreaming: #" + i + ":: " + gotStreamingValues[i] + "\033[00m");
       }
       assertEquals(wantStreamingValues, gotStreamingValues);
     }
@@ -205,24 +215,50 @@ public void reset() {
 
   public static class MethodAndRequestId {
     String method;
-    String requestId;
+    XGoogSpannerRequestId requestId;
 
-    public MethodAndRequestId(String method, String requestId) {
+    public MethodAndRequestId(String method, XGoogSpannerRequestId requestId) {
       this.method = method;
       this.requestId = requestId;
     }
 
     public String toString() {
-      return "{" + this.method + ":" + this.requestId + "}";
+      return "{" + this.method + ":" + this.requestId.toString() + "}";
     }
+
+    @Override
+    public boolean equals(Object o) {
+      if (!(o instanceof MethodAndRequestId)) {
+          return false;
+      }
+      MethodAndRequestId other = (MethodAndRequestId) o;
+      return Objects.equals(this.method, other.method) && Objects.equals(this.requestId, other.requestId);
+    }
+  }
+
+  static class MethodAndRequestIdComparator implements Comparator<MethodAndRequestId> {
+      @Override
+      public int compare(MethodAndRequestId mr1, MethodAndRequestId mr2) {
+          int cmpMethod = mr1.method.compareTo(mr2.method);
+          if (cmpMethod != 0) {
+              return cmpMethod;
+          }
+          if (Objects.equals(mr1.requestId, mr2.requestId)) {
+             return 0;
+          } 
+          if (mr1.requestId.isGreaterThan(mr2.requestId)) {
+              return +1;
+          }
+          return -1;
+      }
   }
 
   public static MethodAndRequestId ofMethodAndRequestId(String method, String reqId) {
-    return new MethodAndRequestId(method, reqId);
+    return new MethodAndRequestId(method, XGoogSpannerRequestId.of(reqId));
   }
 
   public static MethodAndRequestId ofMethodAndRequestId(
       String method, XGoogSpannerRequestId reqId) {
-    return new MethodAndRequestId(method, reqId.toString());
+    return new MethodAndRequestId(method, reqId);
   }
 }