io.grpc.LoadBalancerProvider Java Examples

/**
 * Called either when the child is just created and in this case updated with the cached {@code
 * resolvedAddresses}, or when priority lb receives a new resolved addresses while the child
 * already exists.
 */
void updateResolvedAddresses() {
  final ResolvedAddresses addresses = resolvedAddresses;
  syncContext.execute(
      new Runnable() {
        @Override
        public void run() {
          PriorityLbConfig config = (PriorityLbConfig) addresses.getLoadBalancingPolicyConfig();
          PolicySelection childPolicySelection = config.childConfigs.get(priority);
          LoadBalancerProvider lbProvider = childPolicySelection.getProvider();
          String newPolicy = lbProvider.getPolicyName();
          if (!newPolicy.equals(policy)) {
            policy = newPolicy;
            lb.switchTo(lbProvider);
          }
          // TODO(zdapeng): Implement address filtering.
          lb.handleResolvedAddresses(
              addresses
                  .toBuilder()
                  .setAddresses(AddressFilter.filter(addresses.getAddresses(), priority))
                  .setLoadBalancingPolicyConfig(childPolicySelection.getConfig())
                  .build());
        }
      });
}
 

private void startFallbackChildPolicy() {
  String defaultTarget = lbPolicyConfig.getRouteLookupConfig().getDefaultTarget();
  fallbackChildPolicyWrapper = refCountedChildPolicyWrapperFactory.createOrGet(defaultTarget);

  LoadBalancerProvider lbProvider =
      lbPolicyConfig.getLoadBalancingPolicy().getEffectiveLbProvider();
  final LoadBalancer lb =
      lbProvider.newLoadBalancer(fallbackChildPolicyWrapper.getHelper());
  final ConfigOrError lbConfig =
      lbProvider
          .parseLoadBalancingPolicyConfig(
              lbPolicyConfig
                  .getLoadBalancingPolicy()
                  .getEffectiveChildPolicy(defaultTarget));
  helper.getSynchronizationContext().execute(
      new Runnable() {
        @Override
        public void run() {
          lb.handleResolvedAddresses(
              childLbResolvedAddressFactory.create(lbConfig.getConfig()));
          lb.requestConnection();
        }
      });
}
 

@Test
public void childLoadBalancingPolicy_noPolicyProvided() {
  LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class);
  when(mockProvider.getPolicyName()).thenReturn("rls");
  when(mockProvider.isAvailable()).thenReturn(true);

  LoadBalancerRegistry.getDefaultRegistry().register(mockProvider);
  try {
    ChildLoadBalancingPolicy.create(
        "targetFieldName",
        ImmutableList.<Map<String, ?>>of(
            ImmutableMap.<String, Object>of(
                "rls", ImmutableMap.of(), "rls2", ImmutableMap.of())));
    fail("parsing exception expected");
  } catch (InvalidChildPolicyConfigException e) {
    assertThat(e).hasMessageThat()
        .contains("childPolicy should have exactly one loadbalancing policy");
  } finally {
    LoadBalancerRegistry.getDefaultRegistry().deregister(mockProvider);
  }
}
 

