core: Convert AutoConfiguredLB to an actual LB

ejona86 · web-flow · commit 4bbf8eee5069 · 2026-01-05T21:30:53.000-08:00
AutoConfiguredLB wasn't able to be a LB because it needed to be able to reject configuration to cause the NameResolver to refresh. Since 4b4cb0b, and especially 9888a54, the LB API now is able to do this directly. The real end-goal of this work is to replace (much of) AutoConfiguredLB with GracefulSwitchLB. The AutoConfiguredLBFactory will still be needed for config handling, but the LB itself could just become an instance of GracefulSwitchLB. Using GracefulSwitchLB will let us reuse more of the config parsing logic, avoids a latency hit when the top-level policy changes, and gets rid of the last usage of ServiceConfigUtil.selectLbPolicyFromList() outside of GracefulSwitchLB. Go and C are already using GracefulSwitchLB for the top-level policy. Moving the defaultProvider creation earlier was to allow parseLoadBalancingPolicyConfig() to never return null. However, that ran into some simple but annoying test failures because the service config was now being detected as changed. That's solveable, but turns out to be more involved than this change itself, so that's left for later. Since the error handling is nicer now and the earlier creation will be needed eventually anyway, I left the earlier creation in-place even though it technically doesn't have to be done as part of this commit.
diff --git a/core/src/main/java/io/grpc/internal/AutoConfiguredLoadBalancerFactory.java b/core/src/main/java/io/grpc/internal/AutoConfiguredLoadBalancerFactory.java
@@ -38,10 +38,10 @@
 import java.util.Map;
 import javax.annotation.Nullable;
 
