Skip to content

Commit

Permalink
core: make Auto config load balancer not depend on service config
Browse files Browse the repository at this point in the history 
Also, add some tests
  • Loading branch information
@carl-mastrangelo
carl-mastrangelo authored Jun 26, 2018
1 parent 7e68a0b commit 81da3eb
Show file tree
Hide file tree
Showing 3 changed files with 209 additions and 26 deletions.
3 changes: 2 additions & 1 deletion core/build.gradle
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ dependencies {
}

testCompile project(':grpc-context').sourceSets.test.output,
project(':grpc-testing')
project(':grpc-testing'),
project(':grpc-grpclb')

signature "org.codehaus.mojo.signature:java16:1.1@signature"
}
Expand Down
42 changes: 18 additions & 24 deletions core/src/main/java/io/grpc/internal/AutoConfiguredLoadBalancerFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
package io.grpc.internal;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import io.grpc.Attributes;
import io.grpc.ConnectivityState;
import io.grpc.ConnectivityStateInfo;
Expand Down Expand Up @@ -67,23 +66,21 @@ static final class AutoConfiguredLoadBalancer extends LoadBalancer {

AutoConfiguredLoadBalancer(Helper helper) {
this.helper = helper;
setDelegateFactory(PickFirstBalancerFactory.getInstance());
setDelegate(getDelegateFactory().newLoadBalancer(helper));
delegateFactory = PickFirstBalancerFactory.getInstance();
delegate = delegateFactory.newLoadBalancer(helper);
}

// Must be run inside ChannelExecutor.
@Override
public void handleResolvedAddressGroups(
List<EquivalentAddressGroup> servers, Attributes attributes) {
Map<String, Object> configMap = attributes.get(GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG);
if (configMap != null) {
Factory newlbf = decideLoadBalancerFactory(servers, configMap);
if (newlbf != null && newlbf != delegateFactory) {
helper.updateBalancingState(ConnectivityState.CONNECTING, new EmptySubchannelPicker());
getDelegate().shutdown();
setDelegateFactory(newlbf);
setDelegate(getDelegateFactory().newLoadBalancer(helper));
}
Factory newlbf = decideLoadBalancerFactory(servers, configMap);
if (newlbf != null && newlbf != delegateFactory) {
helper.updateBalancingState(ConnectivityState.CONNECTING, new EmptySubchannelPicker());
delegate.shutdown();
delegateFactory = newlbf;
delegate = delegateFactory.newLoadBalancer(helper);
}
getDelegate().handleResolvedAddressGroups(servers, attributes);
}
Expand All @@ -100,8 +97,8 @@ public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo s

@Override
public void shutdown() {
getDelegate().shutdown();
setDelegate(null);
delegate.shutdown();
delegate = null;
}

@VisibleForTesting
Expand All @@ -110,20 +107,15 @@ LoadBalancer getDelegate() {
}

@VisibleForTesting
void setDelegate(LoadBalancer delegate) {
this.delegate = delegate;
void setDelegate(LoadBalancer lb) {
delegate = lb;
}

@VisibleForTesting
LoadBalancer.Factory getDelegateFactory() {
return delegateFactory;
}

@VisibleForTesting
void setDelegateFactory(LoadBalancer.Factory delegateFactory) {
this.delegateFactory = delegateFactory;
}

/**
* Picks a load balancer based on given criteria. In order of preference:
*
Expand All @@ -141,8 +133,7 @@ void setDelegateFactory(LoadBalancer.Factory delegateFactory) {
@Nullable
@VisibleForTesting
static LoadBalancer.Factory decideLoadBalancerFactory(
List<EquivalentAddressGroup> servers, Map<String, Object> config) {
Preconditions.checkNotNull(config);
List<EquivalentAddressGroup> servers, @Nullable Map<String, Object> config) {
// Check for balancer addresses
boolean haveBalancerAddress = false;
for (EquivalentAddressGroup s : servers) {
Expand All @@ -164,8 +155,11 @@ static LoadBalancer.Factory decideLoadBalancerFactory(
}
}

String serviceConfigChoiceBalancingPolicy =
ServiceConfigUtil.getLoadBalancingPolicyFromServiceConfig(config);
String serviceConfigChoiceBalancingPolicy = null;
if (config != null) {
serviceConfigChoiceBalancingPolicy =
ServiceConfigUtil.getLoadBalancingPolicyFromServiceConfig(config);
}

// Check for an explicitly present lb choice
if (serviceConfigChoiceBalancingPolicy != null) {
Expand Down
190 changes: 189 additions & 1 deletion core/src/test/java/io/grpc/internal/AutoConfiguredLoadBalancerFactoryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
package io.grpc.internal;

import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import io.grpc.Attributes;
import io.grpc.ConnectivityState;
Expand All @@ -30,8 +32,15 @@
import io.grpc.NameResolver.Factory;
import io.grpc.PickFirstBalancerFactory;
import io.grpc.Status;
import io.grpc.grpclb.GrpclbLoadBalancerFactory;
import io.grpc.internal.AutoConfiguredLoadBalancerFactory.AutoConfiguredLoadBalancer;
import io.grpc.util.RoundRobinLoadBalancerFactory;
import java.net.SocketAddress;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.Nonnull;
import org.junit.Test;
Expand All @@ -41,7 +50,6 @@
/**
* Unit tests for {@link AutoConfiguredLoadBalancerFactory}.
*/
// TODO(carl-mastrangelo): Add tests for selection logic.
@RunWith(JUnit4.class)
public class AutoConfiguredLoadBalancerFactoryTest {
private final AutoConfiguredLoadBalancerFactory lbf = new AutoConfiguredLoadBalancerFactory();
Expand Down Expand Up @@ -98,6 +106,158 @@ public void shutdown() {
assertThat(calls.getAndSet(0)).isEqualTo(3);
}

@Test
public void handleResolvedAddressGroups_keepOldBalancer() {
final List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(new SocketAddress(){}, Attributes.EMPTY));
Helper helper = new TestHelper() {
@Override
public Subchannel createSubchannel(EquivalentAddressGroup addrs, Attributes attrs) {
assertThat(addrs).isEqualTo(servers.get(0));
return new TestSubchannel(addrs, attrs);
}

@Override
public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
// noop
}
};
AutoConfiguredLoadBalancer lb =
(AutoConfiguredLoadBalancer) lbf.newLoadBalancer(helper);
LoadBalancer oldDelegate = lb.getDelegate();

lb.handleResolvedAddressGroups(servers, Attributes.EMPTY);

assertThat(lb.getDelegate()).isSameAs(oldDelegate);
}

@Test
public void handleResolvedAddressGroups_shutsDownOldBalancer() {
Map<String, Object> serviceConfig = new HashMap<String, Object>();
serviceConfig.put("loadBalancingPolicy", "round_robin");
Attributes serviceConfigAttrs =
Attributes.newBuilder()
.set(GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG, serviceConfig)
.build();
final List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(
new SocketAddress(){},
Attributes.EMPTY));
Helper helper = new TestHelper() {
@Override
public Subchannel createSubchannel(final EquivalentAddressGroup addrs, Attributes attrs) {
assertThat(addrs).isEqualTo(servers.get(0));
return new TestSubchannel(addrs, attrs);
}

@Override
public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
// noop
}
};
AutoConfiguredLoadBalancer lb =
(AutoConfiguredLoadBalancer) lbf.newLoadBalancer(helper);
final AtomicBoolean shutdown = new AtomicBoolean();
TestLoadBalancer testlb = new TestLoadBalancer() {

@Override
public void handleNameResolutionError(Status error) {
// noop
}

@Override
public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
// noop
}

@Override
public void shutdown() {
shutdown.set(true);
}
};
lb.setDelegate(testlb);