@VisibleForTesting
static LoadBalancerProvider newRoundRobinProvider() {
  try {
    Class<? extends LoadBalancerProvider> rrProviderClass =
        Class.forName("io.grpc.util.SecretRoundRobinLoadBalancerProvider$Provider")
            .asSubclass(LoadBalancerProvider.class);
    return rrProviderClass.getDeclaredConstructor().newInstance();
  } catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

@SuppressWarnings("unchecked")
@VisibleForTesting
static LoadBalancerProvider newRoundRobinProvider() {
  try {
    Class<LoadBalancerProvider> rrProviderClass =
        (Class<LoadBalancerProvider>) Class.forName(
            "io.grpc.util.SecretRoundRobinLoadBalancerProvider$Provider");
    return rrProviderClass.getDeclaredConstructor().newInstance();
  } catch (Exception e) {
    Throwables.throwIfUnchecked(e);
    throw new RuntimeException(e);
  }
}
 

@VisibleForTesting
ChildLoadBalancingPolicy(
    String targetFieldName,
    Map<String, Object> effectiveRawChildPolicy,
    LoadBalancerProvider effectiveLbProvider) {
  checkArgument(
      targetFieldName != null && !targetFieldName.isEmpty(),
      "targetFieldName cannot be empty or null");
  this.targetFieldName = targetFieldName;
  this.effectiveRawChildPolicy =
      checkNotNull(effectiveRawChildPolicy, "effectiveRawChildPolicy");
  this.effectiveLbProvider = checkNotNull(effectiveLbProvider, "effectiveLbProvider");
}
 

/** Creates ChildLoadBalancingPolicy. */
@SuppressWarnings("unchecked")
static ChildLoadBalancingPolicy create(
    String childPolicyConfigTargetFieldName, List<Map<String, ?>> childPolicies)
    throws InvalidChildPolicyConfigException {
  Map<String, Object> effectiveChildPolicy = null;
  LoadBalancerProvider effectiveLbProvider = null;
  List<String> policyTried = new ArrayList<>();

  LoadBalancerRegistry lbRegistry = LoadBalancerRegistry.getDefaultRegistry();
  for (Map<String, ?> childPolicy : childPolicies) {
    if (childPolicy.isEmpty()) {
      continue;
    }
    if (childPolicy.size() != 1) {
      throw
          new InvalidChildPolicyConfigException(
              "childPolicy should have exactly one loadbalancing policy");
    }
    String policyName = childPolicy.keySet().iterator().next();
    LoadBalancerProvider provider = lbRegistry.getProvider(policyName);
    if (provider != null) {
      effectiveLbProvider = provider;
      effectiveChildPolicy = Collections.unmodifiableMap(childPolicy);
      break;
    }
    policyTried.add(policyName);
  }
  if (effectiveChildPolicy == null) {
    throw
        new InvalidChildPolicyConfigException(
            String.format("no valid childPolicy found, policy tried: %s", policyTried));
  }
  return
      new ChildLoadBalancingPolicy(
          childPolicyConfigTargetFieldName,
          (Map<String, Object>) effectiveChildPolicy.values().iterator().next(),
          effectiveLbProvider);
}
 

private void createChildLbPolicy() {
  ChildLoadBalancingPolicy childPolicy = lbPolicyConfig.getLoadBalancingPolicy();
  LoadBalancerProvider lbProvider = childPolicy.getEffectiveLbProvider();
  ConfigOrError lbConfig =
      lbProvider
          .parseLoadBalancingPolicyConfig(
              childPolicy.getEffectiveChildPolicy(childPolicyWrapper.getTarget()));

  LoadBalancer lb = lbProvider.newLoadBalancer(childPolicyWrapper.getHelper());
  lb.handleResolvedAddresses(childLbResolvedAddressFactory.create(lbConfig.getConfig()));
  lb.requestConnection();
}
 

@Test
public void childLoadBalancingPolicy_effectiveChildPolicy() {
  LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class);
  ChildLoadBalancingPolicy childLbPolicy =
      new ChildLoadBalancingPolicy(
          "targetFieldName",
          ImmutableMap.<String, Object>of("foo", "bar"),
          mockProvider);

  assertThat(childLbPolicy.getEffectiveChildPolicy("target"))
      .containsExactly("foo", "bar", "targetFieldName", "target");
  assertThat(childLbPolicy.getEffectiveLbProvider()).isEqualTo(mockProvider);
}
 

@SuppressWarnings("ExpectedExceptionChecker")
@Test
public void priorityLbConfig_missingChildConfig() {
  Map<String, PolicySelection> childConfigs =
      ImmutableMap.of("p1", new PolicySelection(mock(LoadBalancerProvider.class), null, null));
  List<String> priorities = ImmutableList.of("p0", "p1");

  thrown.expect(IllegalArgumentException.class);
  new PriorityLbConfig(childConfigs, priorities);
}
 

@Test
public void registry() {
  LoadBalancerProvider hcRoundRobin =
      LoadBalancerRegistry.getDefaultRegistry().getProvider("round_robin");
  assertThat(hcRoundRobin).isInstanceOf(
      HealthCheckingRoundRobinLoadBalancerProvider.class);
}
 

