use the new API in channelProviders

Author: AgraVator

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

@@ -103,7 +103,6 @@ protected NewChannelBuilderResult newChannelBuilder(String target, ChannelCreden
   protected NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentials creds,
                                                       NameResolverRegistry nameResolverRegistry,
                                                       NameResolverProvider nameResolverProvider) {
-    // Implementation note: Currently delegates to the simplified version
     return newChannelBuilder(target, creds);
   }
 

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

@@ -152,6 +152,8 @@ public static ManagedChannelBuilder<?> forTarget(String target) {
 
   private final List<ClientInterceptor> interceptors = new ArrayList<>();
   NameResolverRegistry nameResolverRegistry = NameResolverRegistry.getDefaultRegistry();
+  @Nullable
+  NameResolverProvider nameResolverProvider;
 
   final List<ClientTransportFilter> transportFilters = new ArrayList<>();
 
@@ -433,11 +435,16 @@ public ManagedChannelImplBuilder nameResolverFactory(NameResolver.Factory resolv
     return this;
   }
 
-  ManagedChannelImplBuilder nameResolverRegistry(NameResolverRegistry resolverRegistry) {
+  public 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,
@@ -719,7 +726,8 @@ public ManagedChannel build() {
     ClientTransportFactory clientTransportFactory =
         clientTransportFactoryBuilder.buildClientTransportFactory();
     ResolvedNameResolver resolvedResolver = getNameResolverProvider(
-        target, nameResolverRegistry, clientTransportFactory.getSupportedSocketAddressTypes());
+        target, nameResolverRegistry, nameResolverProvider,
+        clientTransportFactory.getSupportedSocketAddressTypes());
     return new ManagedChannelOrphanWrapper(new ManagedChannelImpl(
         this,
         clientTransportFactory,
@@ -826,7 +834,8 @@ public ResolvedNameResolver(URI targetUri, NameResolverProvider provider) {
   @VisibleForTesting
   static ResolvedNameResolver getNameResolverProvider(
       String target, NameResolverRegistry nameResolverRegistry,
-      Collection<Class<? extends SocketAddress>> channelTransportSocketAddressTypes) {
+      NameResolverProvider nameResolverProvider,
+          Collection<Class<? extends SocketAddress>> channelTransportSocketAddressTypes) {
     // Finding a NameResolver. Try using the target string as the URI. If that fails, try prepending
     // "dns:///".
     NameResolverProvider provider = null;
@@ -841,19 +850,28 @@ static ResolvedNameResolver getNameResolverProvider(
     if (targetUri != null) {
       // For "localhost:8080" this would likely cause provider to be null, because "localhost" is
       // parsed as the scheme. Will hit the next case and try "dns:///localhost:8080".
-      provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      provider = nameResolverProvider;
+      if (provider == null) {
+        provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      }
     }
 
-    if (provider == null && !URI_PATTERN.matcher(target).matches()) {
-      // It doesn't look like a URI target. Maybe it's an authority string. Try with the default
-      // scheme from the registry.
+    if (!URI_PATTERN.matcher(target).matches()) {
+      // It doesn't look like a URI target. Maybe it's an authority string. Try with
+      // the default scheme from the registry (if provider is not specified) or
+      // the provider's default scheme (if provider is specified).
+      String scheme = (provider != null)
+          ? provider.getDefaultScheme()
+          : nameResolverRegistry.getDefaultScheme();
       try {
-        targetUri = new URI(nameResolverRegistry.getDefaultScheme(), "", "/" + target, null);
+        targetUri = new URI(scheme, "", "/" + target, null);
       } catch (URISyntaxException e) {
-        // Should not be possible.
+        // Should not happen because we just validated the URI.
         throw new IllegalArgumentException(e);
       }
-      provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      if (provider == null) {
+        provider = nameResolverRegistry.getProviderForScheme(targetUri.getScheme());
+      }
     }
 
     if (provider == null) {

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

@@ -96,7 +96,7 @@ public void validTargetNoProvider() {
     NameResolverRegistry nameResolverRegistry = new NameResolverRegistry();
     try {
       ManagedChannelImplBuilder.getNameResolverProvider(
-          "foo.googleapis.com:8080", nameResolverRegistry,
+          "foo.googleapis.com:8080", nameResolverRegistry, null,
           Collections.singleton(InetSocketAddress.class));
       fail("Should fail");
     } catch (IllegalArgumentException e) {
@@ -109,7 +109,7 @@ public void validTargetProviderAddrTypesNotSupported() {
     NameResolverRegistry nameResolverRegistry = getTestRegistry("testscheme");
     try {
       ManagedChannelImplBuilder.getNameResolverProvider(
-          "testscheme:///foo.googleapis.com:8080", nameResolverRegistry,
+          "testscheme:///foo.googleapis.com:8080", nameResolverRegistry, null,
           Collections.singleton(CustomSocketAddress.class));
       fail("Should fail");
     } catch (IllegalArgumentException e) {
@@ -123,7 +123,7 @@ private void testValidTarget(String target, String expectedUriString, URI expect
     NameResolverRegistry nameResolverRegistry = getTestRegistry(expectedUri.getScheme());
     ManagedChannelImplBuilder.ResolvedNameResolver resolved =
         ManagedChannelImplBuilder.getNameResolverProvider(
-            target, nameResolverRegistry, Collections.singleton(InetSocketAddress.class));
+            target, nameResolverRegistry, null, Collections.singleton(InetSocketAddress.class));
     assertThat(resolved.provider).isInstanceOf(FakeNameResolverProvider.class);
     assertThat(resolved.targetUri).isEqualTo(expectedUri);
     assertThat(resolved.targetUri.toString()).isEqualTo(expectedUriString);
@@ -135,7 +135,7 @@ private void testInvalidTarget(String target) {
     try {
       ManagedChannelImplBuilder.ResolvedNameResolver resolved =
           ManagedChannelImplBuilder.getNameResolverProvider(
-              target, nameResolverRegistry, Collections.singleton(InetSocketAddress.class));
+              target, nameResolverRegistry, null, Collections.singleton(InetSocketAddress.class));
       FakeNameResolverProvider nameResolverProvider = (FakeNameResolverProvider) resolved.provider;
       fail("Should have failed, but got resolver provider " + nameResolverProvider);
     } catch (IllegalArgumentException e) {

netty/src/main/java/io/grpc/netty/NettyChannelBuilder.java

@@ -38,6 +38,8 @@
 import io.grpc.HttpConnectProxiedSocketAddress;
 import io.grpc.Internal;
 import io.grpc.ManagedChannelBuilder;
+import io.grpc.NameResolverProvider;
+import io.grpc.NameResolverRegistry;
 import io.grpc.internal.AtomicBackoff;
 import io.grpc.internal.ClientTransportFactory;
 import io.grpc.internal.ConnectionClientTransport;
@@ -708,6 +710,24 @@ NettyChannelBuilder setTransportTracerFactory(TransportTracer.Factory transportT
     return this;
   }
 
+  /**
+   * Sets the registry used for looking up name resolvers.
+   */
+  @CanIgnoreReturnValue
+  public NettyChannelBuilder nameResolverRegistry(NameResolverRegistry registry) {
+    managedChannelImplBuilder.nameResolverRegistry(registry);
+    return this;
+  }
+
+  /**
+   * Sets the {@link io.grpc.NameResolverProvider} to use.
+   */
+  @CanIgnoreReturnValue
+  public NettyChannelBuilder nameResolverProvider(NameResolverProvider provider) {
+    managedChannelImplBuilder.nameResolverProvider(provider);
+    return this;
+  }
+
   static Collection<Class<? extends SocketAddress>> getSupportedSocketAddressTypes() {
     return Collections.singleton(InetSocketAddress.class);
   }

netty/src/main/java/io/grpc/netty/NettyChannelProvider.java

@@ -19,6 +19,8 @@
 import io.grpc.ChannelCredentials;
 import io.grpc.Internal;
 import io.grpc.ManagedChannelProvider;
+import io.grpc.NameResolverProvider;
+import io.grpc.NameResolverRegistry;
 import java.net.SocketAddress;
 import java.util.Collection;
 
@@ -55,6 +57,25 @@ public NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentia
         new NettyChannelBuilder(target, creds, result.callCredentials, result.negotiator));
   }
 
+  @Override
+  public NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentials creds,
+                                                   NameResolverRegistry nameResolverRegistry,
+                                                   NameResolverProvider nameResolverProvider) {
+    ProtocolNegotiators.FromChannelCredentialsResult result = ProtocolNegotiators.from(creds);
+    if (result.error != null) {
+      return NewChannelBuilderResult.error(result.error);
+    }
+    NettyChannelBuilder builder = new NettyChannelBuilder(target, creds,
+        result.callCredentials, result.negotiator);
+    if (nameResolverRegistry != null) {
+      builder.nameResolverRegistry(nameResolverRegistry);
+    }
+    if (nameResolverProvider != null) {
+      builder.nameResolverProvider(nameResolverProvider);
+    }
+    return NewChannelBuilderResult.channelBuilder(builder);
+  }
+
   @Override
   protected Collection<Class<? extends SocketAddress>> getSupportedSocketAddressTypes() {
     return NettyChannelBuilder.getSupportedSocketAddressTypes();

netty/src/main/java/io/grpc/netty/UdsNettyChannelProvider.java

@@ -20,6 +20,8 @@
 import io.grpc.ChannelCredentials;
 import io.grpc.Internal;
 import io.grpc.ManagedChannelProvider;
+import io.grpc.NameResolverProvider;
+import io.grpc.NameResolverRegistry;
 import io.grpc.internal.SharedResourcePool;
 import io.netty.channel.unix.DomainSocketAddress;
 import java.net.SocketAddress;
@@ -62,6 +64,21 @@ public NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentia
     return result;
   }
 
+  @Override
+  public NewChannelBuilderResult newChannelBuilder(String target, ChannelCredentials creds,
+                                                   NameResolverRegistry nameResolverRegistry,
+                                                   NameResolverProvider nameResolverProvider) {
+    Preconditions.checkState(isAvailable());
+    NewChannelBuilderResult result = new NettyChannelProvider().newChannelBuilder(
+        target, creds, nameResolverRegistry, nameResolverProvider);
+    if (result.getChannelBuilder() != null) {
+      ((NettyChannelBuilder) result.getChannelBuilder())
+          .eventLoopGroupPool(SharedResourcePool.forResource(Utils.DEFAULT_WORKER_EVENT_LOOP_GROUP))
+          .channelType(Utils.EPOLL_DOMAIN_CLIENT_CHANNEL_TYPE, DomainSocketAddress.class);
+    }
+    return result;
+  }
+
   @Override
   protected Collection<Class<? extends SocketAddress>> getSupportedSocketAddressTypes() {
     return Collections.singleton(DomainSocketAddress.class);

netty/src/test/java/io/grpc/netty/NettyChannelProviderTest.java

@@ -87,4 +87,164 @@ public void newChannelBuilder_fail() {
         TlsChannelCredentials.newBuilder().requireFakeFeature().build());
     assertThat(result.getError()).contains("FAKE");
   }
+
+  @Test
+  public void newChannelBuilder_withRegistry() {
+    io.grpc.NameResolverRegistry registry = new io.grpc.NameResolverRegistry();
+    NewChannelBuilderResult result = provider.newChannelBuilder(
+        "localhost:443", TlsChannelCredentials.create(), registry, null);
+    assertThat(result.getChannelBuilder()).isInstanceOf(NettyChannelBuilder.class);
+  }
+
+  @Test
+  public void newChannelBuilder_withProvider() {
+    io.grpc.NameResolverProvider resolverProvider = new io.grpc.NameResolverProvider() {
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return "dns";
+      }
+
+      @Override
+      public io.grpc.NameResolver newNameResolver(java.net.URI targetUri,
+          io.grpc.NameResolver.Args args) {
+        return null;
+      }
+    };
+    NewChannelBuilderResult result = provider.newChannelBuilder(
+        "localhost:443", TlsChannelCredentials.create(), null,
+        resolverProvider);
+    assertThat(result.getChannelBuilder()).isInstanceOf(NettyChannelBuilder.class);
+  }
+
+  @Test
+  public void newChannelBuilder_registryPropagation_e2e() {
+    String scheme = "testscheme";
+    final io.grpc.NameResolverRegistry registry = new io.grpc.NameResolverRegistry();
+    final java.util.concurrent.atomic.AtomicReference<io.grpc.NameResolverRegistry>
+        capturedRegistry = new java.util.concurrent.atomic.AtomicReference<>();
+
+    final io.grpc.NameResolverProvider resolverProvider = new io.grpc.NameResolverProvider() {
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return scheme;
+      }
+
+      @Override
+      public io.grpc.NameResolver newNameResolver(java.net.URI targetUri,
+          io.grpc.NameResolver.Args args) {
+        capturedRegistry.set(args.getNameResolverRegistry());
+        return new io.grpc.NameResolver() {
+          @Override
+          public String getServiceAuthority() {
+            return "authority";
+          }
+
+          @Override
+          public void start(Listener2 listener) {
+          }
+
+          @Override
+          public void shutdown() {
+          }
+        };
+      }
+    };
+    registry.register(resolverProvider);
+
+    NewChannelBuilderResult result = provider.newChannelBuilder(
+        scheme + ":///target", TlsChannelCredentials.create(), registry,
+        null);
+    assertThat(result.getChannelBuilder()).isInstanceOf(NettyChannelBuilder.class);
+    // Verify build() succeeds
+    result.getChannelBuilder().build();
+
+    // Verify the registry passed to args is the exact same instance
+    assertSame("Registry should be propagated to NameResolver.Args", registry,
+        capturedRegistry.get());
+
+    // Verify default registry (empty) fails
+    NewChannelBuilderResult defaultResult = provider.newChannelBuilder(
+        scheme + ":///target", TlsChannelCredentials.create(),
+        new io.grpc.NameResolverRegistry(), null);
+    // The provider might still return a builder, but build() should fail if it
+    // can't find the resolver.
+    // However, NettyChannelProvider just delegates to NettyChannelBuilder.
+    // NettyChannelBuilder delegates to ManagedChannelImplBuilder.
+    // ManagedChannelImplBuilder.build() calls getNameResolverProvider(), which
+    // throws if not found.
+    try {
+      defaultResult.getChannelBuilder().build();
+      fail("Should have failed to build() without correct registry");
+    } catch (IllegalArgumentException e) {
+      // Expected
+    }
+  }
+
+  @Test
+  public void newChannelBuilder_providerPropagation_e2e() {
+    String scheme = "otherscheme";
+    final io.grpc.NameResolverProvider resolverProvider = new io.grpc.NameResolverProvider() {
+      @Override
+      protected boolean isAvailable() {
+        return true;
+      }
+
+      @Override
+      protected int priority() {
+        return 5;
+      }
+
+      @Override
+      public String getDefaultScheme() {
+        return scheme;
+      }
+
+      @Override
+      public io.grpc.NameResolver newNameResolver(java.net.URI targetUri,
+          io.grpc.NameResolver.Args args) {
+        return new io.grpc.NameResolver() {
+          @Override
+          public String getServiceAuthority() {
+            return "authority";
+          }
+
+          @Override
+          public void start(Listener2 listener) {
+          }
+
+          @Override
+          public void shutdown() {
+          }
+        };
+      }
+    };
+
+    // Pass explicit provider, null registry
+    NewChannelBuilderResult result = provider.newChannelBuilder(
+        scheme + ":///target", TlsChannelCredentials.create(),
+        null, resolverProvider);
+    assertThat(result.getChannelBuilder()).isInstanceOf(NettyChannelBuilder.class);
+    // Should succeed because we passed the specific provider
+    result.getChannelBuilder().build();
+  }
 }