xds: Avoid switchTo in WrrLocalityLb and WeightedTargetLb

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

@@ -79,17 +79,11 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
     WeightedTargetConfig weightedTargetConfig = (WeightedTargetConfig) lbConfig;
     Map<String, WeightedPolicySelection> newTargets = weightedTargetConfig.targets;
     for (String targetName : newTargets.keySet()) {
-      WeightedPolicySelection weightedChildLbConfig = newTargets.get(targetName);
       if (!targets.containsKey(targetName)) {
         ChildHelper childHelper = new ChildHelper(targetName);
         GracefulSwitchLoadBalancer childBalancer = new GracefulSwitchLoadBalancer(childHelper);
-        childBalancer.switchTo(weightedChildLbConfig.policySelection.getProvider());
         childHelpers.put(targetName, childHelper);
         childBalancers.put(targetName, childBalancer);
-      } else if (!weightedChildLbConfig.policySelection.getProvider().equals(
-          targets.get(targetName).policySelection.getProvider())) {
-        childBalancers.get(targetName)
-            .switchTo(weightedChildLbConfig.policySelection.getProvider());
       }
     }
     targets = newTargets;
@@ -97,7 +91,7 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
       childBalancers.get(targetName).handleResolvedAddresses(
           resolvedAddresses.toBuilder()
               .setAddresses(AddressFilter.filter(resolvedAddresses.getAddresses(), targetName))
-              .setLoadBalancingPolicyConfig(targets.get(targetName).policySelection.getConfig())
+              .setLoadBalancingPolicyConfig(targets.get(targetName).childConfig)
               .setAttributes(resolvedAddresses.getAttributes().toBuilder()
                 .set(CHILD_NAME, targetName)
                 .build())

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

@@ -26,11 +26,8 @@ import io.grpc.LoadBalancerRegistry;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
 import io.grpc.internal.JsonUtil;
-import io.grpc.internal.ServiceConfigUtil;
+import io.grpc.util.GracefulSwitchLoadBalancer;
-import io.grpc.internal.ServiceConfigUtil.LbConfig;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -97,22 +94,16 @@ public final class WeightedTargetLoadBalancerProvider extends LoadBalancerProvid
           return ConfigOrError.fromError(Status.INTERNAL.withDescription(
               "Wrong weight for target " + name + " in weighted_target LB policy:\n " + rawConfig));
         }
-        List<LbConfig> childConfigCandidates = ServiceConfigUtil.unwrapLoadBalancingConfigList(
-            JsonUtil.getListOfObjects(rawWeightedTarget, "childPolicy"));
-        if (childConfigCandidates == null || childConfigCandidates.isEmpty()) {
-          return ConfigOrError.fromError(Status.INTERNAL.withDescription(
-              "No child policy for target " + name + " in weighted_target LB policy:\n "
-                  + rawConfig));
-        }
         LoadBalancerRegistry lbRegistry =
             this.lbRegistry == null ? LoadBalancerRegistry.getDefaultRegistry() : this.lbRegistry;
-        ConfigOrError selectedConfig =
+        ConfigOrError childConfig = GracefulSwitchLoadBalancer.parseLoadBalancingPolicyConfig(
-            ServiceConfigUtil.selectLbPolicyFromList(childConfigCandidates, lbRegistry);
+            JsonUtil.getListOfObjects(rawWeightedTarget, "childPolicy"), lbRegistry);
-        if (selectedConfig.getError() != null) {
+        if (childConfig.getError() != null) {
-          return selectedConfig;
+          return ConfigOrError.fromError(Status.INTERNAL
+              .withDescription("Could not parse weighted_target's child policy:" + name)
+              .withCause(childConfig.getError().asRuntimeException()));
         }
-        PolicySelection policySelection = (PolicySelection) selectedConfig.getConfig();
+        parsedChildConfigs.put(name, new WeightedPolicySelection(weight, childConfig.getConfig()));
-        parsedChildConfigs.put(name, new WeightedPolicySelection(weight, policySelection));
       }
       return ConfigOrError.fromConfig(new WeightedTargetConfig(parsedChildConfigs));
     } catch (RuntimeException e) {
@@ -125,11 +116,11 @@ public final class WeightedTargetLoadBalancerProvider extends LoadBalancerProvid
   static final class WeightedPolicySelection {
 
     final int weight;
-    final PolicySelection policySelection;
+    final Object childConfig;
 
-    WeightedPolicySelection(int weight, PolicySelection policySelection) {
+    WeightedPolicySelection(int weight, Object childConfig) {
       this.weight = weight;
-      this.policySelection = policySelection;
+      this.childConfig = childConfig;
     }
 
     @Override
@@ -141,19 +132,19 @@ public final class WeightedTargetLoadBalancerProvider extends LoadBalancerProvid
         return false;
       }
       WeightedPolicySelection that = (WeightedPolicySelection) o;
-      return weight == that.weight && Objects.equals(policySelection, that.policySelection);
+      return weight == that.weight && Objects.equals(childConfig, that.childConfig);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(weight, policySelection);
+      return Objects.hash(weight, childConfig);
     }
 
     @Override
     public String toString() {
       return MoreObjects.toStringHelper(this)
           .add("weight", weight)
-          .add("policySelection", policySelection)
+          .add("childConfig", childConfig)
           .toString();
     }
   }

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