@Test
public void policyName() {
  LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
  assertThat(hcRoundRobin.getPolicyName())
      .isEqualTo(
          HealthCheckingRoundRobinLoadBalancerProvider.newRoundRobinProvider().getPolicyName());
}
 

@Test
public void priority() {
  LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
  assertThat(hcRoundRobin.getPriority())
      .isEqualTo(
          HealthCheckingRoundRobinLoadBalancerProvider.newRoundRobinProvider().getPriority() + 1);
}
 

/**
 * Parses and selects a load balancing policy from a non-empty list of raw configs. If selection
 * is successful, the returned ConfigOrError object will include a {@link
 * ServiceConfigUtil.PolicySelection} as its config value.
 */
public static ConfigOrError selectLbPolicyFromList(
    List<LbConfig> lbConfigs, LoadBalancerRegistry lbRegistry) {
  List<String> policiesTried = new ArrayList<>();
  for (LbConfig lbConfig : lbConfigs) {
    String policy = lbConfig.getPolicyName();
    LoadBalancerProvider provider = lbRegistry.getProvider(policy);
    if (provider == null) {
      policiesTried.add(policy);
    } else {
      if (!policiesTried.isEmpty()) {
        Logger.getLogger(ServiceConfigUtil.class.getName()).log(
            Level.FINEST,
            "{0} specified by Service Config are not available", policiesTried);
      }
      ConfigOrError parsedLbPolicyConfig =
          provider.parseLoadBalancingPolicyConfig(lbConfig.getRawConfigValue());
      if (parsedLbPolicyConfig.getError() != null) {
        return parsedLbPolicyConfig;
      }
      return ConfigOrError.fromConfig(new PolicySelection(
          provider, lbConfig.rawConfigValue, parsedLbPolicyConfig.getConfig()));
    }
  }
  return ConfigOrError.fromError(
      Status.UNKNOWN.withDescription(
          "None of " + policiesTried + " specified by Service Config are available."));
}
 

/** Constructs a PolicySelection with selected LB provider, a copy of raw config and the deeply
 * parsed LB config. */
public PolicySelection(
    LoadBalancerProvider provider,
    @Nullable Map<String, ?> rawConfig,
    @Nullable Object config) {
  this.provider = checkNotNull(provider, "provider");
  this.rawConfig = rawConfig;
  this.config = config;
}
 

private LoadBalancerProvider getProviderOrThrow(String policy, String choiceReason)
    throws PolicyException {
  LoadBalancerProvider provider = registry.getProvider(policy);
  if (provider == null) {
    throw new PolicyException(
        "Trying to load '" + policy + "' because " + choiceReason + ", but it's unavailable");
  }
  return provider;
}
 

@Before
public void setUp() {
  for (String lbPolicy : lbPolicies) {
    LoadBalancerProvider lbProvider = new FakeLoadBalancerProvider(lbPolicy);
    lbProviders.put(lbPolicy, lbProvider);
    lbRegistry.register(lbProvider);
  }
}
 

private static LoadBalancerProvider getProviderOrThrow(String policy, String reason)
    throws PolicyNotFoundException {
  LoadBalancerProvider provider = registry.getProvider(policy);
  if (provider == null) {
    throw new PolicyNotFoundException(policy, reason);
  }
  return provider;
}
 

@Test
public void registry() {
  LoadBalancerProvider hcRoundRobin =
      LoadBalancerRegistry.getDefaultRegistry().getProvider("round_robin");
  assertThat(hcRoundRobin).isInstanceOf(
      HealthCheckingRoundRobinLoadBalancerProvider.class);
}
 

@Test
public void policyName() {
  LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
  assertThat(hcRoundRobin.getPolicyName())
      .isEqualTo(
          HealthCheckingRoundRobinLoadBalancerProvider.newRoundRobinProvider().getPolicyName());
}
 

@Test
public void priority() {
  LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
  assertThat(hcRoundRobin.getPriority())
      .isEqualTo(
          HealthCheckingRoundRobinLoadBalancerProvider.newRoundRobinProvider().getPriority() + 1);
}
 

