Implement load-balancing options for BigtableChannelPool. Experiments have shown that least-in-flight balancing is particularly effective when some channels or backends have problems.

Author: neilw4

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java

@@ -34,6 +34,7 @@
 import java.time.Clock;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.Executors;
@@ -71,6 +72,7 @@ public class BigtableChannelPool extends ManagedChannel {
   private final ChannelPoolHealthChecker channelPoolHealthChecker;
   private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
+  private final Random rng = new Random();
 
   public static BigtableChannelPool create(
       BigtableChannelPoolSettings settings,
@@ -138,15 +140,94 @@ public String authority() {
   }
 
   /**
-   * Create a {@link ClientCall} on a Channel from the pool chosen in a round-robin fashion to the
-   * remote operation specified by the given {@link MethodDescriptor}. The returned {@link
-   * ClientCall} does not trigger any remote behavior until {@link
-   * ClientCall#start(ClientCall.Listener, io.grpc.Metadata)} is invoked.
+   * Create a {@link ClientCall} on a Channel from the pool to the remote operation specified by the
+   * given {@link MethodDescriptor}. The returned {@link ClientCall} does not trigger any remote
+   * behavior until {@link ClientCall#start(ClientCall.Listener, io.grpc.Metadata)} is invoked.
    */
   @Override
   public <ReqT, RespT> ClientCall<ReqT, RespT> newCall(
       MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions) {
-    return getChannel(indexTicker.getAndIncrement()).newCall(methodDescriptor, callOptions);
+    return new AffinityChannel(pickEntryIndex()).newCall(methodDescriptor, callOptions);
+  }
+
+  /** Pick an entry to use for the next call. */
+  private int pickEntryIndex() {
+    switch (settings.getLoadBalancingStrategy()) {
+      case ROUND_ROBIN:
+        return pickEntryIndexRoundRobin();
+      case LEAST_IN_FLIGHT:
+        return pickEntryIndexLeastInFlight();
+      case POWER_OF_TWO_LEAST_IN_FLIGHT:
+        return pickEntryIndexPowerOfTwoLeastInFlight();
+      default:
+          LOG.warning(String.format(
+            "Unknown load balancing strategy %s, falling back to ROUND_ROBIN.",
+            settings.getLoadBalancingStrategy()));
+          return pickEntryIndexRoundRobin();
+
+    }
+  }
+
+  /** Pick an entry using the Round Robin algorithm. */
+  private int pickEntryIndexRoundRobin() {
+    return indexTicker.getAndIncrement();
+  }
+
+  /** Pick an entry at random. */
+  private int pickEntryIndexRandom() {
+    List<Entry> localEntries = entries.get();
+    for (int attempt = 0; attempt < 5; attempt++) {
+      int choice = rng.nextInt(localEntries.size());
+      if (localEntries.get(choice).retainable()) {
+        return choice;
+      }
+    }
+    LOG.warning("couldn't find retainable channel, picking at random");
+    return rng.nextInt(localEntries.size());
+  }
+  
+  /** Pick an entry using the least-in-flight algorithm. */
+  private int pickEntryIndexLeastInFlight() {
+    List<Entry> localEntries = entries.get();
+    int minRpcs = Integer.MAX_VALUE;
+    List<Integer> candidates = new ArrayList<>();
+
+    for (int i = 0; i < localEntries.size(); i++) {
+      Entry entry = localEntries.get(i);
+      if (entry.retainable()) {
+        int rpcs = entry.outstandingRpcs.get();
+        if (rpcs < minRpcs) {
+          minRpcs = rpcs;
+          candidates.clear();
+          candidates.add(i);
+        } else if (rpcs == minRpcs) {
+          candidates.add(i);
+        }
+      }
+    }
+    if (candidates.isEmpty()) {
+      LOG.warning(
+        "Least-in-flight picker couldn't find available channel. Picking at random.");
+      return pickEntryIndexRandom();
+    }
+    // If there are multiple matching entries, pick one at random.
+    return candidates.get(rng.nextInt(candidates.size()));
+  }
+
+  /** Pick an entry using the power-of-two algorithm. */
+  private int pickEntryIndexPowerOfTwoLeastInFlight() {
+    List<Entry> localEntries = entries.get();
+    int choice1 = pickEntryIndexRandom();
+    int choice2 = pickEntryIndexRandom();
+    if (choice1 == choice2) {
+      // Try to pick two different entries. If this picks the same entry again, it's likely that
+      // there's only one healthy channel in the pool and we should proceed anyway.
+      choice2 = pickEntryIndexRandom();
+    }
+
+    Entry entry1 = localEntries.get(choice1);
+    Entry entry2 = localEntries.get(choice2);
+    return entry1.outstandingRpcs.get() < entry2.outstandingRpcs.get() ? choice1 : choice2;
   }
 
   Channel getChannel(int affinity) {
@@ -392,16 +473,23 @@ void refresh() {
   }
 
   /**
-   * Get and retain a Channel Entry. The returned Entry will have its rpc count incremented,
-   * preventing it from getting recycled.
+   * 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 most
+   *     of the time, if the channel pool was refreshed since the last call, a new channel will be
+   *     returned. The reverse is not true: Two calls with different affinities might return the
+   *     same channel. However, the implementation should attempt to spread load evenly.
    */
-  Entry getRetainedEntry(int affinity) {
+  private Entry getEntry(int affinity) {
     // The maximum number of concurrent calls to this method for any given time span is at most 2,
     // so the loop can actually be 2 times. But going for 5 times for a safety margin for potential
     // code evolving
     for (int i = 0; i < 5; i++) {
-      Entry entry = getEntry(affinity);
-      if (entry.retain()) {
+      List<Entry> localEntries = entries.get();
+      int index = Math.abs(affinity % localEntries.size());
+      Entry entry = localEntries.get(index);
+      if (entry.retainable()) {
         return entry;
       }
     }
@@ -412,23 +500,6 @@ Entry getRetainedEntry(int affinity) {
     throw new IllegalStateException("Bug: failed to retain a channel");
   }
 
-  /**
-   * 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 most
-   *     of the time, if the channel pool was refreshed since the last call, a new channel will be
-   *     returned. The reverse is not true: Two calls with different affinities might return the
-   *     same channel. However, the implementation should attempt to spread load evenly.
-   */
-  private Entry getEntry(int affinity) {
-    List<Entry> localEntries = entries.get();
-
-    int index = Math.abs(affinity % localEntries.size());
-
-    return localEntries.get(index);
-  }
-
   /** Bundles a gRPC {@link ManagedChannel} with some usage accounting. */
   static class Entry {
     private final ManagedChannel channel;
@@ -480,14 +551,23 @@ ManagedChannel getManagedChannel() {
     int getAndResetMaxOutstanding() {
       return maxOutstanding.getAndSet(outstandingRpcs.get());
     }
+   /**
+     * The method will return false if the channel is closing and the caller should pick a different
+     * channel. If the method returned true, the channel is useable and can be retained.
+     */
+    private boolean retainable() {
+      if (shutdownRequested.get()) {
+        requestShutdown();
+        return false;
+      }
+      return true;
+    }
 
     /**
-     * Try to increment the outstanding RPC count. The method will return false if the channel is
-     * closing and the caller should pick a different channel. If the method returned true, the
-     * channel has been successfully retained and it is the responsibility of the caller to release
-     * it.
+     * Try to increment the outstanding RPC count. It is the responsibility of the caller to check
+     * that the channel is retainable, and to release it after RPC completion.
      */
-    private boolean retain() {
+    private void retain() {
       // register desire to start RPC
       int currentOutstanding = outstandingRpcs.incrementAndGet();
 
@@ -496,13 +576,6 @@ private boolean retain() {
       if (currentOutstanding > prevMax) {
         maxOutstanding.incrementAndGet();
       }
-
-      // abort if the channel is closing
-      if (shutdownRequested.get()) {
-        release();
-        return false;
-      }
-      return true;
     }
 
     /**
@@ -517,8 +590,8 @@ private void release() {
 
       // Must check outstandingRpcs after shutdownRequested (in reverse order of retain()) to ensure
       // mutual exclusion.
-      if (shutdownRequested.get() && outstandingRpcs.get() == 0) {
-        shutdown();
+      if (shutdownRequested.get()) {
+        requestShutdown();
       }
     }
 
@@ -557,9 +630,13 @@ public String authority() {
     @Override
     public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
         MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
-
-      Entry entry = getRetainedEntry(affinity);
-
+      Entry entry = getEntry(affinity);
+      if (entry.retainable()) {
+        entry.retain();
+      } else {
+        LOG.warning("newCall() called on channel entry that is shutting down. The call may not "
+        + "succeed.");
+      }
       return new ReleasingClientCall<>(entry.channel.newCall(methodDescriptor, callOptions), entry);
     }
   }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings.java

@@ -18,7 +18,10 @@
 import com.google.api.core.BetaApi;
 import com.google.api.gax.grpc.ChannelPoolSettings;
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Strings;
+import java.util.logging.Logger;
 import java.time.Duration;
 
 /**
@@ -41,12 +44,33 @@
 @BetaApi("surface for channel pool sizing is not yet stable")
 @AutoValue
 public abstract class BigtableChannelPoolSettings {
+  @VisibleForTesting
+  static final Logger LOG = Logger.getLogger(BigtableChannelPoolSettings.class.getName());
+
   /** How often to check and possibly resize the {@link BigtableChannelPool}. */
   static final Duration RESIZE_INTERVAL = Duration.ofMinutes(1);
 
   /** The maximum number of channels that can be added or removed at a time. */
   static final int MAX_RESIZE_DELTA = 2;
 
+  /** Environment variable used to set load balancing strategy. */
+  private static final String CBT_LOAD_BALANCING_STRATEGY_ENV_VAR = "CBT_LOAD_BALANCING_STRATEGY";
+
+  /** Load balancing strategy to use if environment variable is unset or invalid. */
+  private static final LoadBalancingStrategy DEFAULT_LOAD_BALANCING_STRATEGY =
+      LoadBalancingStrategy.ROUND_ROBIN;
+
+  /** Supported load-balancing strategies. */
+  public enum LoadBalancingStrategy {
+    // Sequentially iterate across all channels.
+    ROUND_ROBIN,
+    // Pick the channel with the fewest in-flight requests. If multiple channels match, pick at
+    // random.
+    LEAST_IN_FLIGHT,
+    // Out of two random channels, pick the channel with the fewest in-flight requests.
+    POWER_OF_TWO_LEAST_IN_FLIGHT,
+  }
+
   /**
    * Threshold to start scaling down the channel pool.
    *
@@ -95,6 +119,9 @@ public abstract class BigtableChannelPoolSettings {
    */
   public abstract boolean isPreemptiveRefreshEnabled();
 
+  /** The load balancing strategy to use for distributing RPCs across channels. */
+  public abstract LoadBalancingStrategy getLoadBalancingStrategy();
+
   /**
    * Helper to check if the {@link BigtableChannelPool} implementation can skip dynamic size logic
    */
@@ -111,6 +138,24 @@ boolean isStaticSize() {
     return false;
   }
 
+  /**
+   * Use environment variable CBT_LOAD_BALANCING_STRATEGY to pick a load-balancing strategy.
+   * @return load-balancing strategy to use.
+   */
+  private static LoadBalancingStrategy loadBalancingStrategyFromEnv() {
+    String strategyString = System.getenv(CBT_LOAD_BALANCING_STRATEGY_ENV_VAR);
+    if (Strings.isNullOrEmpty(strategyString)) {
+      return DEFAULT_LOAD_BALANCING_STRATEGY;
+    }
+    try {
+      return LoadBalancingStrategy.valueOf(strategyString.trim().toUpperCase());
+    } catch (IllegalArgumentException e) {
+      LOG.warning(String.format("Invalid load-balancing strategy %s, falling back to default %s.",
+          strategyString, DEFAULT_LOAD_BALANCING_STRATEGY));
+      return DEFAULT_LOAD_BALANCING_STRATEGY;
+    }
+  }
+
   public abstract Builder toBuilder();
 
   public static BigtableChannelPoolSettings copyFrom(ChannelPoolSettings externalSettings) {
@@ -121,6 +166,7 @@ public static BigtableChannelPoolSettings copyFrom(ChannelPoolSettings externalS
         .setMaxChannelCount(externalSettings.getMaxChannelCount())
         .setInitialChannelCount(externalSettings.getInitialChannelCount())
         .setPreemptiveRefreshEnabled(externalSettings.isPreemptiveRefreshEnabled())
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv())
         .build();
   }
 
@@ -131,6 +177,7 @@ public static BigtableChannelPoolSettings staticallySized(int size) {
         .setMaxRpcsPerChannel(Integer.MAX_VALUE)
         .setMinChannelCount(size)
         .setMaxChannelCount(size)
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv())
         .build();
   }
 
@@ -141,7 +188,8 @@ public static Builder builder() {
         .setMaxChannelCount(200)
         .setMinRpcsPerChannel(0)
         .setMaxRpcsPerChannel(Integer.MAX_VALUE)
-        .setPreemptiveRefreshEnabled(false);
+        .setPreemptiveRefreshEnabled(false)
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv());
   }
 
   @AutoValue.Builder
@@ -158,6 +206,8 @@ public abstract static class Builder {
 
     public abstract Builder setPreemptiveRefreshEnabled(boolean enabled);
 
+    public abstract Builder setLoadBalancingStrategy(LoadBalancingStrategy strategy);
+
     abstract BigtableChannelPoolSettings autoBuild();
 
     public BigtableChannelPoolSettings build() {