refactor getNameResolverProvider()

Author: AgraVator

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

@@ -155,6 +155,7 @@ public static ManagedChannelBuilder<?> forTarget(String target) {
 
   private final List<ClientInterceptor> interceptors = new ArrayList<>();
   NameResolverRegistry nameResolverRegistry = NameResolverRegistry.getDefaultRegistry();
+  boolean registryProvided;
   @Nullable
   NameResolverProvider nameResolverProvider;
 
@@ -343,6 +344,7 @@ public ManagedChannelImplBuilder(
     }
     if (nameResolverRegistry != null) {
       this.nameResolverRegistry = nameResolverRegistry;
+      this.registryProvided = true;
     }
     if (nameResolverProvider != null) {
       this.nameResolverProvider = nameResolverProvider;
@@ -467,6 +469,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) {
@@ -483,6 +487,7 @@ public ManagedChannelImplBuilder nameResolverFactory(NameResolver.Factory resolv
 
   ManagedChannelImplBuilder nameResolverRegistry(NameResolverRegistry resolverRegistry) {
     this.nameResolverRegistry = resolverRegistry;
+    this.registryProvided = true;
     return this;
   }
 
@@ -918,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

@@ -48,6 +48,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;
@@ -800,4 +801,121 @@ public void uriPattern() {
   }
 
   private static class CustomSocketAddress extends SocketAddress {}
+
+  @Test
+  public void nameResolverFactory_notAllowedAfterRegistry() {
+    NameResolverRegistry registry = new NameResolverRegistry();
+    builder.nameResolverRegistry(registry);
+    try {
+      builder.nameResolverFactory(mock(NameResolver.Factory.class));
+      fail("Should throw");
+    } catch (IllegalStateException e) {
+      assertThat(e).hasMessageThat().contains("nameResolverRegistry is already set");
+    }
+  }
+
+  @Test
+  public void getNameResolverProvider_explicitProviderWithIpTarget() {
+    String target = "127.0.0.1:8080";
+    NameResolverRegistry registry = new NameResolverRegistry();
+    NameResolverProvider explicitProvider = mock(NameResolverProvider.class);
+    when(explicitProvider.getDefaultScheme()).thenReturn("dns");
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved;
+    resolved = ManagedChannelImplBuilder
+        .getNameResolverProvider(target, registry, explicitProvider);
+
+    assertThat(resolved.provider).isSameInstanceAs(explicitProvider);
+    assertThat(resolved.targetUri.toString()).isEqualTo("dns:///127.0.0.1:8080");
+  }
+
+  @Test
+  public void getNameResolverProvider_explicitProviderWithInvalidUri() {
+    String target = "::1";
+    NameResolverRegistry registry = new NameResolverRegistry();
+    NameResolverProvider explicitProvider = mock(NameResolverProvider.class);
+    when(explicitProvider.getDefaultScheme()).thenReturn("dns");
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved;
+    resolved = ManagedChannelImplBuilder
+        .getNameResolverProvider(target, registry, explicitProvider);
+
+    assertThat(resolved.provider).isSameInstanceAs(explicitProvider);
+    assertThat(resolved.targetUri.toString()).isEqualTo("dns:///::1");
+  }
+
+  @Test
+  public void getNameResolverProvider_explicitProviderWithValidUri() {
+    String target = "dns:///localhost";
+    NameResolverRegistry registry = new NameResolverRegistry();
+    NameResolverProvider explicitProvider = new NameResolverProvider() {
+      @Override
+      public NameResolver newNameResolver(URI targetUri, NameResolver.Args args) {
+        return null;
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return "dns";
+      }
+
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+    };
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved = ManagedChannelImplBuilder.getNameResolverProvider(
+        target, registry, explicitProvider);
+
+    // Should prefer explicit provider if scheme matches?
+    // Current logic: provider passed to getNameResolverProvider is prioritized for
+    // SCHEME determination for fallback
+    // BUT for valid URI, it logic matches URI scheme.
+    // If explicit provider is passed, it is used if target not valid URI.
+    // If target IS valid URI, it checks if provider != null.
+    // Wait, the code:
+    // if (provider == null) { ... }
+    // If explicit 'provider' arg is NOT null, logic uses it?
+    // Let's re-read ManagedChannelImplBuilder.getNameResolverProvider
+    assertThat(resolved.provider).isSameInstanceAs(explicitProvider);
+  }
+
+  @Test
+  public void getNameResolverProvider_registryFallback() {
+    String target = "dns:///localhost";
+    final NameResolverProvider registryProvider = new NameResolverProvider() {
+      @Override
+      public NameResolver newNameResolver(URI targetUri, NameResolver.Args args) {
+        return null;
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return "dns";
+      }
+
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+    };
+    NameResolverRegistry registry = new NameResolverRegistry();
+    registry.register(registryProvider);
+
+    ManagedChannelImplBuilder.ResolvedNameResolver resolved = ManagedChannelImplBuilder.getNameResolverProvider(
+        target, registry, null);
+
+    assertThat(resolved.provider).isSameInstanceAs(registryProvider);
+  }
 }