@Override
public void handleResolvedAddressGroups(
    List<EquivalentAddressGroup> servers, Attributes attributes) {
  Map<String, Object> configMap = attributes.get(GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG);
  LoadBalancerProvider newlbp;
  try {
    newlbp = decideLoadBalancerProvider(servers, configMap);
  } catch (PolicyNotFoundException e) {
    Status s = Status.INTERNAL.withDescription(e.getMessage());
    helper.updateBalancingState(ConnectivityState.TRANSIENT_FAILURE, new FailingPicker(s));
    delegate.shutdown();
    delegateProvider = null;
    delegate = new NoopLoadBalancer();
    return;
  }

  if (delegateProvider == null
      || !newlbp.getPolicyName().equals(delegateProvider.getPolicyName())) {
    helper.updateBalancingState(ConnectivityState.CONNECTING, new EmptyPicker());
    delegate.shutdown();
    delegateProvider = newlbp;
    LoadBalancer old = delegate;
    delegate = delegateProvider.newLoadBalancer(helper);
    helper.getChannelLogger().log(
        ChannelLogLevel.INFO, "Load balancer changed from {0} to {1}",
        old.getClass().getSimpleName(), delegate.getClass().getSimpleName());
  }
  getDelegate().handleResolvedAddressGroups(servers, attributes);
}
 

/**
 * Picks a load balancer based on given criteria.  In order of preference:
 *
 * <ol>
 *   <li>User provided lb on the channel.  This is a degenerate case and not handled here.</li>
 *   <li>"grpclb" if any gRPC LB balancer addresses are present</li>
 *   <li>The policy picked by the service config</li>
 *   <li>"pick_first" if the service config choice does not specify</li>
 * </ol>
 *
 * @param servers The list of servers reported
 * @param config the service config object
 * @return the new load balancer factory, never null
 */
@VisibleForTesting
static LoadBalancerProvider decideLoadBalancerProvider(
    List<EquivalentAddressGroup> servers, @Nullable Map<String, Object> config)
    throws PolicyNotFoundException {
  // Check for balancer addresses
  boolean haveBalancerAddress = false;
  for (EquivalentAddressGroup s : servers) {
    if (s.getAttributes().get(GrpcAttributes.ATTR_LB_ADDR_AUTHORITY) != null) {
      haveBalancerAddress = true;
      break;
    }
  }

  if (haveBalancerAddress) {
    return getProviderOrThrow("grpclb", "NameResolver has returned balancer addresses");
  }

  String serviceConfigChoiceBalancingPolicy = null;
  if (config != null) {
    serviceConfigChoiceBalancingPolicy =
        ServiceConfigUtil.getLoadBalancingPolicyFromServiceConfig(config);
    if (serviceConfigChoiceBalancingPolicy != null) {
      // Handle ASCII specifically rather than relying on the implicit default locale of the str
      return getProviderOrThrow(
          Ascii.toLowerCase(serviceConfigChoiceBalancingPolicy),
          "service-config specifies load-balancing policy");
    }
  }
  return getProviderOrThrow(DEFAULT_POLICY, "Using default policy");
}
 

@SuppressWarnings("ExpectedExceptionChecker")
@Test
public void priorityLbConfig_emptyPriorities() {
  Map<String, PolicySelection> childConfigs =
      ImmutableMap.of("p0", new PolicySelection(mock(LoadBalancerProvider.class), null, null));
  List<String> priorities = ImmutableList.of();

  thrown.expect(IllegalArgumentException.class);
  new PriorityLbConfig(childConfigs, priorities);
}
 

@Test
public void decideLoadBalancerProvider_noBalancerAddresses_noServiceConfig_pickFirst()
    throws Exception {
  Map<String, Object> serviceConfig = null;
  List<EquivalentAddressGroup> servers =
      Collections.singletonList(
          new EquivalentAddressGroup(new SocketAddress(){}, Attributes.EMPTY));
  LoadBalancerProvider provider =
      AutoConfiguredLoadBalancer.decideLoadBalancerProvider(servers, serviceConfig);

  assertThat(provider).isInstanceOf(PickFirstLoadBalancerProvider.class);
}
 

