diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml index a5f9d8c3e604..0d8c8eb370a6 100644 --- a/google-cloud-bigtable/clirr-ignored-differences.xml +++ b/google-cloud-bigtable/clirr-ignored-differences.xml @@ -456,4 +456,16 @@ *sendPrimeRequestsAsync* com.google.api.core.ApiFuture + + + 7013 + com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings + com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$LoadBalancingStrategy getLoadBalancingStrategy() + + + + 7013 + com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettings$Builder + com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$Builder setLoadBalancingStrategy(com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPoolSettings$LoadBalancingStrategy) + 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 index c8ced11158c9..173722f2f4bf 100644 --- 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 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 ClientCall newCall( MethodDescriptor 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 localEntries = entries.get(); + int minRpcs = Integer.MAX_VALUE; + List 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 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 ClientCall newCall( MethodDescriptor 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 index e94a7665e81e..4ef21418ed89 100644 --- 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() {