lb.handleResolvedAddressGroups(servers, serviceConfigAttrs);

assertThat(lb.getDelegateFactory()).isEqualTo(RoundRobinLoadBalancerFactory.getInstance());
assertTrue(shutdown.get());
}

@Test
public void decideLoadBalancerFactory_noBalancerAddresses_noServiceConfig_pickFirst() {
Map<String, Object> serviceConfig = null;
List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(new SocketAddress(){}, Attributes.EMPTY));
LoadBalancer.Factory factory =
AutoConfiguredLoadBalancer.decideLoadBalancerFactory(servers, serviceConfig);

assertThat(factory).isInstanceOf(PickFirstBalancerFactory.class);
}

@Test
public void decideLoadBalancerFactory_oneBalancer_noServiceConfig_grpclb() {
Map<String, Object> serviceConfig = null;
List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(
new SocketAddress(){},
Attributes.newBuilder().set(GrpcAttributes.ATTR_LB_ADDR_AUTHORITY, "ok").build()));
LoadBalancer.Factory factory = AutoConfiguredLoadBalancer.decideLoadBalancerFactory(
servers, serviceConfig);

assertThat(factory).isInstanceOf(GrpclbLoadBalancerFactory.class);
}

@Test
public void decideLoadBalancerFactory_grpclbOverridesServiceConfig() {
Map<String, Object> serviceConfig = new HashMap<String, Object>();
serviceConfig.put("loadBalancingPolicy", "round_robin");
List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(
new SocketAddress(){},
Attributes.newBuilder().set(GrpcAttributes.ATTR_LB_ADDR_AUTHORITY, "ok").build()));
LoadBalancer.Factory factory = AutoConfiguredLoadBalancer.decideLoadBalancerFactory(
servers, serviceConfig);

