refactor getNameResolverProvider()

Author: AgraVator

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

@@ -1472,8 +1472,9 @@ public ClientTransportFactory buildClientTransportFactory() {
               channelCreds,
               callCredentials,
               transportFactoryBuilder,
-              new FixedPortProvider(nameResolverArgs.getDefaultPort()))
-              .nameResolverRegistry(nameResolverRegistry);
+              new FixedPortProvider(nameResolverArgs.getDefaultPort()),
+              nameResolverRegistry,
+              null /* nameResolverProvider */);
         }
 
         @Override

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

@@ -177,6 +177,8 @@ public static ManagedChannelBuilder<?> forTarget(String target) {
   @Nullable
   String authorityOverride;
 
+  boolean registryProvided;
+
   String defaultLbPolicy = GrpcUtil.DEFAULT_LB_POLICY;
 
   boolean fullStreamDecompression;
@@ -343,6 +345,7 @@ public ManagedChannelImplBuilder(
     }
     if (nameResolverRegistry != null) {
       this.nameResolverRegistry = nameResolverRegistry;
+      this.registryProvided = true;
     }
     if (nameResolverProvider != null) {
       this.nameResolverProvider = nameResolverProvider;
@@ -467,6 +470,8 @@ public ManagedChannelImplBuilder nameResolverFactory(NameResolver.Factory resolv
     Preconditions.checkState(directServerAddress == null,
         "directServerAddress is set (%s), which forbids the use of NameResolverFactory",
         directServerAddress);
+    Preconditions.checkState(!registryProvided,
+        "NameResolverRegistry is already set, which forbids the use of NameResolverFactory");
     if (resolverFactory != null) {
       NameResolverRegistry reg = new NameResolverRegistry();
       if (resolverFactory instanceof NameResolverProvider) {
@@ -481,16 +486,11 @@ public ManagedChannelImplBuilder nameResolverFactory(NameResolver.Factory resolv
     return this;
   }
 
-  public ManagedChannelImplBuilder nameResolverRegistry(NameResolverRegistry resolverRegistry) {
+  ManagedChannelImplBuilder nameResolverRegistry(NameResolverRegistry resolverRegistry) {
     this.nameResolverRegistry = resolverRegistry;
     return this;
   }
 
-  public ManagedChannelImplBuilder nameResolverProvider(NameResolverProvider provider) {
-    this.nameResolverProvider = provider;
-    return this;
-  }
-
   @Override
   public ManagedChannelImplBuilder defaultLoadBalancingPolicy(String policy) {
     Preconditions.checkState(directServerAddress == null,
@@ -923,15 +923,21 @@ static ResolvedNameResolver getNameResolverProvider(
       // the provider's default scheme (if provider is specified).
       String scheme = (provider != null)
           ? provider.getDefaultScheme()
-          : nameResolverRegistry.getDefaultScheme();
+          : (nameResolverProvider != null
+              ? nameResolverProvider.getDefaultScheme()
+              : nameResolverRegistry.getDefaultScheme());
       try {
         targetUri = new URI(scheme, "", "/" + target, null);
       } catch (URISyntaxException e) {
         // Should not happen because we just validated the URI.
         throw new IllegalArgumentException(e);
       }
       if (provider == null) {
-        provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+        if (nameResolverProvider != null) {
+          provider = nameResolverProvider;
+        } else {
+          provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+        }
       }
     }
 

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

@@ -26,6 +26,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -48,6 +49,7 @@
 import io.grpc.MethodDescriptor;
 import io.grpc.MetricSink;
 import io.grpc.NameResolver;
+import io.grpc.NameResolverProvider;
 import io.grpc.NameResolverRegistry;
 import io.grpc.StaticTestingClassLoader;
 import io.grpc.internal.ManagedChannelImplBuilder.ChannelBuilderDefaultPortProvider;
@@ -249,6 +251,39 @@ public void nameResolverFactory_notAllowedWithDirectAddress() {
     directAddressBuilder.nameResolverFactory(mock(NameResolver.Factory.class));
   }
 
+  @Test(expected = IllegalStateException.class)
+  @SuppressWarnings("deprecation")
+  public void nameResolverFactory_notAllowedWithRegistry() {
+    builder = new ManagedChannelImplBuilder(
+        DUMMY_TARGET,
+        null,
+        null,
+        new UnsupportedClientTransportFactoryBuilder(),
+        null,
+        new NameResolverRegistry(),
+        null);
+    builder.nameResolverFactory(mock(NameResolver.Factory.class));
+  }
+
+  @Test
+  public void getNameResolverProvider_explicitProviderWithIpTarget() {
+    NameResolverProvider explicitProvider = mock(NameResolverProvider.class);
+    when(explicitProvider.getDefaultScheme()).thenReturn("explicit");
+    when(explicitProvider.newNameResolver(any(URI.class), any(NameResolver.Args.class)))
+        .thenReturn(mock(NameResolver.class));
+
+    // IP address target, likely to fail URI parsing or not match URI pattern
+    String target = "127.0.0.1:8080";
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved = 
+        ManagedChannelImplBuilder.getNameResolverProvider(
+            target,
+            new NameResolverRegistry(),
+            explicitProvider);
+
+    assertSame(explicitProvider, resolved.provider);
+  }
+
   @Test
   public void defaultLoadBalancingPolicy_default() {
     assertEquals("pick_first", builder.defaultLbPolicy);