googleapis · pongad · Mar 15, 2018 · Mar 12, 2018 · Mar 13, 2018 · Mar 13, 2018
diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java
@@ -45,7 +45,7 @@
  */
 class ChannelPool extends ManagedChannel {
   private final ImmutableList<ManagedChannel> channels;
-  private final AtomicInteger requestCount = new AtomicInteger();
+  private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
 
   /**
@@ -140,8 +140,24 @@ public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedE
    * @return A {@link ManagedChannel} that can be used for a single RPC call.
    */
   private ManagedChannel getNextChannel() {
-    int currentRequestNum = requestCount.getAndIncrement();
-    int index = Math.abs(currentRequestNum % channels.size());
+    return getChannel(indexTicker.getAndIncrement());
+  }
+
+  /**
+   * Returns one of the channels managed by this pool. The pool continues to "own" the channel, and
+   * the caller should not shut it down.
+   *
+   * @param affinity Two calls to this method with the same affinity returns the same channel. The
+   *     reverse is not true: Two calls with different affinities might return the same channel.
+   *     However, the implementation should attempt to spread load evenly.
+   */
+  ManagedChannel getChannel(int affinity) {
+    int index = affinity % channels.size();
+    index = Math.abs(index);
+    // If index is the most negative int, abs(index) is still negative.
+    if (index < 0) {
+      index = 0;
+    }
     return channels.get(index);
   }
 }
diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
@@ -63,26 +63,29 @@ public final class GrpcCallContext implements ApiCallContext {
   private final CallOptions callOptions;
   @Nullable private final Duration streamWaitTimeout;
   @Nullable private final Duration streamIdleTimeout;
+  @Nullable private final Integer channelAffinity;
 
   /** Returns an empty instance with a null channel and default {@link CallOptions}. */
   public static GrpcCallContext createDefault() {
-    return new GrpcCallContext(null, CallOptions.DEFAULT, null, null);
+    return new GrpcCallContext(null, CallOptions.DEFAULT, null, null, null);
   }
 
   /** Returns an instance with the given channel and {@link CallOptions}. */
   public static GrpcCallContext of(Channel channel, CallOptions callOptions) {
-    return new GrpcCallContext(channel, callOptions, null, null);
+    return new GrpcCallContext(channel, callOptions, null, null, null);
   }
 
   private GrpcCallContext(
       Channel channel,
       CallOptions callOptions,
       @Nullable Duration streamWaitTimeout,
-      @Nullable Duration streamIdleTimeout) {
+      @Nullable Duration streamIdleTimeout,
+      @Nullable Integer channelAffinity) {
     this.channel = channel;
     this.callOptions = Preconditions.checkNotNull(callOptions);
     this.streamWaitTimeout = streamWaitTimeout;
     this.streamIdleTimeout = streamIdleTimeout;
+    this.channelAffinity = channelAffinity;
   }
 
   /**
@@ -155,7 +158,8 @@ public GrpcCallContext withStreamWaitTimeout(@Nullable Duration streamWaitTimeou
           streamWaitTimeout.compareTo(Duration.ZERO) >= 0, "Invalid timeout: < 0 s");
     }
 
-    return new GrpcCallContext(channel, callOptions, streamWaitTimeout, streamIdleTimeout);
+    return new GrpcCallContext(
+        channel, callOptions, streamWaitTimeout, streamIdleTimeout, channelAffinity);
   }
 
   @Override
@@ -165,7 +169,13 @@ public GrpcCallContext withStreamIdleTimeout(@Nullable Duration streamIdleTimeou
           streamIdleTimeout.compareTo(Duration.ZERO) >= 0, "Invalid timeout: < 0 s");
     }
 
-    return new GrpcCallContext(channel, callOptions, streamWaitTimeout, streamIdleTimeout);
+    return new GrpcCallContext(
+        channel, callOptions, streamWaitTimeout, streamIdleTimeout, channelAffinity);
+  }
+
+  public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) {
+    return new GrpcCallContext(
+        channel, callOptions, streamWaitTimeout, streamIdleTimeout, affinity);
   }
 
   @Override
@@ -205,11 +215,16 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
       newStreamIdleTimeout = this.streamIdleTimeout;
     }
 
+    Integer newChannelAffinity = grpcCallContext.channelAffinity;
+    if (newChannelAffinity == null) {
+      newChannelAffinity = this.channelAffinity;
+    }
+
     CallOptions newCallOptions =
         this.callOptions.withCallCredentials(newCallCredentials).withDeadline(newDeadline);
 
     return new GrpcCallContext(
-        newChannel, newCallOptions, newStreamWaitTimeout, newStreamIdleTimeout);
+        newChannel, newCallOptions, newStreamWaitTimeout, newStreamIdleTimeout, newChannelAffinity);
   }
 
   /** The {@link Channel} set on this context. */
@@ -244,16 +259,35 @@ public Duration getStreamIdleTimeout() {
     return streamIdleTimeout;
   }
 
+  /**
+   * The channel affinity for this context.
+   *
+   * @see ApiCallContext#withStreamIdleTimeout(Duration)
+   */
+  @BetaApi("The surface for streaming is not stable yet and may change in the future.")
+  @Nullable
+  public Integer getChannelAffinity() {
+    return channelAffinity;
+  }
+
   /** Returns a new instance with the channel set to the given channel. */
   public GrpcCallContext withChannel(Channel newChannel) {
     return new GrpcCallContext(
-        newChannel, this.callOptions, this.streamWaitTimeout, this.streamIdleTimeout);
+        newChannel,
+        this.callOptions,
+        this.streamWaitTimeout,
+        this.streamIdleTimeout,
+        this.channelAffinity);
   }
 
   /** Returns a new instance with the call options set to the given call options. */
   public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
     return new GrpcCallContext(
-        this.channel, newCallOptions, this.streamWaitTimeout, this.streamIdleTimeout);
+        this.channel,
+        newCallOptions,
+        this.streamWaitTimeout,
+        this.streamIdleTimeout,
+        this.channelAffinity);
   }
 
   public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams) {
@@ -281,11 +315,13 @@ public boolean equals(Object o) {
     return Objects.equals(channel, that.channel)
         && Objects.equals(callOptions, that.callOptions)
         && Objects.equals(streamWaitTimeout, that.streamWaitTimeout)
-        && Objects.equals(streamIdleTimeout, that.streamIdleTimeout);
+        && Objects.equals(streamIdleTimeout, that.streamIdleTimeout)
+        && Objects.equals(channelAffinity, that.channelAffinity);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(channel, callOptions, streamWaitTimeout, streamIdleTimeout);
+    return Objects.hash(
+        channel, callOptions, streamWaitTimeout, streamIdleTimeout, channelAffinity);
   }
 }
diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcClientCalls.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcClientCalls.java
@@ -32,6 +32,7 @@
 import com.google.api.client.util.Preconditions;
 import com.google.api.gax.rpc.ApiCallContext;
 import io.grpc.CallOptions;
+import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
 
@@ -57,6 +58,11 @@ public static <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
     CallOptions callOptions = grpcContext.getCallOptions();
     Preconditions.checkNotNull(callOptions);
 
-    return grpcContext.getChannel().newCall(descriptor, callOptions);
+    Channel channel = grpcContext.getChannel();
+    if (grpcContext.getChannelAffinity() != null && channel instanceof ChannelPool) {
+      channel = ((ChannelPool) channel).getChannel(grpcContext.getChannelAffinity());
+    }
+
+    return channel.newCall(descriptor, callOptions);
   }
 }
diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectCallable.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectCallable.java
@@ -34,9 +34,6 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.common.base.Preconditions;
-import io.grpc.CallOptions;
-import io.grpc.Channel;
-import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
 import io.grpc.stub.ClientCalls;
 
@@ -56,17 +53,9 @@ class GrpcDirectCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, Re
   public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputContext) {
     Preconditions.checkNotNull(request);
     Preconditions.checkNotNull(inputContext);
-    GrpcCallContext context = GrpcCallContext.createDefault().nullToSelf(inputContext);
 
     return new ListenableFutureToApiFuture<>(
-        ClientCalls.futureUnaryCall(
-            newCall(context.getChannel(), context.getCallOptions()), request));
-  }
-
-  private ClientCall<RequestT, ResponseT> newCall(Channel channel, CallOptions callOptions) {
-    Preconditions.checkNotNull(channel);
-    Preconditions.checkNotNull(callOptions);
-    return channel.newCall(descriptor, callOptions);
+        ClientCalls.futureUnaryCall(GrpcClientCalls.newCall(descriptor, inputContext), request));
   }
 
   @Override

diff --git a/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java b/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.grpc;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.gax.grpc.testing.FakeServiceGrpc;
+import com.google.type.Color;
+import com.google.type.Money;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ManagedChannel;
+import io.grpc.MethodDescriptor;
+import java.util.Arrays;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class GrpcClientCallsTest {
+  @Test
+  public void testAffinity() {
+    MethodDescriptor<Color, Money> descriptor = FakeServiceGrpc.METHOD_RECOGNIZE;
+
+    @SuppressWarnings("unchecked")
+    ClientCall<Color, Money> clientCall0 = Mockito.mock(ClientCall.class);
+
+    @SuppressWarnings("unchecked")
+    ClientCall<Color, Money> clientCall1 = Mockito.mock(ClientCall.class);
+
+    ManagedChannel channel0 = Mockito.mock(ManagedChannel.class);
+    ManagedChannel channel1 = Mockito.mock(ManagedChannel.class);
+
+    Mockito.when(channel0.newCall(Mockito.eq(descriptor), Mockito.<CallOptions>any()))
+        .thenReturn(clientCall0);
+    Mockito.when(channel1.newCall(Mockito.eq(descriptor), Mockito.<CallOptions>any()))
+        .thenReturn(clientCall1);
+
+    Channel pool = new ChannelPool(Arrays.asList(channel0, channel1));
+    GrpcCallContext context = GrpcCallContext.createDefault().withChannel(pool);
+
+    ClientCall<Color, Money> gotCallA =
+        GrpcClientCalls.newCall(descriptor, context.withChannelAffinity(0));
+    ClientCall<Color, Money> gotCallB =
+        GrpcClientCalls.newCall(descriptor, context.withChannelAffinity(0));
+    ClientCall<Color, Money> gotCallC =
+        GrpcClientCalls.newCall(descriptor, context.withChannelAffinity(1));
+
+    assertThat(gotCallA).isSameAs(gotCallB);
+    assertThat(gotCallA).isNotSameAs(gotCallC);
+  }
+}