assertThat(factory).isInstanceOf(GrpclbLoadBalancerFactory.class);
}

@Test
public void decideLoadBalancerFactory_serviceConfigOverridesDefault() {
Map<String, Object> serviceConfig = new HashMap<String, Object>();
serviceConfig.put("loadBalancingPolicy", "round_robin");
List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(
new SocketAddress(){},
Attributes.EMPTY));
LoadBalancer.Factory factory = AutoConfiguredLoadBalancer.decideLoadBalancerFactory(
servers, serviceConfig);

assertThat(factory).isInstanceOf(RoundRobinLoadBalancerFactory.class);
}

@Test
public void decideLoadBalancerFactory_serviceConfigFailsOnUnknown() {
Map<String, Object> serviceConfig = new HashMap<String, Object>();
serviceConfig.put("loadBalancingPolicy", "MAGIC_BALANCER");
List<EquivalentAddressGroup> servers =
Collections.singletonList(
new EquivalentAddressGroup(
new SocketAddress(){},
Attributes.EMPTY));
try {
AutoConfiguredLoadBalancer.decideLoadBalancerFactory(servers, serviceConfig);
fail();
} catch (IllegalArgumentException e) {
// expected
}
}

public static class ForwardingLoadBalancer extends LoadBalancer {
private final LoadBalancer delegate;

Expand Down Expand Up @@ -186,4 +346,32 @@ private static class TestHelper extends ForwardingLoadBalancerHelper {
super(null);
}
}

private static class TestSubchannel extends Subchannel {
TestSubchannel(EquivalentAddressGroup addrs, Attributes attrs) {
this.addrs = addrs;
this.attrs = attrs;
}

final EquivalentAddressGroup addrs;
final Attributes attrs;

@Override
public void shutdown() {
}

@Override
public void requestConnection() {
}

@Override
public EquivalentAddressGroup getAddresses() {
return addrs;
}

@Override
public Attributes getAttributes() {
return attrs;
}
}
}

0 comments on commit 81da3eb

Please sign in to comment.