core: Validate policy name in defaultLoadBalancingPolicy() immediately

Author: manmohak07

api/src/main/java/io/grpc/ManagedChannelBuilder.java

@@ -274,6 +274,11 @@ public T useTransportSecurity() {
    * <p>Policy implementations are looked up in the
    * {@link LoadBalancerRegistry#getDefaultRegistry default LoadBalancerRegistry}.
    *
+   * <p>The provided policy name is validated against the
+   * {@link LoadBalancerRegistry#getDefaultRegistry default LoadBalancerRegistry} immediately.
+   * If no provider is found for the given policy name, an {@link IllegalArgumentException}
+   * is thrown.
+   * 
    * <p>This method is implemented by all stock channel builders that are shipped with gRPC, but may
    * not be implemented by custom channel builders, in which case this method will throw.
    *

core/src/main/java/io/grpc/internal/ManagedChannelImplBuilder.java

@@ -39,6 +39,8 @@
 import io.grpc.InternalChannelz;
 import io.grpc.InternalConfiguratorRegistry;
 import io.grpc.InternalFeatureFlags;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.LoadBalancerRegistry;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
@@ -447,7 +449,9 @@ public ManagedChannelImplBuilder defaultLoadBalancingPolicy(String policy) {
         "directServerAddress is set (%s), which forbids the use of load-balancing policy",
         directServerAddress);
     Preconditions.checkArgument(policy != null, "policy cannot be null");
-    this.defaultLbPolicy = policy;
+    LoadBalancerProvider provider = LoadBalancerRegistry.getDefaultRegistry().getProvider(policy);
+    Preconditions.checkArgument(provider != null, "No provider available for the '%s' load balancing policy.", policy);
+    this.defaultLbPolicy = provider.getPolicyName();
     return this;
   }
 

core/src/test/java/io/grpc/internal/ManagedChannelImplBuilderTest.java

@@ -43,6 +43,9 @@
 import io.grpc.InternalConfiguratorRegistry;
 import io.grpc.InternalFeatureFlags;
 import io.grpc.InternalManagedChannelBuilder.InternalInterceptorFactory;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.LoadBalancerRegistry;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.MethodDescriptor;
@@ -67,6 +70,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Pattern;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -85,6 +89,7 @@ public class ManagedChannelImplBuilderTest {
   private static final String DUMMY_TARGET = "fake-target";
   private static final String DUMMY_AUTHORITY_VALID = "valid:1234";
   private static final String DUMMY_AUTHORITY_INVALID = "[ : : 1]";
+  private static final String FAKE_POLICY_NAME = "magic_balancer";
   private static final ClientInterceptor DUMMY_USER_INTERCEPTOR =
       new ClientInterceptor() {
         @Override
@@ -101,6 +106,29 @@ public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
           return next.newCall(method, callOptions);
         }
       };
+  
+  private final LoadBalancerProvider fakeProvider =
+      new LoadBalancerProvider() {
+        @Override
+        public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
+          return mock(LoadBalancer.class);
+        }
+
+        @Override
+        public boolean isAvailable() {
+          return true;
+        }
+
+        @Override
+        public int getPriority() {
+          return 5;
+        }
+
+        @Override
+        public String getPolicyName() {
+          return FAKE_POLICY_NAME;
+        }
+      };
 
   @Parameters(name = "enableRfc3986UrisParam={0}")
   public static Iterable<Object[]> data() {
@@ -141,6 +169,13 @@ public void setUp() throws Exception {
         DUMMY_TARGET,
         new UnsupportedClientTransportFactoryBuilder(),
         new FixedPortProvider(DUMMY_PORT));
+
+    LoadBalancerRegistry.getDefaultRegistry().register(fakeProvider);
+  }
+
+  @After
+  public void tearDown() {
+    LoadBalancerRegistry.getDefaultRegistry().deregister(fakeProvider);
   }
 
   /** Ensure getDefaultPort() returns default port when no custom implementation provided. */
@@ -249,6 +284,14 @@ public void nameResolverFactory_notAllowedWithDirectAddress() {
     directAddressBuilder.nameResolverFactory(mock(NameResolver.Factory.class));
   }
 
+  @Test
+  public void defaultLoadBalancingPolicy_unregisteredPolicy() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> builder.defaultLoadBalancingPolicy("unregistered_balancer")
+    );
+  }
+
   @Test
   public void defaultLoadBalancingPolicy_default() {
     assertEquals("pick_first", builder.defaultLbPolicy);