Look for a multi-endpoint specified in gRPC context.

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

@@ -30,6 +30,7 @@
 import io.grpc.ClientCall;
 import io.grpc.ClientCall.Listener;
 import io.grpc.ConnectivityState;
+import io.grpc.Context;
 import io.grpc.Grpc;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
@@ -118,6 +119,7 @@
 public class GcpMultiEndpointChannel extends ManagedChannel {
 
   public static final CallOptions.Key<String> ME_KEY = CallOptions.Key.create("MultiEndpoint");
+  public static final Context.Key<String> ME_CONTEXT_KEY = Context.key("MultiEndpoint");
   private final LabelKey endpointKey =
       LabelKey.create("endpoint", "Endpoint address.");
   private final Map<String, MultiEndpoint> multiEndpoints = new ConcurrentHashMap<>();
@@ -428,7 +430,10 @@ public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedE
   @Override
   public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
-    final String multiEndpointKey = callOptions.getOption(ME_KEY);
+    String multiEndpointKey = callOptions.getOption(ME_KEY);
+    if (multiEndpointKey == null) {
+      multiEndpointKey = ME_CONTEXT_KEY.get(Context.current());
+    }
     MultiEndpoint me = defaultMultiEndpoint;
     if (multiEndpointKey != null) {
       me = multiEndpoints.getOrDefault(multiEndpointKey, defaultMultiEndpoint);

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

@@ -1234,7 +1234,7 @@ public void testStateNotifications() throws InterruptedException {
     gcpChannel.notifyWhenStateChanged(ConnectivityState.SHUTDOWN, () ->
         immediateCallbackCalled.set(true));
 
-    TimeUnit.MILLISECONDS.sleep(1);
+    TimeUnit.MILLISECONDS.sleep(2);
 
     assertThat(immediateCallbackCalled.get()).isTrue();
 

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.grpc;
 
+import static com.google.cloud.grpc.GcpMultiEndpointChannel.ME_CONTEXT_KEY;
 import static com.google.cloud.grpc.GcpMultiEndpointChannel.ME_KEY;
 import static com.google.cloud.spanner.SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY;
 import static com.google.common.base.Preconditions.checkState;
@@ -654,8 +655,8 @@ public void testSpannerMultiEndpointClient() throws IOException, InterruptedExce
             followerPoolIndex + ": Binding \\d+ key\\(s\\) to channel \\d:.*"
         )).count()).isEqualTo(0);
 
-    // Create context for using follower-first multi-endpoint.
-    Function<String, Context> contextFor = meName -> Context.current()
+    // Function for creating a context with a specific multi-endpoint set in call options.
+    Function<String, Context> callContextFor = meName -> Context.current()
         .withValue(CALL_CONTEXT_CONFIGURATOR_KEY,
             new CallContextConfigurator() {
               @Nullable
@@ -667,8 +668,12 @@ public <ReqT, RespT> ApiCallContext configure(ApiCallContext context, ReqT reque
               }
             });
 
+    // Function for creating a context with a specific multi-endpoint set in the context key.
+    Function<String, Context> contextFor = meName ->
+        Context.current().withValue(ME_CONTEXT_KEY, meName);
+
     assertThat(getOkCallsCount(fakeRegistry, followerEndpoint)).isEqualTo(0);
-    // Use follower, make sure it is used.
+    // Use follower, make sure it is used. (multi-endpoint is set in the context)
     contextFor.apply("follower").run(readQuery);
     assertThat(getOkCallsCount(fakeRegistry, followerEndpoint)).isEqualTo(1);
 
@@ -713,10 +718,17 @@ public <ReqT, RespT> ApiCallContext configure(ApiCallContext context, ReqT reque
 
     // Make sure that the follower endpoint still works if specified.
     assertThat(getOkCallsCount(fakeRegistry, followerEndpoint)).isEqualTo(2);
-    // Use follower, make sure it is used.
-    contextFor.apply("follower-2").run(readQuery);
+    // Use follower, make sure it is used. (multi-endpoint is set in the call options)
+    callContextFor.apply("follower-2").run(readQuery);
     assertThat(getOkCallsCount(fakeRegistry, followerEndpoint)).isEqualTo(3);
 
+    // Use leader, make sure it is used. (multi-endpoint from the call options overrides context-set
+    // multi-endpoint)
+    contextFor.apply("follower-2").run(() ->
+      callContextFor.apply("leader").run(readQuery)
+    );
+    assertThat(getOkCallsCount(fakeRegistry, newLeaderEndpoint)).isEqualTo(2);
+
     gcpMultiEndpointChannel.shutdown();
     spanner.close();
   }
@@ -1002,9 +1014,6 @@ public void testSessionsCreatedWithoutRoundRobin() throws Exception {
     List<ListenableFuture<Session>> futures = new ArrayList<>();
     assertEquals(ConnectivityState.IDLE, gcpChannel.getState(false));
 
-    // Initial log messages count.
-    int logCount = logRecords.size();
-
     // Should create one session per channel.
     CreateSessionRequest req = CreateSessionRequest.newBuilder().setDatabase(DATABASE_PATH).build();
     for (int i = 0; i < MAX_CHANNEL; i++) {
@@ -1016,8 +1025,6 @@ public void testSessionsCreatedWithoutRoundRobin() throws Exception {
           poolIndex + ": Channel " + i + " created.");
       assertThat(lastLogMessage()).isEqualTo(
           poolIndex + ": Channel " + i + " picked for bind operation.");
-      logCount += 3;
-      assertThat(logRecords.size()).isEqualTo(logCount);
       assertThat(lastLogLevel()).isEqualTo(Level.FINEST);
     }
     // Each channel should have 1 active stream with the CreateSession request because we create them concurrently.
@@ -1046,7 +1053,8 @@ public void testSessionsCreatedWithoutRoundRobin() throws Exception {
     // Verify the channel is in use.
     assertEquals(1, currentChannel.getActiveStreamsCount());
 
-    logCount = logRecords.size();
+    // Initial log messages count.
+    int logCount = logRecords.size();
 
     // Create another 1 session per channel sequentially.
     // Without the round-robin it won't use the currentChannel as it has more active streams (1) than other channels.