-public final class AutoConfiguredLoadBalancerFactory {
+public final class AutoConfiguredLoadBalancerFactory extends LoadBalancerProvider {
 
   private final LoadBalancerRegistry registry;
-  private final String defaultPolicy;
+  private final LoadBalancerProvider defaultProvider;
 
   public AutoConfiguredLoadBalancerFactory(String defaultPolicy) {
     this(LoadBalancerRegistry.getDefaultRegistry(), defaultPolicy);
@@ -50,71 +50,47 @@ public AutoConfiguredLoadBalancerFactory(String defaultPolicy) {
   @VisibleForTesting
   AutoConfiguredLoadBalancerFactory(LoadBalancerRegistry registry, String defaultPolicy) {
     this.registry = checkNotNull(registry, "registry");
-    this.defaultPolicy = checkNotNull(defaultPolicy, "defaultPolicy");
+    LoadBalancerProvider provider =
+        registry.getProvider(checkNotNull(defaultPolicy, "defaultPolicy"));
+    if (provider == null) {
+      Status status = Status.INTERNAL.withDescription("Could not find policy '" + defaultPolicy
+          + "'. Make sure its implementation is either registered to LoadBalancerRegistry or"
+          + " included in META-INF/services/io.grpc.LoadBalancerProvider from your jar files.");
+      provider = new FixedPickerLoadBalancerProvider(
+          ConnectivityState.TRANSIENT_FAILURE,
+          new LoadBalancer.FixedResultPicker(PickResult.withError(status)),
+          status);
+    }
+    this.defaultProvider = provider;
   }
 
+  @Override
   public AutoConfiguredLoadBalancer newLoadBalancer(Helper helper) {
     return new AutoConfiguredLoadBalancer(helper);
   }
 
-  private static final class NoopLoadBalancer extends LoadBalancer {
-
-    @Override
-    @Deprecated
-    @SuppressWarnings("InlineMeSuggester")
-    public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
-    }
-
-    @Override
-    public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
-      return Status.OK;
-    }
-
-    @Override
-    public void handleNameResolutionError(Status error) {}
-
-    @Override
-    public void shutdown() {}
-  }
-
   @VisibleForTesting
-  public final class AutoConfiguredLoadBalancer {
+  public final class AutoConfiguredLoadBalancer extends LoadBalancer {
     private final Helper helper;
     private LoadBalancer delegate;
     private LoadBalancerProvider delegateProvider;
 
     AutoConfiguredLoadBalancer(Helper helper) {
       this.helper = helper;
-      delegateProvider = registry.getProvider(defaultPolicy);
-      if (delegateProvider == null) {
-        throw new IllegalStateException("Could not find policy '" + defaultPolicy
-            + "'. Make sure its implementation is either registered to LoadBalancerRegistry or"
-            + " included in META-INF/services/io.grpc.LoadBalancerProvider from your jar files.");
-      }
+      this.delegateProvider = defaultProvider;
       delegate = delegateProvider.newLoadBalancer(helper);
     }
 
     /**
      * Returns non-OK status if the delegate rejects the resolvedAddresses (e.g. if it does not
      * support an empty list).
      */
-    Status tryAcceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+    @Override
+    public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       PolicySelection policySelection =
           (PolicySelection) resolvedAddresses.getLoadBalancingPolicyConfig();
 
       if (policySelection == null) {
-        LoadBalancerProvider defaultProvider;
-        try {
-          defaultProvider = getProviderOrThrow(defaultPolicy, "using default policy");
-        } catch (PolicyException e) {
-          Status s = Status.INTERNAL.withDescription(e.getMessage());
-          helper.updateBalancingState(
-              ConnectivityState.TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(s)));
-          delegate.shutdown();
-          delegateProvider = null;
-          delegate = new NoopLoadBalancer();
-          return Status.OK;
-        }
         policySelection =
             new PolicySelection(defaultProvider, /* config= */ null);
       }
@@ -145,20 +121,24 @@ Status tryAcceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
               .build());
     }
 
-    void handleNameResolutionError(Status error) {
+    @Override
+    public void handleNameResolutionError(Status error) {
       getDelegate().handleNameResolutionError(error);
     }
 
+    @Override
     @Deprecated
-    void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
+    public void handleSubchannelState(Subchannel subchannel, ConnectivityStateInfo stateInfo) {
       getDelegate().handleSubchannelState(subchannel, stateInfo);
     }
 
-    void requestConnection() {
+    @Override
+    public void requestConnection() {
       getDelegate().requestConnection();
     }
 
-    void shutdown() {
+    @Override
+    public void shutdown() {
       delegate.shutdown();
       delegate = null;
     }
@@ -179,16 +159,6 @@ LoadBalancerProvider getDelegateProvider() {
     }
   }
 
-  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;
-  }
-
   /**
    * Parses first available LoadBalancer policy from service config. Available LoadBalancer should
    * be registered to {@link LoadBalancerRegistry}. If the first available LoadBalancer policy is
@@ -209,8 +179,11 @@ private LoadBalancerProvider getProviderOrThrow(String policy, String choiceReas
    *
    * @return the parsed {@link PolicySelection}, or {@code null} if no selection could be made.
    */
+  // TODO(ejona): The Provider API doesn't allow null, but ScParser can handle this and it will need
+  // tweaking to ManagedChannelImpl.defaultServiceConfig to fix.
   @Nullable
-  ConfigOrError parseLoadBalancerPolicy(Map<String, ?> serviceConfig) {
+  @Override
+  public ConfigOrError parseLoadBalancingPolicyConfig(Map<String, ?> serviceConfig) {
     try {
       List<LbConfig> loadBalancerConfigs = null;
       if (serviceConfig != null) {
@@ -228,12 +201,18 @@ ConfigOrError parseLoadBalancerPolicy(Map<String, ?> serviceConfig) {
     }
   }
 
-  @VisibleForTesting
-  static final class PolicyException extends Exception {
-    private static final long serialVersionUID = 1L;
+  @Override
+  public boolean isAvailable() {
+    return true;
+  }
 
-    private PolicyException(String msg) {
-      super(msg);
-    }
+  @Override
+  public int getPriority() {
+    return 5;
+  }
+
+  @Override
+  public String getPolicyName() {
+    return "auto_configured_internal";
   }
 }
diff --git a/core/src/main/java/io/grpc/internal/FixedPickerLoadBalancerProvider.java b/core/src/main/java/io/grpc/internal/FixedPickerLoadBalancerProvider.java
@@ -0,0 +1,80 @@
+/*
+ * Copyright 2026 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.internal;
+
+import static java.util.Objects.requireNonNull;
+
+import io.grpc.ConnectivityState;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.Status;
+
+/** A LB provider whose LB always uses the same picker. */
+final class FixedPickerLoadBalancerProvider extends LoadBalancerProvider {
+  private final ConnectivityState state;
+  private final LoadBalancer.SubchannelPicker picker;
+  private final Status acceptAddressesStatus;
+
+  public FixedPickerLoadBalancerProvider(
+      ConnectivityState state, LoadBalancer.SubchannelPicker picker, Status acceptAddressesStatus) {
+    this.state = requireNonNull(state, "state");
+    this.picker = requireNonNull(picker, "picker");
+    this.acceptAddressesStatus = requireNonNull(acceptAddressesStatus, "acceptAddressesStatus");
+  }
+
+  @Override
+  public boolean isAvailable() {
+    return true;
+  }
+
+  @Override
+  public int getPriority() {
+    return 5;
+  }
+
+  @Override
+  public String getPolicyName() {
+    return "fixed_picker_lb_internal";
+  }
+
+  @Override
+  public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
+    return new FixedPickerLoadBalancer(helper);
+  }
+
+  private final class FixedPickerLoadBalancer extends LoadBalancer {
+    private final Helper helper;
+
+    public FixedPickerLoadBalancer(Helper helper) {
+      this.helper = requireNonNull(helper, "helper");
+    }
+
+    @Override
+    public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+      helper.updateBalancingState(state, picker);
+      return acceptAddressesStatus;
+    }
+
+    @Override
+    public void handleNameResolutionError(Status error) {
+      helper.updateBalancingState(state, picker);
+    }
+
+    @Override
+    public void shutdown() {}
+  }
+}
diff --git a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
@@ -69,6 +69,7 @@
 import io.grpc.LoadBalancer.ResolvedAddresses;
 import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.LoadBalancer.SubchannelStateListener;
+import io.grpc.LoadBalancerProvider;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.Metadata;
@@ -85,7 +86,6 @@
 import io.grpc.StatusOr;
 import io.grpc.SynchronizationContext;
 import io.grpc.SynchronizationContext.ScheduledHandle;
-import io.grpc.internal.AutoConfiguredLoadBalancerFactory.AutoConfiguredLoadBalancer;
 import io.grpc.internal.ClientCallImpl.ClientStreamProvider;
 import io.grpc.internal.ClientTransportFactory.SwapChannelCredentialsResult;
 import io.grpc.internal.ManagedChannelImplBuilder.ClientTransportFactoryBuilder;
@@ -163,7 +163,7 @@ public Result selectConfig(PickSubchannelArgs args) {
   private final URI targetUri;
   private final NameResolverProvider nameResolverProvider;
   private final NameResolver.Args nameResolverArgs;
-  private final AutoConfiguredLoadBalancerFactory loadBalancerFactory;
+  private final LoadBalancerProvider loadBalancerFactory;
   private final ClientTransportFactory originalTransportFactory;
   @Nullable
   private final ChannelCredentials originalChannelCreds;
@@ -1342,7 +1342,7 @@ void remove(RetriableStream<?> retriableStream) {
   }
 
   private final class LbHelperImpl extends LoadBalancer.Helper {
-    AutoConfiguredLoadBalancer lb;
+    LoadBalancer lb;
 
     @Override
     public AbstractSubchannel createSubchannel(CreateSubchannelArgs args) {
@@ -1727,7 +1727,7 @@ public Status onResult2(final ResolutionResult resolutionResult) {
             .setAddresses(serversOrError.getValue())
             .setAttributes(attributes)
             .setLoadBalancingPolicyConfig(effectiveServiceConfig.getLoadBalancingConfig());
-        Status addressAcceptanceStatus = helper.lb.tryAcceptResolvedAddresses(
+        Status addressAcceptanceStatus = helper.lb.acceptResolvedAddresses(
             resolvedAddresses.build());
         return addressAcceptanceStatus;
       }
diff --git a/core/src/main/java/io/grpc/internal/ScParser.java b/core/src/main/java/io/grpc/internal/ScParser.java
@@ -19,6 +19,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.annotations.VisibleForTesting;
+import io.grpc.LoadBalancerProvider;
 import io.grpc.NameResolver;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
@@ -31,26 +32,28 @@ public final class ScParser extends NameResolver.ServiceConfigParser {
   private final boolean retryEnabled;
   private final int maxRetryAttemptsLimit;
   private final int maxHedgedAttemptsLimit;
-  private final AutoConfiguredLoadBalancerFactory autoLoadBalancerFactory;
+  private final LoadBalancerProvider parser;
 
   /** Creates a parse with global retry settings and an auto configured lb factory. */
   public ScParser(
       boolean retryEnabled,
       int maxRetryAttemptsLimit,
       int maxHedgedAttemptsLimit,
-      AutoConfiguredLoadBalancerFactory autoLoadBalancerFactory) {
+      LoadBalancerProvider parser) {
     this.retryEnabled = retryEnabled;
     this.maxRetryAttemptsLimit = maxRetryAttemptsLimit;
     this.maxHedgedAttemptsLimit = maxHedgedAttemptsLimit;
-    this.autoLoadBalancerFactory = checkNotNull(autoLoadBalancerFactory, "autoLoadBalancerFactory");
+    this.parser = checkNotNull(parser, "parser");
   }
 
   @Override
   public ConfigOrError parseServiceConfig(Map<String, ?> rawServiceConfig) {
     try {
       Object loadBalancingPolicySelection;
       ConfigOrError choiceFromLoadBalancer =
-          autoLoadBalancerFactory.parseLoadBalancerPolicy(rawServiceConfig);
+          parser.parseLoadBalancingPolicyConfig(rawServiceConfig);
+      // TODO(ejona): The Provider API doesn't allow null, but AutoConfiguredLoadBalancerFactory can
+      // return null and it will need tweaking to ManagedChannelImpl.defaultServiceConfig to fix.
       if (choiceFromLoadBalancer == null) {
         loadBalancingPolicySelection = null;
       } else if (choiceFromLoadBalancer.getError() != null) {
diff --git a/core/src/test/java/io/grpc/internal/AutoConfiguredLoadBalancerFactoryTest.java b/core/src/test/java/io/grpc/internal/AutoConfiguredLoadBalancerFactoryTest.java
diff --git a/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java b/core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