@@ -27,7 +27,6 @@ import io.grpc.InternalLogId;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.Status;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedPolicySelection;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedTargetConfig;
@@ -108,13 +107,15 @@ final class WrrLocalityLoadBalancer extends LoadBalancer {
     for (String locality : localityWeights.keySet()) {
       weightedPolicySelections.put(locality,
           new WeightedPolicySelection(localityWeights.get(locality),
-              wrrLocalityConfig.childPolicy));
+              wrrLocalityConfig.childConfig));
     }
 
-    switchLb.switchTo(lbRegistry.getProvider(WEIGHTED_TARGET_POLICY_NAME));
+    Object switchConfig = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+        lbRegistry.getProvider(WEIGHTED_TARGET_POLICY_NAME),
+        new WeightedTargetConfig(weightedPolicySelections));
     switchLb.handleResolvedAddresses(
         resolvedAddresses.toBuilder()
-            .setLoadBalancingPolicyConfig(new WeightedTargetConfig(weightedPolicySelections))
+            .setLoadBalancingPolicyConfig(switchConfig)
             .build());
 
     return Status.OK;
@@ -136,10 +137,10 @@ final class WrrLocalityLoadBalancer extends LoadBalancer {
    */
   static final class WrrLocalityConfig {
 
-    final PolicySelection childPolicy;
+    final Object childConfig;
 
-    WrrLocalityConfig(PolicySelection childPolicy) {
+    WrrLocalityConfig(Object childConfig) {
-      this.childPolicy = childPolicy;
+      this.childConfig = childConfig;
     }
 
     @Override
@@ -151,17 +152,17 @@ final class WrrLocalityLoadBalancer extends LoadBalancer {
         return false;
       }
       WrrLocalityConfig that = (WrrLocalityConfig) o;
-      return Objects.equals(childPolicy, that.childPolicy);
+      return Objects.equals(childConfig, that.childConfig);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hashCode(childPolicy);
+      return Objects.hashCode(childConfig);
     }
 
     @Override
     public String toString() {
-      return MoreObjects.toStringHelper(this).add("childPolicy", childPolicy).toString();
+      return MoreObjects.toStringHelper(this).add("childConfig", childConfig).toString();
     }
   }
 }

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

@@ -24,11 +24,8 @@ import io.grpc.LoadBalancerRegistry;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
 import io.grpc.internal.JsonUtil;
