From 3c01bfe77f3d64d9ab97b28bc924105b3e72e8a7 Mon Sep 17 00:00:00 2001 From: Terry Wilson Date: Fri, 28 Apr 2023 15:50:18 -0700 Subject: [PATCH] core: Optional address shuffle in PickFirstLoadBalancer (#10110) If provided with the new PickFirstLoadBalancerConfig, PickFirstLoadBalancer will shuffle the list of addresses it receives from the name resolver. PickFirstLoadBalancerProvider will now support the new config if enabled by an env variable. --- .../grpc/internal/PickFirstLoadBalancer.java | 34 +++++++++++ .../PickFirstLoadBalancerProvider.java | 15 ++++- .../PickFirstLoadBalancerProviderTest.java | 61 +++++++++++++++++++ .../internal/PickFirstLoadBalancerTest.java | 44 +++++++++++++ 4 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java 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(