xds: refactor XdsConfig to use PolicySelection

Author: dapengzhang0

xds/src/main/java/io/grpc/xds/CdsLoadBalancer.java

@@ -27,10 +27,11 @@
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.InternalLogId;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.Status;
 import io.grpc.internal.ObjectPool;
-import io.grpc.internal.ServiceConfigUtil.LbConfig;
+import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.util.ForwardingLoadBalancerHelper;
 import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.CdsLoadBalancerProvider.CdsConfig;
@@ -283,10 +284,13 @@ public void onClusterChanged(ClusterUpdate newUpdate) {
       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 XdsConfig edsConfig =
           new XdsConfig(
               /* cluster = */ newUpdate.getClusterName(),
-              new LbConfig(newUpdate.getLbPolicy(), ImmutableMap.<String, Object>of()),
+              new PolicySelection(lbProvider, ImmutableMap.<String, Object>of(), lbConfig),
               /* fallbackPolicy = */ null,
               /* edsServiceName = */ newUpdate.getEdsServiceName(),
               /* lrsServerName = */ newUpdate.getLrsServerName());

xds/src/main/java/io/grpc/xds/EdsLoadBalancer.java

@@ -121,8 +121,10 @@ public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
           "Received EDS lb config: cluster={0}, child_policy={1}, fallback_policy={2}, "
               + "eds_service_name={3}, report_load={4}",
           newXdsConfig.cluster,
-          newXdsConfig.childPolicy != null ? newXdsConfig.childPolicy.getPolicyName() : "",
-          newXdsConfig.fallbackPolicy != null ? newXdsConfig.fallbackPolicy.getPolicyName() : "",
+          newXdsConfig.endpointPickingPolicy != null
+              ? newXdsConfig.endpointPickingPolicy.getProvider().getPolicyName() : "",
+          newXdsConfig.fallbackPolicy != null
+              ? newXdsConfig.fallbackPolicy.getProvider().getPolicyName() : "",
           newXdsConfig.edsServiceName,
           newXdsConfig.lrsServerName != null);
     }

xds/src/main/java/io/grpc/xds/FallbackLb.java

@@ -19,37 +19,24 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 
-import com.google.common.annotations.VisibleForTesting;
-import io.grpc.Attributes;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
-import io.grpc.LoadBalancerRegistry;
-import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
-import io.grpc.internal.ServiceConfigUtil.LbConfig;
+import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.util.ForwardingLoadBalancer;
 import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.XdsLoadBalancerProvider.XdsConfig;
 import io.grpc.xds.XdsSubchannelPickers.ErrorPicker;
 import java.util.List;
-import java.util.Map;
 
 /** Fallback load balancer. Handles fallback policy changes. */
 final class FallbackLb extends ForwardingLoadBalancer {
 
   private final Helper fallbackLbHelper;
-  private final LoadBalancerRegistry lbRegistry;
   private final GracefulSwitchLoadBalancer fallbackPolicyLb;
 
   FallbackLb(Helper fallbackLbHelper) {
-    this(checkNotNull(fallbackLbHelper, "fallbackLbHelper"),
-        LoadBalancerRegistry.getDefaultRegistry());
-  }
-
-  @VisibleForTesting
-  FallbackLb(Helper fallbackLbHelper, LoadBalancerRegistry lbRegistry) {
-    this.fallbackLbHelper = fallbackLbHelper;
-    this.lbRegistry = lbRegistry;
+    this.fallbackLbHelper = checkNotNull(fallbackLbHelper, "fallbackLbHelper");
     fallbackPolicyLb = new GracefulSwitchLoadBalancer(fallbackLbHelper);
   }
 
@@ -58,55 +45,18 @@ protected LoadBalancer delegate() {
     return fallbackPolicyLb;
   }
 
-  @SuppressWarnings("deprecation")
   @Override
   public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
-    Attributes attributes = resolvedAddresses.getAttributes();
-    XdsConfig xdsConfig;
-    Object lbConfig = resolvedAddresses.getLoadBalancingPolicyConfig();
-    if (lbConfig != null) {
-      if (!(lbConfig instanceof XdsConfig)) {
-        fallbackLbHelper.updateBalancingState(
-            TRANSIENT_FAILURE,
-            new ErrorPicker(Status.UNAVAILABLE.withDescription(
-                "Load balancing config '" + lbConfig + "' is not an XdsConfig")));
-        return;
-      }
-      xdsConfig = (XdsConfig) lbConfig;
-    } else {
-      // In the future, in all cases xdsConfig can be obtained directly by
-      // resolvedAddresses.getLoadBalancingPolicyConfig().
-      Map<String, ?> newRawLbConfig = attributes.get(ATTR_LOAD_BALANCING_CONFIG);
-      if (newRawLbConfig == null) {
-        // This will not happen when the service config error handling is implemented.
-        // For now simply go to TRANSIENT_FAILURE.
-        fallbackLbHelper.updateBalancingState(
-            TRANSIENT_FAILURE,
-            new ErrorPicker(
-                Status.UNAVAILABLE.withDescription("ATTR_LOAD_BALANCING_CONFIG not available")));
-        return;
-      }
-      ConfigOrError cfg =
-          XdsLoadBalancerProvider.parseLoadBalancingConfigPolicy(newRawLbConfig, lbRegistry);
-      if (cfg.getError() != null) {
-        // This will not happen when the service config error handling is implemented.
-        // For now simply go to TRANSIENT_FAILURE.
-        fallbackLbHelper.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(cfg.getError()));
-        return;
-      }
-      xdsConfig = (XdsConfig) cfg.getConfig();
-    }
-
-    LbConfig fallbackPolicy = xdsConfig.fallbackPolicy;
+    XdsConfig xdsConfig = (XdsConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
+    PolicySelection fallbackPolicy = xdsConfig.fallbackPolicy;
     if (fallbackPolicy == null) {
       // In the latest xDS design, fallback is not supported.
       fallbackLbHelper.updateBalancingState(
           TRANSIENT_FAILURE,
           new ErrorPicker(Status.UNAVAILABLE.withDescription("Fallback is not supported")));
       return;
     }
-    String newFallbackPolicyName = fallbackPolicy.getPolicyName();
-    fallbackPolicyLb.switchTo(lbRegistry.getProvider(newFallbackPolicyName));
+    fallbackPolicyLb.switchTo(fallbackPolicy.getProvider());
 
     List<EquivalentAddressGroup> servers = resolvedAddresses.getAddresses();
     // TODO(zhangkun83): FIXME(#5496): this is a temporary hack.
@@ -119,8 +69,7 @@ public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       // TODO(carl-mastrangelo): propagate the load balancing config policy
       ResolvedAddresses fallbackResolvedAddresses = resolvedAddresses.toBuilder()
           .setAddresses(servers)
-          .setAttributes(attributes.toBuilder()
-              .set(ATTR_LOAD_BALANCING_CONFIG, fallbackPolicy.getRawConfigValue()).build())
+          .setLoadBalancingPolicyConfig(fallbackPolicy.getConfig())
           .build();
       fallbackPolicyLb.handleResolvedAddresses(fallbackResolvedAddresses);
     }

