api: use <scheme,provider> map in nameResoverRegistry

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

@@ -22,7 +22,7 @@
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Comparator;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.logging.Level;
@@ -47,9 +47,10 @@ public final class NameResolverRegistry {
 
   @GuardedBy("this")
   private final LinkedHashSet<NameResolverProvider> allProviders = new LinkedHashSet<>();
-  /** Immutable, sorted version of {@code allProviders}. Is replaced instead of mutating. */
+  /** Generated from {@code allProviders}. Is mapping from scheme key to the highest priority
+   * {@link NameResolverProvider}. */
   @GuardedBy("this")
-  private List<NameResolverProvider> effectiveProviders = Collections.emptyList();
+  private LinkedHashMap<String, NameResolverProvider> effectiveProviders = new LinkedHashMap<>();
 
   /**
    * Register a provider.
@@ -81,16 +82,16 @@ public synchronized void deregister(NameResolverProvider provider) {
   }
 
   private synchronized void refreshProviders() {
-    List<NameResolverProvider> providers = new ArrayList<>(allProviders);
-    // Sort descending based on priority.
-    // sort() must be stable, as we prefer first-registered providers
-    Collections.sort(providers, Collections.reverseOrder(new Comparator<NameResolverProvider>() {
-      @Override
-      public int compare(NameResolverProvider o1, NameResolverProvider o2) {
-        return o1.priority() - o2.priority();
+    LinkedHashMap<String, NameResolverProvider> refreshedProviders = new LinkedHashMap<>();
+    // We prefer first-registered providers
+    for (NameResolverProvider provider : allProviders) {
+      String scheme = provider.getDefaultScheme();
+      NameResolverProvider existing = refreshedProviders.get(scheme);
+      if (existing == null || existing.priority() < provider.priority()) {
+        refreshedProviders.put(scheme, provider);
       }
-    }));
-    effectiveProviders = Collections.unmodifiableList(providers);
+    }
+    effectiveProviders = refreshedProviders;
   }
 
   /**
@@ -120,11 +121,12 @@ public static synchronized NameResolverRegistry getDefaultRegistry() {
   }
 
   /**
-   * Returns effective providers, in priority order.
+   * Returns effective providers map from scheme to the highest priority NameResolverProvider of
+   * that scheme.
    */
   @VisibleForTesting
-  synchronized List<NameResolverProvider> providers() {
-    return effectiveProviders;
+  synchronized LinkedHashMap<String, NameResolverProvider> providers() {
+    return new LinkedHashMap<>(effectiveProviders);
   }
 
   public NameResolver.Factory asFactory() {
@@ -149,23 +151,17 @@ private final class NameResolverFactory extends NameResolver.Factory {
     @Override
     @Nullable
     public NameResolver newNameResolver(URI targetUri, NameResolver.Args args) {
-      List<NameResolverProvider> providers = providers();
-      for (NameResolverProvider provider : providers) {
-        NameResolver resolver = provider.newNameResolver(targetUri, args);
-        if (resolver != null) {
-          return resolver;
-        }
-      }
-      return null;
+      NameResolverProvider provider = providers().get(targetUri.getScheme());
+      return provider == null ? null : provider.newNameResolver(targetUri, args);
     }
 
     @Override
     public String getDefaultScheme() {
-      List<NameResolverProvider> providers = providers();
+      LinkedHashMap<String, NameResolverProvider> providers = providers();
       if (providers.isEmpty()) {
         return "unknown";
       }
-      return providers.get(0).getDefaultScheme();
+      return providers.entrySet().iterator().next().getKey();
     }
   }
 

api/src/test/java/io/grpc/NameResolverRegistryTest.java

@@ -24,6 +24,7 @@
 import io.grpc.internal.DnsNameResolverProvider;
 import java.lang.Thread.UncaughtExceptionHandler;
 import java.net.URI;
+import java.util.LinkedHashMap;
 import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -59,12 +60,16 @@ public void deregister() {
     NameResolverProvider p1 = new BaseProvider(true, 5);
     NameResolverProvider p2 = new BaseProvider(true, 5);
     NameResolverProvider p3 = new BaseProvider(true, 5);
+    String sameScheme = p1.getDefaultScheme();
     reg.register(p1);
     reg.register(p2);
     reg.register(p3);
-    assertThat(reg.providers()).containsExactly(p1, p2, p3).inOrder();
+    assertThat(reg.providers().get(sameScheme)).isSameInstanceAs(p1);
     reg.deregister(p2);
-    assertThat(reg.providers()).containsExactly(p1, p3).inOrder();
+    assertThat(reg.providers().get(sameScheme)).isSameInstanceAs(p1);
+    reg.deregister(p1);
+    assertThat(reg.providers().get(sameScheme)).isSameInstanceAs(p3);
+
   }
 
   @Test
@@ -75,12 +80,13 @@ public void provider_sorted() {
     NameResolverProvider p3 = new BaseProvider(true, 8);
     NameResolverProvider p4 = new BaseProvider(true, 3);
     NameResolverProvider p5 = new BaseProvider(true, 8);
+    String sameScheme = p1.getDefaultScheme();
     reg.register(p1);
     reg.register(p2);
     reg.register(p3);
     reg.register(p4);
     reg.register(p5);
-    assertThat(reg.providers()).containsExactly(p3, p5, p1, p2, p4).inOrder();
+    assertThat(reg.providers().get(sameScheme)).isSameInstanceAs(p3);
   }
 
   @Test
@@ -107,7 +113,7 @@ public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs)
   @Test
   public void newNameResolver_providerReturnsNonNull() {
     NameResolverRegistry registry = new NameResolverRegistry();
-    registry.register(new BaseProvider(true, 5) {
+    registry.register(new BaseProvider(true, 5, uri.getScheme()) {
       @Override
       public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs) {
         return null;
@@ -127,38 +133,74 @@ public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs)
       }
     };
     registry.register(
-        new BaseProvider(true, 4) {
+        new BaseProvider(true, 4, uri.getScheme()) {
           @Override
           public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs) {
             return nr;
           }
         });
     registry.register(
-        new BaseProvider(true, 3) {
+        new BaseProvider(true, 3, uri.getScheme()) {
           @Override
           public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs) {
             fail("Should not be called");
             throw new AssertionError();
           }
         });
-    assertThat(registry.asFactory().newNameResolver(uri, args)).isSameInstanceAs(nr);
+    assertThat(registry.asFactory().newNameResolver(uri, args)).isNull();
+  }
+
+  @Test
+  public void newNameResolver_multipleScheme() {
+    NameResolverRegistry registry = new NameResolverRegistry();
+    registry.register(new BaseProvider(true, 5, uri.getScheme()) {
+      @Override
+      public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs) {
+        return null;
+      }
+    });
+    final NameResolver nr = new NameResolver() {
+      @Override public String getServiceAuthority() {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override public void start(Listener2 listener) {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override public void shutdown() {
+        throw new UnsupportedOperationException();
+      }
+    };
+    registry.register(
+        new BaseProvider(true, 4, "other") {
+          @Override
+          public NameResolver newNameResolver(URI passedUri, NameResolver.Args passedArgs) {
+            return nr;
+          }
+        });
+
+    assertThat(registry.asFactory().newNameResolver(uri, args)).isNull();
+    assertThat(registry.asFactory().newNameResolver(URI.create("other:///0.0.0.0:80"), args))
+            .isSameInstanceAs(nr);
+    assertThat(registry.asFactory().getDefaultScheme()).isEqualTo("dns");
   }
 
   @Test
   public void newNameResolver_noProvider() {
     NameResolver.Factory factory = new NameResolverRegistry().asFactory();
     assertThat(factory.newNameResolver(uri, args)).isNull();
+    assertThat(factory.getDefaultScheme()).isEqualTo("unknown");
   }
 
   @Test
   public void baseProviders() {
-    List<NameResolverProvider> providers = NameResolverRegistry.getDefaultRegistry().providers();
-    assertThat(providers).hasSize(2);
-    // 2 name resolvers from grpclb and core, ordered with decreasing priorities.
-    assertThat(providers.get(0).getClass().getName())
+    LinkedHashMap<String, NameResolverProvider> providers =
+            NameResolverRegistry.getDefaultRegistry().providers();
+    assertThat(providers).hasSize(1);
+    // 2 name resolvers from grpclb and core, higher priority one is returned.
+    assertThat(providers.get("dns").getClass().getName())
         .isEqualTo("io.grpc.grpclb.SecretGrpclbNameResolverProvider$Provider");
-    assertThat(providers.get(1).getClass().getName())
-        .isEqualTo("io.grpc.internal.DnsNameResolverProvider");
     assertThat(NameResolverRegistry.getDefaultRegistry().asFactory().getDefaultScheme())
         .isEqualTo("dns");
   }
@@ -184,10 +226,18 @@ NameResolverProvider.class, getClass().getClassLoader())) {
   private static class BaseProvider extends NameResolverProvider {
     private final boolean isAvailable;
     private final int priority;
+    private final String scheme;
 
     public BaseProvider(boolean isAvailable, int priority) {
       this.isAvailable = isAvailable;
       this.priority = priority;
+      this.scheme = null;
+    }
+
+    public BaseProvider(boolean isAvailable, int priority, String scheme) {
+      this.isAvailable = isAvailable;
+      this.priority = priority;
+      this.scheme = scheme;
     }
 
     @Override
@@ -207,7 +257,7 @@ public NameResolver newNameResolver(URI targetUri, NameResolver.Args args) {
 
     @Override
     public String getDefaultScheme() {
-      return "scheme" + getClass().getSimpleName();
+      return scheme == null ? "scheme" + getClass().getSimpleName() : scheme;
     }
   }
 }

xds/src/test/java/io/grpc/xds/XdsSdsClientServerTest.java

@@ -88,8 +88,11 @@ public class XdsSdsClientServerTest {
   private TlsContextManagerImpl tlsContextManagerForServer;
 
   @Before
-  public void setUp() throws IOException {
+  public void setUp() throws Exception {
     port = XdsServerTestHelper.findFreePort();
+    URI expectedUri = new URI("sdstest://localhost:" + port);
+    fakeNameResolverFactory = new FakeNameResolverFactory.Builder(expectedUri).build();
+    NameResolverRegistry.getDefaultRegistry().register(fakeNameResolverFactory);
   }
 
   @After
@@ -410,9 +413,6 @@ static EnvoyServerProtoData.Listener buildListener(
   private SimpleServiceGrpc.SimpleServiceBlockingStub getBlockingStub(
       final UpstreamTlsContext upstreamTlsContext, String overrideAuthority)
       throws URISyntaxException {
-    URI expectedUri = new URI("sdstest://localhost:" + port);
-    fakeNameResolverFactory = new FakeNameResolverFactory.Builder(expectedUri).build();
-    NameResolverRegistry.getDefaultRegistry().register(fakeNameResolverFactory);
     ManagedChannelBuilder<?> channelBuilder =
         Grpc.newChannelBuilder(
             "sdstest://localhost:" + port,