Add disable affinity feature

grpc-gcp/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java

@@ -34,6 +34,7 @@
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
 import io.grpc.ConnectivityState;
+import io.grpc.Context;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
@@ -72,6 +73,10 @@ public class GcpManagedChannel extends ManagedChannel {
   static final AtomicInteger channelPoolIndex = new AtomicInteger();
   static final int DEFAULT_MAX_CHANNEL = 10;
   static final int DEFAULT_MAX_STREAM = 100;
+  public static final Context.Key<Boolean> DISABLE_AFFINITY_CTX_KEY =
+    Context.keyWithDefault("DisableAffinity", false);
+  public static final CallOptions.Key<Boolean> DISABLE_AFFINITY_KEY =
+    CallOptions.Key.createWithDefault("DisableAffinity", false);
 
   @GuardedBy("this")
   private Integer bindingIndex = -1;
@@ -1224,7 +1229,10 @@ public String authority() {
   public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
       MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions) {
     AffinityConfig affinity = methodToAffinity.get(methodDescriptor.getFullMethodName());
-    if (affinity == null) {
+    if (
+        affinity == null
+        || callOptions.getOption(DISABLE_AFFINITY_KEY)
+        || DISABLE_AFFINITY_CTX_KEY.get(Context.current())) {
       return new GcpClientCall.SimpleGcpClientCall<>(
           getChannelRef(null), methodDescriptor, callOptions);
     }

grpc-gcp/src/test/java/com/google/cloud/grpc/SpannerIntegrationTest.java

@@ -1258,6 +1258,99 @@ public void testExecuteStreamingSqlAsync() throws Exception {
     deleteAsyncSessions(stub, respNames);
   }
 
+  @Test
+  public void testExecuteStreamingSqlWithAffinityDisabledViaContext() throws Exception {
+    SpannerStub stub = getSpannerStub();
+    String sessionName = createAsyncSessions(stub).get(0);
+    AsyncResponseObserver<PartialResultSet> resp = new AsyncResponseObserver<>();
+    stub.executeStreamingSql(
+        ExecuteSqlRequest
+          .newBuilder()
+          .setSession(sessionName)
+          .setSql("select * FROM Users")
+          .build(),
+          resp);
+    // The ChannelRef which is bound with the current affinity key.
+    ChannelRef currentChannel =
+        gcpChannel.affinityKeyToChannelRef.get(sessionName);
+    // Verify the channel is in use.
+    assertEquals(1, currentChannel.getActiveStreamsCount());
+    PartialResultSet response = resp.get();
+    assertEquals(USERNAME, response.getValues(1).getStringValue());
+    assertEquals(0, currentChannel.getActiveStreamsCount());
+
+    Context ctx = Context.current().withValue(GcpManagedChannel.DISABLE_AFFINITY_CTX_KEY, true);
+    List<AsyncResponseObserver<PartialResultSet>> resps = new ArrayList<>();
+    for (int i = 0; i < MAX_CHANNEL; i++) {
+      // Execute the call in the context where affinity is disabled.
+      ctx.run(() -> {
+        AsyncResponseObserver<PartialResultSet> r = new AsyncResponseObserver<>();
+        resps.add(r);
+        stub.executeStreamingSql(
+          ExecuteSqlRequest.newBuilder().setSession(sessionName).setSql("select * FROM Users").build(),
+          r);
+      });
+    }
+    // Verify calls with disabled affinity are distributed accross all channels.
+    for (ChannelRef ch : gcpChannel.channelRefs) {
+      assertEquals(1, ch.getActiveStreamsCount());
+    }
+
+    for (AsyncResponseObserver<PartialResultSet> r : resps) {
+      response = r.get();
+      assertEquals(USERNAME, response.getValues(1).getStringValue());
+      assertEquals(0, currentChannel.getActiveStreamsCount());
+    }
+  }
+
+  @Test
+  public void testExecuteStreamingSqlWithAffinityDisabledViaCallOptions() throws Exception {
+    SpannerStub stub = getSpannerStub();
+    String sessionName = createAsyncSessions(stub).get(0);
+    AsyncResponseObserver<PartialResultSet> resp = new AsyncResponseObserver<>();
+    stub.executeStreamingSql(
+        ExecuteSqlRequest
+          .newBuilder()
+          .setSession(sessionName)
+          .setSql("select * FROM Users")
+          .build(),
+          resp);
+    // The ChannelRef which is bound with the current affinity key.
+    ChannelRef currentChannel =
+        gcpChannel.affinityKeyToChannelRef.get(sessionName);
+    // Verify the channel is in use.
+    assertEquals(1, currentChannel.getActiveStreamsCount());
+    PartialResultSet response = resp.get();
+    assertEquals(USERNAME, response.getValues(1).getStringValue());
+    assertEquals(0, currentChannel.getActiveStreamsCount());
+
+    List<AsyncResponseObserver<PartialResultSet>> resps = new ArrayList<>();
+    for (int i = 0; i < MAX_CHANNEL; i++) {
+      AsyncResponseObserver<PartialResultSet> r = new AsyncResponseObserver<>();
+      resps.add(r);
+      // Execute the call with affinity disabled via call option.
+      stub
+        .withOption(GcpManagedChannel.DISABLE_AFFINITY_KEY, true)
+        .executeStreamingSql(
+          ExecuteSqlRequest
+            .newBuilder()
+            .setSession(sessionName)
+            .setSql("select * FROM Users")
+            .build(),
+          r);
+    }
+    // Verify calls with disabled affinity are distributed accross all channels.
+    for (ChannelRef ch : gcpChannel.channelRefs) {
+      assertEquals(1, ch.getActiveStreamsCount());
+    }
+
+    for (AsyncResponseObserver<PartialResultSet> r : resps) {
+      response = r.get();
+      assertEquals(USERNAME, response.getValues(1).getStringValue());
+      assertEquals(0, currentChannel.getActiveStreamsCount());
+    }
+  }
+
   @Test
   public void testPartitionQueryAsync() throws Exception {
     SpannerStub stub = getSpannerStub();