xds/src/main/java/io/grpc/xds/XdsLoadBalancerProvider.java

@@ -16,7 +16,8 @@
 
 package io.grpc.xds;
 
-import com.google.common.annotations.VisibleForTesting;
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Objects;
 import com.google.common.collect.ImmutableMap;
@@ -30,6 +31,8 @@
 import io.grpc.internal.JsonUtil;
 import io.grpc.internal.ServiceConfigUtil;
 import io.grpc.internal.ServiceConfigUtil.LbConfig;
+import io.grpc.internal.ServiceConfigUtil.PolicySelection;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -42,10 +45,7 @@
 @Internal
 public final class XdsLoadBalancerProvider extends LoadBalancerProvider {
 
-  static final String XDS_POLICY_NAME = "xds_experimental";
-
-  private static final LbConfig DEFAULT_FALLBACK_POLICY =
-      new LbConfig("round_robin", ImmutableMap.<String, Void>of());
+  private static final String XDS_POLICY_NAME = "xds_experimental";
 
   @Override
   public boolean isAvailable() {
@@ -78,8 +78,38 @@ static ConfigOrError parseLoadBalancingConfigPolicy(
       Map<String, ?> rawLoadBalancingPolicyConfig, LoadBalancerRegistry registry) {
     try {
       String cluster = JsonUtil.getString(rawLoadBalancingPolicyConfig, "cluster");
-      LbConfig childPolicy = selectChildPolicy(rawLoadBalancingPolicyConfig, registry);
-      LbConfig fallbackPolicy = selectFallbackPolicy(rawLoadBalancingPolicyConfig, registry);
+
+      LbConfig roundRobinConfig = new LbConfig("round_robin", ImmutableMap.<String, Object>of());
+      List<LbConfig> endpointPickingConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
+          JsonUtil.getListOfObjects(rawLoadBalancingPolicyConfig, "endpointPickingPolicy"));
+      if (endpointPickingConfigs == null) {
+        endpointPickingConfigs = new ArrayList<>(1);
+      } else {
+        endpointPickingConfigs = new ArrayList<>(endpointPickingConfigs);
+      }
+      endpointPickingConfigs.add(roundRobinConfig);
+      ConfigOrError childConfigOrError =
+          ServiceConfigUtil.selectLbPolicyFromList(endpointPickingConfigs, registry);
+      if (childConfigOrError.getError() != null) {
+        return childConfigOrError;
+      }
+      PolicySelection childPolicy = (PolicySelection) childConfigOrError.getConfig();
+
+      List<LbConfig> fallbackConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
+          JsonUtil.getListOfObjects(rawLoadBalancingPolicyConfig, "fallbackPolicy"));
+      if (fallbackConfigs == null) {
+        fallbackConfigs = new ArrayList<>(1);
+      } else {
+        fallbackConfigs = new ArrayList<>(fallbackConfigs);
+      }
+      fallbackConfigs.add(roundRobinConfig);
+      ConfigOrError fallbackConfigOrError =
+          ServiceConfigUtil.selectLbPolicyFromList(fallbackConfigs, registry);
+      if (fallbackConfigOrError.getError() != null) {
+        return fallbackConfigOrError;
+      }
+      PolicySelection fallbackPolicy = (PolicySelection) fallbackConfigOrError.getConfig();
+
       String edsServiceName = JsonUtil.getString(rawLoadBalancingPolicyConfig, "edsServiceName");
       String lrsServerName =
           JsonUtil.getString(rawLoadBalancingPolicyConfig, "lrsLoadReportingServerName");
@@ -92,51 +122,16 @@ static ConfigOrError parseLoadBalancingConfigPolicy(
     }
   }
 
-  @VisibleForTesting
-  static LbConfig selectFallbackPolicy(
-      Map<String, ?> rawLoadBalancingPolicyConfig, LoadBalancerRegistry lbRegistry) {
-    List<LbConfig> fallbackConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
-        JsonUtil.getListOfObjects(rawLoadBalancingPolicyConfig, "fallbackPolicy"));
-    LbConfig fallbackPolicy = selectSupportedLbPolicy(fallbackConfigs, lbRegistry);
-    return fallbackPolicy == null ? DEFAULT_FALLBACK_POLICY : fallbackPolicy;
-  }
-
-  @Nullable
-  @VisibleForTesting
-  static LbConfig selectChildPolicy(
-      Map<String, ?> rawLoadBalancingPolicyConfig, LoadBalancerRegistry lbRegistry) {
-    List<LbConfig> childConfigs = ServiceConfigUtil.unwrapLoadBalancingConfigList(
-        JsonUtil.getListOfObjects(rawLoadBalancingPolicyConfig, "childPolicy"));
-    return selectSupportedLbPolicy(childConfigs, lbRegistry);
-  }
-
-  @Nullable
-  private static LbConfig selectSupportedLbPolicy(
-      @Nullable List<LbConfig> lbConfigs, LoadBalancerRegistry lbRegistry) {
-    if (lbConfigs == null) {
-      return null;
-    }
-    for (LbConfig lbConfig : lbConfigs) {
-      String lbPolicy = lbConfig.getPolicyName();
-      if (lbRegistry.getProvider(lbPolicy) != null) {
-        return lbConfig;
-      }
-    }
-    return null;
-  }
-
   /**
    * Represents a successfully parsed and validated LoadBalancingConfig for XDS.
    */
   static final class XdsConfig {
     // FIXME(chengyuanzhang): make cluster name required.
     @Nullable
     final String cluster;
-    // TODO(carl-mastrangelo): make these Object's containing the fully parsed child configs.
-    @Nullable
-    final LbConfig childPolicy;
+    final PolicySelection endpointPickingPolicy;
     @Nullable
-    final LbConfig fallbackPolicy;
+    final PolicySelection fallbackPolicy;
     // Optional. Name to use in EDS query. If not present, defaults to the server name from the
     // target URI.
     @Nullable
@@ -149,12 +144,12 @@ static final class XdsConfig {
 
     XdsConfig(
         @Nullable String cluster,
-        @Nullable LbConfig childPolicy,
-        @Nullable LbConfig fallbackPolicy,
+        PolicySelection endpointPickingPolicy,
+        @Nullable PolicySelection fallbackPolicy,
         @Nullable String edsServiceName,
         @Nullable String lrsServerName) {
       this.cluster = cluster;
-      this.childPolicy = childPolicy;
+      this.endpointPickingPolicy = checkNotNull(endpointPickingPolicy, "endpointPickingPolicy");
       this.fallbackPolicy = fallbackPolicy;
       this.edsServiceName = edsServiceName;
       this.lrsServerName = lrsServerName;
@@ -164,7 +159,7 @@ static final class XdsConfig {
     public String toString() {
       return MoreObjects.toStringHelper(this)
           .add("cluster", cluster)
-          .add("childPolicy", childPolicy)
+          .add("endpointPickingPolicy", endpointPickingPolicy)
           .add("fallbackPolicy", fallbackPolicy)
           .add("edsServiceName", edsServiceName)
           .add("lrsServerName", lrsServerName)
@@ -178,15 +173,16 @@ public boolean equals(Object obj) {
       }
       XdsConfig that = (XdsConfig) obj;
       return Objects.equal(this.cluster, that.cluster)
-          && Objects.equal(this.childPolicy, that.childPolicy)
+          && Objects.equal(this.endpointPickingPolicy, that.endpointPickingPolicy)
           && Objects.equal(this.fallbackPolicy, that.fallbackPolicy)
           && Objects.equal(this.edsServiceName, that.edsServiceName)
           && Objects.equal(this.lrsServerName, that.lrsServerName);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hashCode(cluster, childPolicy, fallbackPolicy, edsServiceName, lrsServerName);
+      return Objects.hashCode(
+          cluster, endpointPickingPolicy, fallbackPolicy, edsServiceName, lrsServerName);
     }
   }
 }