@Test
public void decideLoadBalancerProvider_oneBalancer_noServiceConfig_grpclb() throws Exception {
  Map<String, Object> serviceConfig = null;
  List<EquivalentAddressGroup> servers =
      Collections.singletonList(
          new EquivalentAddressGroup(
              new SocketAddress(){},
              Attributes.newBuilder().set(GrpcAttributes.ATTR_LB_ADDR_AUTHORITY, "ok").build()));
  LoadBalancerProvider provider = AutoConfiguredLoadBalancer.decideLoadBalancerProvider(
      servers, serviceConfig);

  assertThat(provider).isInstanceOf(GrpclbLoadBalancerProvider.class);
}
 

@Test
public void decideLoadBalancerProvider_grpclbOverridesServiceConfig() throws Exception {
  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()));
  LoadBalancerProvider provider = AutoConfiguredLoadBalancer.decideLoadBalancerProvider(
      servers, serviceConfig);

  assertThat(provider).isInstanceOf(GrpclbLoadBalancerProvider.class);
}
 

@Test
public void decideLoadBalancerProvider_serviceConfigOverridesDefault() throws Exception {
  Map<String, Object> serviceConfig = new HashMap<String, Object>();
  serviceConfig.put("loadBalancingPolicy", "round_robin");
  List<EquivalentAddressGroup> servers =
      Collections.singletonList(
          new EquivalentAddressGroup(
              new SocketAddress(){},
              Attributes.EMPTY));
  LoadBalancerProvider provider = AutoConfiguredLoadBalancer.decideLoadBalancerProvider(
      servers, serviceConfig);

  assertThat(provider.getClass().getName()).isEqualTo(
      "io.grpc.util.SecretRoundRobinLoadBalancerProvider$Provider");
}
 

void reactivate(LoadBalancerProvider policyProvider) {
  if (deletionTimer != null && deletionTimer.isPending()) {
    deletionTimer.cancel();
    deactivated = false;
    logger.log(XdsLogLevel.DEBUG, "Route action {0} reactivated", name);
  }
  if (!this.policyProvider.getPolicyName().equals(policyProvider.getPolicyName())) {
    logger.log(
        XdsLogLevel.DEBUG,
        "Action {0} switching policy from {1} to {2}",
        name, this.policyProvider.getPolicyName(), policyProvider.getPolicyName());
    lb.switchTo(policyProvider);
    this.policyProvider = policyProvider;
  }
}
 

@Override
public void onClusterChanged(ClusterUpdate newUpdate) {
  if (logger.isLoggable(XdsLogLevel.INFO)) {
    logger.log(
        XdsLogLevel.INFO,
        "Received cluster update from xDS client {0}: "
            + "cluster_name={1}, eds_service_name={2}, lb_policy={3}, report_load={4}",
        xdsClient, newUpdate.getClusterName(), newUpdate.getEdsServiceName(),
        newUpdate.getLbPolicy(), newUpdate.getLrsServerName() != null);
  }
  checkArgument(
      newUpdate.getLbPolicy().equals("round_robin"), "can only support round_robin policy");

  LoadBalancerProvider lbProvider = lbRegistry.getProvider(newUpdate.getLbPolicy());
  Object lbConfig =
      lbProvider.parseLoadBalancingPolicyConfig(ImmutableMap.<String, Object>of()).getConfig();
  final EdsConfig edsConfig =
      new EdsConfig(
          /* clusterName = */ newUpdate.getClusterName(),
          /* edsServiceName = */ newUpdate.getEdsServiceName(),
          /* lrsServerName = */ newUpdate.getLrsServerName(),
          new PolicySelection(lbProvider, ImmutableMap.<String, Object>of(), lbConfig));
  if (isXdsSecurityEnabled()) {
    updateSslContextProvider(newUpdate.getUpstreamTlsContext());
  }
  if (edsBalancer == null) {
    edsBalancer = lbRegistry.getProvider(EDS_POLICY_NAME).newLoadBalancer(helper);
  }
  edsBalancer.handleResolvedAddresses(
      resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(edsConfig).build());
}