-import io.grpc.internal.ServiceConfigUtil;
+import io.grpc.util.GracefulSwitchLoadBalancer;
-import io.grpc.internal.ServiceConfigUtil.LbConfig;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.xds.WrrLocalityLoadBalancer.WrrLocalityConfig;
-import java.util.List;
 import java.util.Map;
 
 /**
@@ -62,21 +59,14 @@ public final class WrrLocalityLoadBalancerProvider extends LoadBalancerProvider
   @Override
   public ConfigOrError parseLoadBalancingPolicyConfig(Map<String, ?> rawConfig) {
     try {
-      List<LbConfig> childConfigCandidates = ServiceConfigUtil.unwrapLoadBalancingConfigList(
+      ConfigOrError childConfig = GracefulSwitchLoadBalancer.parseLoadBalancingPolicyConfig(
           JsonUtil.getListOfObjects(rawConfig, "childPolicy"));
-      if (childConfigCandidates == null || childConfigCandidates.isEmpty()) {
+      if (childConfig.getError() != null) {
-        return ConfigOrError.fromError(Status.INTERNAL.withDescription(
+        return ConfigOrError.fromError(Status.INTERNAL
-            "No child policy in wrr_locality LB policy: "
+            .withDescription("Failed to parse child policy in wrr_locality LB policy: " + rawConfig)
-                + rawConfig));
+            .withCause(childConfig.getError().asRuntimeException()));
       }
-      ConfigOrError selectedConfig =
+      return ConfigOrError.fromConfig(new WrrLocalityConfig(childConfig.getConfig()));
-          ServiceConfigUtil.selectLbPolicyFromList(childConfigCandidates,
-              LoadBalancerRegistry.getDefaultRegistry());
-      if (selectedConfig.getError() != null) {
-        return selectedConfig;
-      }
-      PolicySelection policySelection = (PolicySelection) selectedConfig.getConfig();
-      return ConfigOrError.fromConfig(new WrrLocalityConfig(policySelection));
     } catch (RuntimeException e) {
       return ConfigOrError.fromError(Status.INTERNAL.withCause(e)
           .withDescription("Failed to parse wrr_locality LB config: " + rawConfig));

xds/src/test/java/io/grpc/xds/ClusterImplLoadBalancerTest.java

@@ -52,6 +52,7 @@ import io.grpc.internal.PickSubchannelArgsImpl;
 import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.protobuf.ProtoUtils;
 import io.grpc.testing.TestMethodDescriptors;
+import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.ClusterImplLoadBalancerProvider.ClusterImplConfig;
 import io.grpc.xds.Endpoints.DropOverload;
 import io.grpc.xds.EnvoyServerProtoData.DownstreamTlsContext;
@@ -119,8 +120,8 @@ public class ClusterImplLoadBalancerTest {
   private final FakeClock fakeClock = new FakeClock();
   private final Locality locality =
       Locality.create("test-region", "test-zone", "test-subzone");
-  private final PolicySelection roundRobin =
+  private final Object roundRobin = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
-      new PolicySelection(new FakeLoadBalancerProvider("round_robin"), null);
+      new FakeLoadBalancerProvider("round_robin"), null);
   private final List<FakeLoadBalancer> downstreamBalancers = new ArrayList<>();
   private final FakeTlsContextManager tlsContextManager = new FakeTlsContextManager();
   private final LoadStatsManager2 loadStatsManager =

xds/src/test/java/io/grpc/xds/ClusterResolverLoadBalancerTest.java

@@ -59,6 +59,8 @@ import io.grpc.internal.FakeClock.ScheduledTask;
 import io.grpc.internal.GrpcUtil;
 import io.grpc.internal.ObjectPool;
 import io.grpc.internal.ServiceConfigUtil.PolicySelection;
+import io.grpc.util.GracefulSwitchLoadBalancer;
+import io.grpc.util.GracefulSwitchLoadBalancerAccessor;
 import io.grpc.util.OutlierDetectionLoadBalancer.OutlierDetectionLoadBalancerConfig;
 import io.grpc.util.OutlierDetectionLoadBalancerProvider;
 import io.grpc.xds.ClusterImplLoadBalancerProvider.ClusterImplConfig;
@@ -159,12 +161,14 @@ public class ClusterResolverLoadBalancerTest {
   private final NameResolverRegistry nsRegistry = new NameResolverRegistry();
   private final PolicySelection roundRobin = new PolicySelection(
       new FakeLoadBalancerProvider("wrr_locality_experimental"), new WrrLocalityConfig(
-          new PolicySelection(new FakeLoadBalancerProvider("round_robin"), null)));
+          GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+              new FakeLoadBalancerProvider("round_robin"), null)));
   private final PolicySelection ringHash = new PolicySelection(
       new FakeLoadBalancerProvider("ring_hash_experimental"), new RingHashConfig(10L, 100L));
   private final PolicySelection leastRequest = new PolicySelection(
       new FakeLoadBalancerProvider("wrr_locality_experimental"), new WrrLocalityConfig(
-          new PolicySelection(new FakeLoadBalancerProvider("least_request_experimental"),
+          GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+              new FakeLoadBalancerProvider("least_request_experimental"),
               new LeastRequestConfig(3))));
   private final List<FakeLoadBalancer> childBalancers = new ArrayList<>();
   private final List<FakeNameResolver> resolvers = new ArrayList<>();
@@ -331,8 +335,9 @@ public class ClusterResolverLoadBalancerTest {
         tlsContext, Collections.<DropOverload>emptyList(), WRR_LOCALITY_POLICY_NAME);
     WrrLocalityConfig wrrLocalityConfig =
         (WrrLocalityConfig) clusterImplConfig.childPolicy.getConfig();
-    assertThat(wrrLocalityConfig.childPolicy.getProvider().getPolicyName()).isEqualTo(
+    LoadBalancerProvider childProvider =
-        "least_request_experimental");
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(wrrLocalityConfig.childConfig);
+    assertThat(childProvider.getPolicyName()).isEqualTo("least_request_experimental");
 
     assertThat(
         childBalancer.addresses.get(0).getAttributes()
@@ -414,8 +419,9 @@ public class ClusterResolverLoadBalancerTest {
         tlsContext, Collections.<DropOverload>emptyList(), WRR_LOCALITY_POLICY_NAME);
     WrrLocalityConfig wrrLocalityConfig =
         (WrrLocalityConfig) clusterImplConfig.childPolicy.getConfig();
-    assertThat(wrrLocalityConfig.childPolicy.getProvider().getPolicyName()).isEqualTo(
+    LoadBalancerProvider childProvider =
-        "least_request_experimental");
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(wrrLocalityConfig.childConfig);
+    assertThat(childProvider.getPolicyName()).isEqualTo("least_request_experimental");
 
     assertThat(
         childBalancer.addresses.get(0).getAttributes()
@@ -484,8 +490,9 @@ public class ClusterResolverLoadBalancerTest {
     assertThat(clusterImplConfig1.childPolicy.getConfig()).isInstanceOf(WrrLocalityConfig.class);
     WrrLocalityConfig wrrLocalityConfig1 =
         (WrrLocalityConfig) clusterImplConfig1.childPolicy.getConfig();
-    assertThat(wrrLocalityConfig1.childPolicy.getProvider().getPolicyName()).isEqualTo(
+    LoadBalancerProvider childProvider1 =
-        "round_robin");
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(wrrLocalityConfig1.childConfig);
+    assertThat(childProvider1.getPolicyName()).isEqualTo("round_robin");
 
     PriorityChildConfig priorityChildConfig2 = priorityLbConfig.childConfigs.get(priority2);
     assertThat(priorityChildConfig2.ignoreReresolution).isTrue();
@@ -498,8 +505,9 @@ public class ClusterResolverLoadBalancerTest {
     assertThat(clusterImplConfig2.childPolicy.getConfig()).isInstanceOf(WrrLocalityConfig.class);
     WrrLocalityConfig wrrLocalityConfig2 =
         (WrrLocalityConfig) clusterImplConfig1.childPolicy.getConfig();
-    assertThat(wrrLocalityConfig2.childPolicy.getProvider().getPolicyName()).isEqualTo(
+    LoadBalancerProvider childProvider2 =
-        "round_robin");
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(wrrLocalityConfig2.childConfig);
+    assertThat(childProvider2.getPolicyName()).isEqualTo("round_robin");
 
     PriorityChildConfig priorityChildConfig3 = priorityLbConfig.childConfigs.get(priority3);
     assertThat(priorityChildConfig3.ignoreReresolution).isTrue();
@@ -512,8 +520,9 @@ public class ClusterResolverLoadBalancerTest {
     assertThat(clusterImplConfig3.childPolicy.getConfig()).isInstanceOf(WrrLocalityConfig.class);
     WrrLocalityConfig wrrLocalityConfig3 =
         (WrrLocalityConfig) clusterImplConfig1.childPolicy.getConfig();
-    assertThat(wrrLocalityConfig3.childPolicy.getProvider().getPolicyName()).isEqualTo(
+    LoadBalancerProvider childProvider3 =
-        "round_robin");
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(wrrLocalityConfig3.childConfig);
+    assertThat(childProvider3.getPolicyName()).isEqualTo("round_robin");
 
     for (EquivalentAddressGroup eag : childBalancer.addresses) {
       if (eag.getAttributes().get(InternalXdsAttributes.ATTR_LOCALITY) == locality1) {

xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerProviderTest.java

@@ -26,7 +26,7 @@ import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.internal.JsonParser;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
+import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedPolicySelection;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedTargetConfig;
 import java.util.Map;
@@ -128,11 +128,13 @@ public class WeightedTargetLoadBalancerProviderTest {
             "target_1",
             new WeightedPolicySelection(
                 10,
-                new PolicySelection(lbProviderFoo, fooConfig)),
+                GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+                    lbProviderFoo, fooConfig)),
             "target_2",
             new WeightedPolicySelection(
                 20,
-                new PolicySelection(lbProviderBar, barConfig)))));
+                GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+                    lbProviderBar, barConfig)))));
     assertThat(parsedConfig).isEqualTo(expectedConfig);
   }
 }

xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java

@@ -49,7 +49,7 @@ import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
+import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.WeightedRandomPicker.WeightedChildPicker;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedPolicySelection;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedTargetConfig;
@@ -146,13 +146,13 @@ public class WeightedTargetLoadBalancerTest {
   };
 
   private final WeightedPolicySelection weightedLbConfig0 = new WeightedPolicySelection(
-      weights[0], new PolicySelection(fooLbProvider, configs[0]));
+      weights[0], newChildConfig(fooLbProvider, configs[0]));
   private final WeightedPolicySelection weightedLbConfig1 = new WeightedPolicySelection(
-      weights[1], new PolicySelection(barLbProvider, configs[1]));
+      weights[1], newChildConfig(barLbProvider, configs[1]));
   private final WeightedPolicySelection weightedLbConfig2 = new WeightedPolicySelection(
-      weights[2],  new PolicySelection(barLbProvider, configs[2]));
+      weights[2],  newChildConfig(barLbProvider, configs[2]));
   private final WeightedPolicySelection weightedLbConfig3 = new WeightedPolicySelection(
-      weights[3], new PolicySelection(fooLbProvider, configs[3]));
+      weights[3], newChildConfig(fooLbProvider, configs[3]));
 
   @Mock
   private Helper helper;
@@ -233,16 +233,16 @@ public class WeightedTargetLoadBalancerTest {
     Map<String, WeightedPolicySelection> newTargets = ImmutableMap.of(
         "target1",
         new WeightedPolicySelection(
-            newWeights[0], new PolicySelection(barLbProvider, newConfigs[0])),
+            newWeights[0], newChildConfig(barLbProvider, newConfigs[0])),
         "target2",
         new WeightedPolicySelection(
-            newWeights[1], new PolicySelection(barLbProvider, newConfigs[1])),
+            newWeights[1], newChildConfig(barLbProvider, newConfigs[1])),
         "target3",
         new WeightedPolicySelection(
-            newWeights[2], new PolicySelection(fooLbProvider, newConfigs[2])),
+            newWeights[2], newChildConfig(fooLbProvider, newConfigs[2])),
         "target4",
         new WeightedPolicySelection(
-            newWeights[3], new PolicySelection(fooLbProvider, newConfigs[3])));
+            newWeights[3], newChildConfig(fooLbProvider, newConfigs[3])));
     weightedTargetLb.handleResolvedAddresses(
         ResolvedAddresses.newBuilder()
             .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
@@ -418,9 +418,9 @@ public class WeightedTargetLoadBalancerTest {
 
     Map<String, WeightedPolicySelection> targets = ImmutableMap.of(
         "target0", new WeightedPolicySelection(
-            weights[0], new PolicySelection(fakeLbProvider, configs[0])),
+            weights[0], newChildConfig(fakeLbProvider, configs[0])),
         "target3", new WeightedPolicySelection(
-            weights[3], new PolicySelection(fakeLbProvider, configs[3])));
+            weights[3], newChildConfig(fakeLbProvider, configs[3])));
     weightedTargetLb.handleResolvedAddresses(
         ResolvedAddresses.newBuilder()
             .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
@@ -434,6 +434,10 @@ public class WeightedTargetLoadBalancerTest {
 
   }
 
+  private Object newChildConfig(LoadBalancerProvider provider, Object config) {
+    return GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(provider, config);
+  }
+
   private static class FakeLoadBalancerProvider extends LoadBalancerProvider {
 
     @Override

xds/src/test/java/io/grpc/xds/WrrLocalityLoadBalancerProviderTest.java

@@ -27,6 +27,7 @@ import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.NameResolver;
+import io.grpc.util.GracefulSwitchLoadBalancerAccessor;
 import io.grpc.xds.WrrLocalityLoadBalancer.WrrLocalityConfig;
 import java.util.Map;
 import org.junit.Test;
@@ -64,6 +65,8 @@ public class WrrLocalityLoadBalancerProviderTest {
     WrrLocalityLoadBalancerProvider provider = new WrrLocalityLoadBalancerProvider();
     NameResolver.ConfigOrError configOrError = provider.parseLoadBalancingPolicyConfig(rawConfig);
     WrrLocalityConfig config = (WrrLocalityConfig) configOrError.getConfig();
-    assertThat(config.childPolicy.getProvider().getPolicyName()).isEqualTo("round_robin");
+    LoadBalancerProvider childProvider =
+        GracefulSwitchLoadBalancerAccessor.getChildProvider(config.childConfig);
+    assertThat(childProvider.getPolicyName()).isEqualTo("round_robin");
   }
 }

xds/src/test/java/io/grpc/xds/WrrLocalityLoadBalancerTest.java

@@ -38,7 +38,7 @@ import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext;
-import io.grpc.internal.ServiceConfigUtil.PolicySelection;
+import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedPolicySelection;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedTargetConfig;
 import io.grpc.xds.WrrLocalityLoadBalancer.WrrLocalityConfig;
@@ -112,7 +112,7 @@ public class WrrLocalityLoadBalancerTest {
     // A two locality cluster with a mock child LB policy.
     String localityOne = "localityOne";
     String localityTwo = "localityTwo";
-    PolicySelection childPolicy = new PolicySelection(mockChildProvider, null);
+    Object childPolicy = newChildConfig(mockChildProvider, null);
 
     // The child config is delivered wrapped in the wrr_locality config and the locality weights
     // in a ResolvedAddresses attribute.
@@ -138,7 +138,7 @@ public class WrrLocalityLoadBalancerTest {
   @Test
   public void handleResolvedAddresses_noLocalityWeights() {
     // A two locality cluster with a mock child LB policy.
-    PolicySelection childPolicy = new PolicySelection(mockChildProvider, null);
+    Object childPolicy = newChildConfig(mockChildProvider, null);
 
     // The child config is delivered wrapped in the wrr_locality config and the locality weights
     // in a ResolvedAddresses attribute.
@@ -163,7 +163,7 @@ public class WrrLocalityLoadBalancerTest {
 
   @Test
   public void handleNameResolutionError_withChildLb() {
-    deliverAddresses(new WrrLocalityConfig(new PolicySelection(mockChildProvider, null)),
+    deliverAddresses(new WrrLocalityConfig(newChildConfig(mockChildProvider, null)),
         ImmutableList.of(makeAddress("addr1", "test-locality", 1)));
     Status status = Status.DEADLINE_EXCEEDED.withDescription("too slow");
     loadBalancer.handleNameResolutionError(status);
@@ -175,7 +175,7 @@ public class WrrLocalityLoadBalancerTest {
 
   @Test
   public void localityWeightAttributeNotPropagated() {
-    PolicySelection childPolicy = new PolicySelection(mockChildProvider, null);
+    Object childPolicy = newChildConfig(mockChildProvider, null);
 
     WrrLocalityConfig wlConfig = new WrrLocalityConfig(childPolicy);
     deliverAddresses(wlConfig, ImmutableList.of(makeAddress("addr1", "test-locality", 1)));
@@ -190,7 +190,7 @@ public class WrrLocalityLoadBalancerTest {
 
   @Test
   public void shutdown() {
-    deliverAddresses(new WrrLocalityConfig(new PolicySelection(mockChildProvider, null)),
+    deliverAddresses(new WrrLocalityConfig(newChildConfig(mockChildProvider, null)),
         ImmutableList.of(makeAddress("addr", "test-locality", 1)));
     loadBalancer.shutdown();
 
@@ -199,17 +199,19 @@ public class WrrLocalityLoadBalancerTest {
 
   @Test
   public void configEquality() {
-    WrrLocalityConfig configOne = new WrrLocalityConfig(
+    WrrLocalityConfig configOne = new WrrLocalityConfig(newChildConfig(mockChildProvider, null));
-        new PolicySelection(mockChildProvider, null));
+    WrrLocalityConfig configTwo = new WrrLocalityConfig(newChildConfig(mockChildProvider, null));
-    WrrLocalityConfig configTwo = new WrrLocalityConfig(
-        new PolicySelection(mockChildProvider, null));
     WrrLocalityConfig differentConfig = new WrrLocalityConfig(
-        new PolicySelection(mockChildProvider, "config"));
+        newChildConfig(mockChildProvider, "config"));
 
     new EqualsTester().addEqualityGroup(configOne, configTwo).addEqualityGroup(differentConfig)
         .testEquals();
   }
 
+  private Object newChildConfig(LoadBalancerProvider provider, Object config) {
+    return GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(provider, config);
+  }
+
   private void deliverAddresses(WrrLocalityConfig config, List<EquivalentAddressGroup> addresses) {
     loadBalancer.handleResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(addresses).setLoadBalancingPolicyConfig(config)