Skip to content

Commit

Permalink
Prepare to switch flag to use new PickFirstLeafLoadBalancer by default (
Browse files Browse the repository at this point in the history 
grpc#10998)

* Fix PickFirstLeafLoadBalancer and tests to work when it is used.
* Actually use EAG attributes for subchannels.
  • Loading branch information
@larry-safran
larry-safran authored Mar 11, 2024
1 parent ebbe067 commit d1c406b
Show file tree
Hide file tree
Showing 13 changed files with 344 additions and 218 deletions.
3 changes: 3 additions & 0 deletions core/src/main/java/io/grpc/internal/GrpcUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -954,6 +954,9 @@ public static boolean getFlag(String envVarName, boolean enableByDefault) {
if (envVar == null) {
envVar = System.getProperty(envVarName);
}
if (envVar != null) {
envVar = envVar.trim();
}
if (enableByDefault) {
return Strings.isNullOrEmpty(envVar) || Boolean.parseBoolean(envVar);
} else {
Expand Down
21 changes: 16 additions & 5 deletions core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,14 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
// If the previous ready subchannel exists in new address list,
// keep this connection and don't create new subchannels
SocketAddress previousAddress = addressIndex.getCurrentAddress();
Attributes prevEagAttrs = addressIndex.getCurrentEagAttributes();
addressIndex.updateGroups(newImmutableAddressGroups);
if (addressIndex.seekTo(previousAddress)) {
if (!addressIndex.getCurrentEagAttributes().equals(prevEagAttrs)) {
log.log(Level.FINE, "EAG attributes changed, need to update subchannel");
SubchannelData subchannelData = subchannels.get(previousAddress);
subchannelData.getSubchannel().updateAddresses(addressIndex.getCurrentEagAsList());
}
return Status.OK;
} else {
addressIndex.reset(); // Previous ready subchannel not in the new list of addresses
Expand Down Expand Up @@ -354,7 +360,7 @@ public void requestConnection() {
currentAddress = addressIndex.getCurrentAddress();
subchannel = subchannels.containsKey(currentAddress)
? subchannels.get(currentAddress).getSubchannel()
: createNewSubchannel(currentAddress);
: createNewSubchannel(currentAddress, addressIndex.getCurrentEagAttributes());

ConnectivityState subchannelState = subchannels.get(currentAddress).getState();
switch (subchannelState) {
Expand Down Expand Up @@ -418,12 +424,12 @@ private void cancelScheduleTask() {
}
}

private Subchannel createNewSubchannel(SocketAddress addr) {
private Subchannel createNewSubchannel(SocketAddress addr, Attributes attrs) {
HealthListener hcListener = new HealthListener();
final Subchannel subchannel = helper.createSubchannel(
CreateSubchannelArgs.newBuilder()
.setAddresses(Lists.newArrayList(
new EquivalentAddressGroup(addr)))
new EquivalentAddressGroup(addr, attrs)))
.addOption(HEALTH_CONSUMER_LISTENER_ARG_KEY, hcListener)
.build());
if (subchannel == null) {
Expand All @@ -433,8 +439,8 @@ private Subchannel createNewSubchannel(SocketAddress addr) {
SubchannelData subchannelData = new SubchannelData(subchannel, IDLE, hcListener);
hcListener.subchannelData = subchannelData;
subchannels.put(addr, subchannelData);
Attributes attrs = subchannel.getAttributes();
if (attrs.get(LoadBalancer.HAS_HEALTH_PRODUCER_LISTENER_KEY) == null) {
Attributes scAttrs = subchannel.getAttributes();
if (scAttrs.get(LoadBalancer.HAS_HEALTH_PRODUCER_LISTENER_KEY) == null) {
hcListener.healthStateInfo = ConnectivityStateInfo.forNonError(READY);
}
subchannel.start(stateInfo -> processSubchannelState(subchannel, stateInfo));
Expand Down Expand Up @@ -584,6 +590,11 @@ public Attributes getCurrentEagAttributes() {
return addressGroups.get(groupIndex).getAttributes();
}

public List<EquivalentAddressGroup> getCurrentEagAsList() {
return Collections.singletonList(
new EquivalentAddressGroup(getCurrentAddress(), getCurrentEagAttributes()));
}

/**
* Update to new groups, resetting the current index.
*/
Expand Down
28 changes: 20 additions & 8 deletions core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,13 @@

package io.grpc.internal;

import com.google.common.base.Strings;
import com.google.common.annotations.VisibleForTesting;
import io.grpc.LoadBalancer;
import io.grpc.LoadBalancerProvider;
import io.grpc.NameResolver;
import io.grpc.NameResolver.ConfigOrError;
import io.grpc.Status;
import io.grpc.internal.PickFirstLeafLoadBalancer.PickFirstLeafLoadBalancerConfig;
import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
import java.util.Map;

Expand All @@ -35,8 +36,7 @@ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider {
private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList";

static boolean enableNewPickFirst =
!Strings.isNullOrEmpty(System.getenv("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST"))
&& Boolean.parseBoolean(System.getenv("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST"));
GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", false);

@Override
public boolean isAvailable() {
Expand All @@ -63,16 +63,28 @@ public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
}

@Override
public ConfigOrError parseLoadBalancingPolicyConfig(
Map<String, ?> rawLoadBalancingPolicyConfig) {
public ConfigOrError parseLoadBalancingPolicyConfig(Map<String, ?> rawLbPolicyConfig) {
try {
return ConfigOrError.fromConfig(
new PickFirstLoadBalancerConfig(JsonUtil.getBoolean(rawLoadBalancingPolicyConfig,
SHUFFLE_ADDRESS_LIST_KEY)));
Object config = getLbPolicyConfig(rawLbPolicyConfig);
return ConfigOrError.fromConfig(config);
} catch (RuntimeException e) {
return ConfigOrError.fromError(
Status.UNAVAILABLE.withCause(e).withDescription(
"Failed parsing configuration for " + getPolicyName()));
}
}

private static Object getLbPolicyConfig(Map<String, ?> rawLbPolicyConfig) {
Boolean shuffleAddressList = JsonUtil.getBoolean(rawLbPolicyConfig, SHUFFLE_ADDRESS_LIST_KEY);
if (enableNewPickFirst) {
return new PickFirstLeafLoadBalancerConfig(shuffleAddressList);
} else {
return new PickFirstLoadBalancerConfig(shuffleAddressList);
}
}

@VisibleForTesting
public static boolean isEnabledNewPickFirst() {
return enableNewPickFirst;
}
}
21 changes: 17 additions & 4 deletions core/src/test/java/io/grpc/internal/AutoConfiguredLoadBalancerFactoryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import io.grpc.NameResolver.ConfigOrError;
import io.grpc.Status;
import io.grpc.internal.AutoConfiguredLoadBalancerFactory.AutoConfiguredLoadBalancer;
import io.grpc.internal.PickFirstLeafLoadBalancer.PickFirstLeafLoadBalancerConfig;
import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
import io.grpc.internal.ServiceConfigUtil.PolicySelection;
import io.grpc.util.ForwardingLoadBalancerHelper;
Expand Down Expand Up @@ -95,6 +96,11 @@ public class AutoConfiguredLoadBalancerFactoryTest {
delegatesTo(
new FakeLoadBalancerProvider("test_lb2", testLbBalancer2, nextParsedConfigOrError2)));

private final Class<? extends LoadBalancer> pfLbClass =
PickFirstLoadBalancerProvider.enableNewPickFirst
? PickFirstLeafLoadBalancer.class
: PickFirstLoadBalancer.class;

@Before
public void setUp() {
when(testLbBalancer.acceptResolvedAddresses(isA(ResolvedAddresses.class))).thenReturn(
Expand Down Expand Up @@ -429,7 +435,7 @@ public Subchannel createSubchannel(CreateSubchannelArgs args) {
.setLoadBalancingPolicyConfig(null)
.build());
assertThat(addressesAcceptanceStatus.isOk()).isTrue();
assertThat(lb.getDelegate()).isInstanceOf(PickFirstLoadBalancer.class);
assertThat(lb.getDelegate()).isInstanceOf(pfLbClass);
}

@Test
Expand Down Expand Up @@ -484,7 +490,7 @@ public Subchannel createSubchannel(CreateSubchannelArgs args) {
verify(channelLogger).log(
eq(ChannelLogLevel.INFO),
eq("Load balancer changed from {0} to {1}"),
eq("PickFirstLoadBalancer"),
eq(pfLbClass.getSimpleName()),
eq(testLbBalancer.getClass().getSimpleName()));

verify(channelLogger).log(
Expand Down Expand Up @@ -628,8 +634,15 @@ public void parseLoadBalancerConfig_lbConfigPropagated() throws Exception {
assertThat(parsed.getConfig()).isNotNull();
PolicySelection policySelection = (PolicySelection) parsed.getConfig();
assertThat(policySelection.provider).isInstanceOf(PickFirstLoadBalancerProvider.class);
assertThat(policySelection.config).isInstanceOf(PickFirstLoadBalancerConfig.class);
assertThat(((PickFirstLoadBalancerConfig) policySelection.config).shuffleAddressList).isTrue();
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
assertThat(policySelection.config).isInstanceOf(PickFirstLeafLoadBalancerConfig.class);
assertThat(((PickFirstLeafLoadBalancerConfig) policySelection.config).shuffleAddressList)
.isTrue();
} else {
assertThat(policySelection.config).isInstanceOf(PickFirstLoadBalancerConfig.class);
assertThat(((PickFirstLoadBalancerConfig) policySelection.config).shuffleAddressList)
.isTrue();
}
verifyNoInteractions(channelLogger);
}

Expand Down
39 changes: 33 additions & 6 deletions core/src/test/java/io/grpc/internal/PickFirstLoadBalancerProviderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import static com.google.common.truth.Truth.assertThat;

import io.grpc.NameResolver.ConfigOrError;
import io.grpc.internal.PickFirstLeafLoadBalancer.PickFirstLeafLoadBalancerConfig;
import io.grpc.internal.PickFirstLoadBalancer.PickFirstLoadBalancerConfig;
import java.util.HashMap;
import java.util.Map;
Expand All @@ -35,20 +36,46 @@ public void parseWithConfig() {
rawConfig.put("shuffleAddressList", true);
ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig(
rawConfig);
PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig();

assertThat(config.shuffleAddressList).isTrue();
assertThat(config.randomSeed).isNull();
Boolean shuffleAddressList;
Long randomSeed;

if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
PickFirstLeafLoadBalancerConfig config =
(PickFirstLeafLoadBalancerConfig) parsedConfig.getConfig();
shuffleAddressList = config.shuffleAddressList;
randomSeed = config.randomSeed;
} else {
PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig();
shuffleAddressList = config.shuffleAddressList;
randomSeed = config.randomSeed;
}

assertThat(shuffleAddressList).isTrue();
assertThat(randomSeed).isNull();
}

@Test
public void parseWithoutConfig() {
Map<String, Object> rawConfig = new HashMap<>();
ConfigOrError parsedConfig = new PickFirstLoadBalancerProvider().parseLoadBalancingPolicyConfig(
rawConfig);
PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig();

assertThat(config.shuffleAddressList).isNull();
assertThat(config.randomSeed).isNull();
Boolean shuffleAddressList;
Long randomSeed;

if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
PickFirstLeafLoadBalancerConfig config =
(PickFirstLeafLoadBalancerConfig) parsedConfig.getConfig();
shuffleAddressList = config.shuffleAddressList;
randomSeed = config.randomSeed;
} else {
PickFirstLoadBalancerConfig config = (PickFirstLoadBalancerConfig) parsedConfig.getConfig();
shuffleAddressList = config.shuffleAddressList;
randomSeed = config.randomSeed;
}

assertThat(shuffleAddressList).isNull();
assertThat(randomSeed).isNull();
}
}
Loading

0 comments on commit d1c406b

Please sign in to comment.