diff --git a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java index 742e5d2f24..c253adb98b 100644 --- a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java +++ b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java @@ -28,8 +28,12 @@ import io.grpc.ConnectivityStateInfo; import io.grpc.EquivalentAddressGroup; import io.grpc.LoadBalancer; import io.grpc.Status; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.Random; import java.util.concurrent.atomic.AtomicBoolean; +import javax.annotation.Nullable; /** * A {@link LoadBalancer} that provides no load-balancing over the addresses from the {@link @@ -55,6 +59,18 @@ final class PickFirstLoadBalancer extends LoadBalancer { return false; } + // We can optionally be configured to shuffle the address list. This can help better distribute + // the load. + if (resolvedAddresses.getLoadBalancingPolicyConfig() instanceof PickFirstLoadBalancerConfig) { + PickFirstLoadBalancerConfig config + = (PickFirstLoadBalancerConfig) resolvedAddresses.getLoadBalancingPolicyConfig(); + if (config.shuffleAddressList != null && config.shuffleAddressList) { + servers = new ArrayList(servers); + Collections.shuffle(servers, + config.randomSeed != null ? new Random(config.randomSeed) : new Random()); + } + } + if (subchannel == null) { final Subchannel subchannel = helper.createSubchannel( CreateSubchannelArgs.newBuilder() @@ -199,4 +215,22 @@ final class PickFirstLoadBalancer extends LoadBalancer { return PickResult.withNoResult(); } } + + public static final class PickFirstLoadBalancerConfig { + + @Nullable + public final Boolean shuffleAddressList; + + // For testing purposes only, not meant to be parsed from a real config. + @Nullable final Long randomSeed; + + public PickFirstLoadBalancerConfig(@Nullable Boolean shuffleAddressList) { + this(shuffleAddressList, null); + } + + PickFirstLoadBalancerConfig(@Nullable Boolean shuffleAddressList, @Nullable Long randomSeed) { + this.shuffleAddressList = shuffleAddressList; + this.randomSeed = randomSeed; + } + } } diff --git a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java index 7f7b366564..35d96cfdcc 100644 --- a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java +++ b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java @@ -16,10 +16,13 @@ package io.grpc.internal; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Strings; import io.grpc.LoadBalancer; import io.grpc.LoadBalancerProvider; import io.grpc.NameResolver; import io.grpc.NameResolver.ConfigOrError; +import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig; import java.util.Map; /** @@ -30,6 +33,10 @@ import java.util.Map; */ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider { private static final String NO_CONFIG = "no service config"; + private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList"; + private static final String CONFIG_FLAG_NAME = "GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG"; + @VisibleForTesting + static boolean enablePickFirstConfig = !Strings.isNullOrEmpty(System.getenv(CONFIG_FLAG_NAME)); @Override public boolean isAvailable() { @@ -54,6 +61,12 @@ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider { @Override public ConfigOrError parseLoadBalancingPolicyConfig( Map rawLoadBalancingPolicyConfig) { - return ConfigOrError.fromConfig(NO_CONFIG); + if (enablePickFirstConfig) { + return ConfigOrError.fromConfig( + new PickFirstLoadBalancerConfig(JsonUtil.getBoolean(rawLoadBalancingPolicyConfig, + SHUFFLE_ADDRESS_LIST_KEY))); + } else { + return ConfigOrError.fromConfig(NO_CONFIG); + } } } diff --git a/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java b/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java new file mode 100644 index 0000000000..1c224eb1c3 --- /dev/null +++ b/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java @@ -0,0 +1,61 @@ +/* + * Copyright 2023 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.internal; + +import static com.google.common.truth.Truth.assertThat; + +import io.grpc.NameResolver.ConfigOrError; +import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig; +import java.util.HashMap; +import java.util.Map; +import org.junit.After; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class PickFirstLoadBalancerProviderTest { + + @After + public void resetConfigFlag() { + PickFirstLoadBalancerProvider.enablePickFirstConfig = false; + } + + @Test + public void parseWithConfigEnabled() { + PickFirstLoadBalancerProvider.enablePickFirstConfig = true; + Map rawConfig = new HashMap<>(); + rawConfig.put("shuffleAddressList", true); + ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig( + rawConfig); + PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig(); + + assertThat(config.shuffleAddressList).isTrue(); + } + + @Test + public void parseWithConfigDisabled() { + PickFirstLoadBalancerProvider.enablePickFirstConfig = false; + Map rawConfig = new HashMap<>(); + rawConfig.put("shuffleAddressList", true); + ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig( + rawConfig); + String config = (String) parsedConfig.getConfig(); + + assertThat(config).isEqualTo("no service config"); + } +} diff --git a/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java b/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java index d9eab84aaa..3e0258f2e4 100644 --- a/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java +++ b/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java @@ -52,6 +52,7 @@ import io.grpc.LoadBalancer.SubchannelStateListener; import io.grpc.Status; import io.grpc.Status.Code; import io.grpc.SynchronizationContext; +import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig; import java.net.SocketAddress; import java.util.List; import org.junit.After; @@ -140,6 +141,49 @@ public class PickFirstLoadBalancerTest { verifyNoMoreInteractions(mockHelper); } + @Test + public void pickAfterResolved_shuffle() throws Exception { + loadBalancer.acceptResolvedAddresses( + ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity) + .setLoadBalancingPolicyConfig(new PickFirstLoadBalancerConfig(true, 123L)).build()); + + verify(mockHelper).createSubchannel(createArgsCaptor.capture()); + CreateSubchannelArgs args = createArgsCaptor.getValue(); + // We should still see the same set of addresses. + assertThat(args.getAddresses()).containsExactlyElementsIn(servers); + // Because we use a fixed seed, the addresses should always be shuffled in this order. + assertThat(args.getAddresses().get(0)).isEqualTo(servers.get(1)); + assertThat(args.getAddresses().get(1)).isEqualTo(servers.get(0)); + assertThat(args.getAddresses().get(2)).isEqualTo(servers.get(2)); + verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verify(mockSubchannel).requestConnection(); + + // Calling pickSubchannel() twice gave the same result + assertEquals(pickerCaptor.getValue().pickSubchannel(mockArgs), + pickerCaptor.getValue().pickSubchannel(mockArgs)); + + verifyNoMoreInteractions(mockHelper); + } + + @Test + public void pickAfterResolved_noShuffle() throws Exception { + loadBalancer.acceptResolvedAddresses( + ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity) + .setLoadBalancingPolicyConfig(new PickFirstLoadBalancerConfig(false)).build()); + + verify(mockHelper).createSubchannel(createArgsCaptor.capture()); + CreateSubchannelArgs args = createArgsCaptor.getValue(); + assertThat(args.getAddresses()).isEqualTo(servers); + verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verify(mockSubchannel).requestConnection(); + + // Calling pickSubchannel() twice gave the same result + assertEquals(pickerCaptor.getValue().pickSubchannel(mockArgs), + pickerCaptor.getValue().pickSubchannel(mockArgs)); + + verifyNoMoreInteractions(mockHelper); + } + @Test public void requestConnectionPicker() throws Exception { loadBalancer.acceptResolvedAddresses(