util: Add GracefulSwitchLb config

This is to replace switchTo(), to allow composing GracefulSwitchLb with
other helpers like MultiChildLb. It also prevents users of
GracefulSwitchLb from needing to use ServiceConfigUtil.

Author: ejona86

util/src/main/java/io/grpc/util/GracefulSwitchLoadBalancer.java

@@ -20,11 +20,18 @@
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Objects;
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.ExperimentalApi;
 import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerRegistry;
+import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
+import io.grpc.internal.ServiceConfigUtil;
+import java.util.List;
+import java.util.Map;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 
@@ -33,9 +40,16 @@
  * other than READY, the new policy will be swapped into place immediately.  Otherwise, the channel
  * will keep using the old policy until the new policy reports READY or the old policy exits READY.
  *
- * <p>The balancer must {@link #switchTo(LoadBalancer.Factory) switch to} a policy prior to {@link
+ * <p>The child balancer and configuration is specified using service config. Config objects are
+ * generally created by calling {@link #parseLoadBalancingPolicyConfig(List)} from a
+ * {@link io.grpc.LoadBalancerProvider#parseLoadBalancingPolicyConfig
+ * provider's parseLoadBalancingPolicyConfig()} implementation.
+ *
+ * <p>Alternatively, the balancer may {@link #switchTo(LoadBalancer.Factory) switch to} a policy
+ * prior to {@link
  * LoadBalancer#handleResolvedAddresses(ResolvedAddresses) handling resolved addresses} for the
- * first time.
+ * first time. This causes graceful switch to ignore the service config and pass through the
+ * resolved addresses directly to the child policy.
  */
 @ExperimentalApi("https://github.com/grpc/grpc-java/issues/5999")
 @NotThreadSafe // Must be accessed in SynchronizationContext
@@ -90,18 +104,51 @@ public String toString() {
   private LoadBalancer pendingLb = defaultBalancer;
   private ConnectivityState pendingState;
   private SubchannelPicker pendingPicker;
+  private boolean switchToCalled;
 
   private boolean currentLbIsReady;
 
   public GracefulSwitchLoadBalancer(Helper helper) {
     this.helper = checkNotNull(helper, "helper");
   }
 
+  @Override
+  public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+    if (switchToCalled) {
+      delegate().handleResolvedAddresses(resolvedAddresses);
+      return;
+    }
+    Config config = (Config) resolvedAddresses.getLoadBalancingPolicyConfig();
+    switchToInternal(config.childFactory);
+    delegate().handleResolvedAddresses(
+        resolvedAddresses.toBuilder()
+          .setLoadBalancingPolicyConfig(config.childConfig)
+          .build());
+  }
+
+  @Override
+  public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+    if (switchToCalled) {
+      return delegate().acceptResolvedAddresses(resolvedAddresses);
+    }
+    Config config = (Config) resolvedAddresses.getLoadBalancingPolicyConfig();
+    switchToInternal(config.childFactory);
+    return delegate().acceptResolvedAddresses(
+        resolvedAddresses.toBuilder()
+          .setLoadBalancingPolicyConfig(config.childConfig)
+          .build());
+  }
+
   /**
    * Gracefully switch to a new policy defined by the given factory, if the given factory isn't
    * equal to the current one.
    */
   public void switchTo(LoadBalancer.Factory newBalancerFactory) {
+    switchToCalled = true;
+    switchToInternal(newBalancerFactory);
+  }
+
+  private void switchToInternal(LoadBalancer.Factory newBalancerFactory) {
     checkNotNull(newBalancerFactory, "newBalancerFactory");
 
     if (newBalancerFactory.equals(pendingBalancerFactory)) {
@@ -185,4 +232,86 @@ public void shutdown() {
   public String delegateType() {
     return delegate().getClass().getSimpleName();
   }
+
+  /**
+   * Provided a JSON list of LoadBalancingConfigs, parse it into a config to pass to GracefulSwitch.
+   */
+  public static ConfigOrError parseLoadBalancingPolicyConfig(
+      List<Map<String, ?>> loadBalancingConfigs) {
+    return parseLoadBalancingPolicyConfig(
+        loadBalancingConfigs, LoadBalancerRegistry.getDefaultRegistry());
+  }
+
+  /**
+   * Provided a JSON list of LoadBalancingConfigs, parse it into a config to pass to GracefulSwitch.
+   */
+  public static ConfigOrError parseLoadBalancingPolicyConfig(
+      List<Map<String, ?>> loadBalancingConfigs, LoadBalancerRegistry lbRegistry) {
+    List<ServiceConfigUtil.LbConfig> childConfigCandidates =
+        ServiceConfigUtil.unwrapLoadBalancingConfigList(loadBalancingConfigs);
+    if (childConfigCandidates == null || childConfigCandidates.isEmpty()) {
+      return ConfigOrError.fromError(
+          Status.INTERNAL.withDescription("No child LB config specified"));
+    }
+    ConfigOrError selectedConfig =
+        ServiceConfigUtil.selectLbPolicyFromList(childConfigCandidates, lbRegistry);
+    if (selectedConfig.getError() != null) {
+      Status error = selectedConfig.getError();
+      return ConfigOrError.fromError(
+          Status.INTERNAL
+              .withCause(error.getCause())
+              .withDescription(error.getDescription())
+              .augmentDescription("Failed to select child config"));
+    }
+    ServiceConfigUtil.PolicySelection selection =
+        (ServiceConfigUtil.PolicySelection) selectedConfig.getConfig();
+    return ConfigOrError.fromConfig(
+        createLoadBalancingPolicyConfig(selection.getProvider(), selection.getConfig()));
+  }
+
+  /**
+   * Directly create a config to pass to GracefulSwitch. The object returned is the same as would be
+   * found in {@code ConfigOrError.getConfig()}.
+   */
+  public static Object createLoadBalancingPolicyConfig(
+      LoadBalancer.Factory childFactory, @Nullable Object childConfig) {
+    return new Config(childFactory, childConfig);
+  }
+
+  static final class Config {
+    final LoadBalancer.Factory childFactory;
+    @Nullable
+    final Object childConfig;
+
+    public Config(LoadBalancer.Factory childFactory, @Nullable Object childConfig) {
+      this.childFactory = checkNotNull(childFactory, "childFactory");
+      this.childConfig = childConfig;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof Config)) {
+        return false;
+      }
+      Config that = (Config) o;
+      return Objects.equal(childFactory, that.childFactory)
+          && Objects.equal(childConfig, that.childConfig);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hashCode(childFactory, childConfig);
+    }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper("GracefulSwitchLoadBalancer.Config")
+          .add("childFactory", childFactory)
+          .add("childConfig", childConfig)
+          .toString();
+    }
+  }
 }

util/src/test/java/io/grpc/util/GracefulSwitchLoadBalancerTest.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.util;
+
+import io.grpc.LoadBalancerProvider;
+
+/**
+ * Accessors for white-box testing involving GracefulSwitchLoadBalancer.
+ */
+public final class GracefulSwitchLoadBalancerAccessor {
+  private GracefulSwitchLoadBalancerAccessor() {
+    // Do not instantiate
+  }
+
+  public static LoadBalancerProvider getChildProvider(Object config) {
+    return (LoadBalancerProvider) ((GracefulSwitchLoadBalancer.Config) config).childFactory;
+  }
+
+  public static Object getChildConfig(Object config) {
+    return ((GracefulSwitchLoadBalancer.Config) config).childConfig;
+  }
+}