feat: Load balancing options for BigtableChannelPool (#2667)

neilw4 · web-flow · commit 5adaa84d80df · 2025-09-22T17:58:48.000Z
Experiments have shown that least-in-flight balancing is particularly effective when some channels or backends have problems. Depends on: https://togithub.com/googleapis/java-bigtable/pull/2651
diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml
@@ -456,4 +456,16 @@
         <method>*sendPrimeRequestsAsync*</method>
         <to>com.google.api.core.ApiFuture</to>
     </difference>
+    <difference>
+        <!-- InternalApi was updated -->
+        <differenceType>7013</differenceType>
+        <className>com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings</className>
+        <method>com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$LoadBalancingStrategy getLoadBalancingStrategy()</method>
+    </difference>
+    <difference>
+        <!-- InternalApi was updated -->
+        <differenceType>7013</differenceType>
+        <className>com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings$Builder</className>
+        <method>com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$Builder setLoadBalancingStrategy(com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$LoadBalancingStrategy)</method>
+    </difference>
 </differences>
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java b/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;
@@ -42,6 +43,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -71,6 +73,8 @@ 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();
+  private final Supplier<Integer> picker;
 
   public static BigtableChannelPool create(
       BigtableChannelPoolSettings settings,
@@ -113,6 +117,23 @@ public static BigtableChannelPool create(
 
     entries.set(initialListBuilder.build());
     authority = entries.get().get(0).channel.authority();
+
+    switch (settings.getLoadBalancingStrategy()) {
+      case ROUND_ROBIN:
+        picker = this::pickEntryIndexRoundRobin;
+        break;
+      case LEAST_IN_FLIGHT:
+        picker = this::pickEntryIndexLeastInFlight;
+        break;
+      case POWER_OF_TWO_LEAST_IN_FLIGHT:
+        picker = this::pickEntryIndexPowerOfTwoLeastInFlight;
+        break;
+      default:
+        throw new IllegalStateException(
+            String.format(
+                "Unknown load balancing strategy %s", settings.getLoadBalancingStrategy()));
+    }
+
     this.executor = executor;
 
     if (!settings.isStaticSize()) {
@@ -138,19 +159,74 @@ 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 the index of an entry to use for the next call. The returned value *should* be within
+   * range, but callers should not assume that this is always the case as race conditions are
+   * possible.
+   */
+  private int pickEntryIndex() {
+    return picker.get();
+  }
+
+  /** Pick an entry using the Round Robin algorithm. */
+  private int pickEntryIndexRoundRobin() {
+    return Math.abs(indexTicker.getAndIncrement() % entries.get().size());
+  }
+
+  /** Pick an entry at random. */
+  private int pickEntryIndexRandom() {
+    return rng.nextInt(entries.get().size());
   }
 
-  Channel getChannel(int affinity) {
-    return new AffinityChannel(affinity);
+  /** 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);
+      int rpcs = entry.outstandingRpcs.get();
+      if (rpcs < minRpcs) {
+        minRpcs = rpcs;
+        candidates.clear();
+        candidates.add(i);
+      } else if (rpcs == minRpcs) {
+        candidates.add(i);
+      }
+    }
+    // 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 index) {
+    return new AffinityChannel(index);
   }
 
   /** {@inheritDoc} */
@@ -395,7 +471,9 @@ void refresh() {
    * Get and retain a Channel Entry. The returned Entry will have its rpc count incremented,
    * preventing it from getting recycled.
    */
-  Entry getRetainedEntry(int affinity) {
+  private Entry getRetainedEntry(int affinity) {
+    // If an entry is not retainable, that usually means that it's about to be replaced and if we
+    // retry we should get a new useable entry.
     // 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
@@ -543,10 +621,10 @@ private void shutdown() {
 
   /** Thin wrapper to ensure that new calls are properly reference counted. */
   private class AffinityChannel extends Channel {
-    private final int affinity;
+    private final int index;
 
-    public AffinityChannel(int affinity) {
-      this.affinity = affinity;
+    public AffinityChannel(int index) {
+      this.index = index;
     }
 
     @Override
@@ -557,9 +635,7 @@ public String authority() {
     @Override
     public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
         MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
-
-      Entry entry = getRetainedEntry(affinity);
-
+      Entry entry = getRetainedEntry(index);
       return new ReleasingClientCall<>(entry.channel.newCall(methodDescriptor, callOptions), entry);
     }
   }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings.java
@@ -16,10 +16,14 @@
 package com.google.cloud.bigtable.gaxx.grpc;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 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.time.Duration;
+import java.util.logging.Logger;
 
 /**
  * Settings to control {@link BigtableChannelPool} behavior.
@@ -41,12 +45,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 +120,10 @@ public abstract class BigtableChannelPoolSettings {
    */
   public abstract boolean isPreemptiveRefreshEnabled();
 
+  /** The load balancing strategy to use for distributing RPCs across channels. */
+  @InternalApi("Use CBT_LOAD_BALANCING_STRATEGY environment variable")
+  public abstract LoadBalancingStrategy getLoadBalancingStrategy();
+
   /**
    * Helper to check if the {@link BigtableChannelPool} implementation can skip dynamic size logic
    */
@@ -111,6 +140,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) {
+      throw new IllegalStateException(
+          String.format("Invalid load-balancing strategy %s", strategyString));
+    }
+  }
+
   public abstract Builder toBuilder();
 
   public static BigtableChannelPoolSettings copyFrom(ChannelPoolSettings externalSettings) {
@@ -121,6 +168,7 @@ public static BigtableChannelPoolSettings copyFrom(ChannelPoolSettings externalS
         .setMaxChannelCount(externalSettings.getMaxChannelCount())
         .setInitialChannelCount(externalSettings.getInitialChannelCount())
         .setPreemptiveRefreshEnabled(externalSettings.isPreemptiveRefreshEnabled())
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv())
         .build();
   }
 
@@ -131,6 +179,7 @@ public static BigtableChannelPoolSettings staticallySized(int size) {
         .setMaxRpcsPerChannel(Integer.MAX_VALUE)
         .setMinChannelCount(size)
         .setMaxChannelCount(size)
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv())
         .build();
   }
 
@@ -141,7 +190,8 @@ public static Builder builder() {
         .setMaxChannelCount(200)
         .setMinRpcsPerChannel(0)
         .setMaxRpcsPerChannel(Integer.MAX_VALUE)
-        .setPreemptiveRefreshEnabled(false);
+        .setPreemptiveRefreshEnabled(false)
+        .setLoadBalancingStrategy(loadBalancingStrategyFromEnv());
   }
 
   @AutoValue.Builder
@@ -158,6 +208,9 @@ public abstract static class Builder {
 
     public abstract Builder setPreemptiveRefreshEnabled(boolean enabled);
 
+    @InternalApi("Use CBT_LOAD_BALANCING_STRATEGY environment variable")
+    public abstract Builder setLoadBalancingStrategy(LoadBalancingStrategy strategy);
+
     abstract BigtableChannelPoolSettings autoBuild();
 
     public BigtableChannelPoolSettings build() {
