xds: allow injecting bootstrapOverride in xdsNameResolverProvider (#8358)

xds/src/main/java/io/grpc/xds/SharedXdsClientPoolProvider.java

@@ -52,7 +52,7 @@ final class SharedXdsClientPoolProvider implements XdsClientPoolFactory {
   private final AtomicReference<Map<String, ?>> bootstrapOverride = new AtomicReference<>();
   private volatile ObjectPool<XdsClient> xdsClientPool;
 
-  private SharedXdsClientPoolProvider() {
+  SharedXdsClientPoolProvider() {
     this(new BootstrapperImpl());
   }
 

xds/src/main/java/io/grpc/xds/XdsNameResolver.java

@@ -124,22 +124,25 @@ final class XdsNameResolver extends NameResolver {
   private ResolveState resolveState;
 
   XdsNameResolver(String name, ServiceConfigParser serviceConfigParser,
-      SynchronizationContext syncContext, ScheduledExecutorService scheduler) {
+      SynchronizationContext syncContext, ScheduledExecutorService scheduler,
+      @Nullable Map<String, ?> bootstrapOverride) {
     this(name, serviceConfigParser, syncContext, scheduler,
         SharedXdsClientPoolProvider.getDefaultProvider(), ThreadSafeRandomImpl.instance,
-        FilterRegistry.getDefaultRegistry());
+        FilterRegistry.getDefaultRegistry(), bootstrapOverride);
   }
 
   @VisibleForTesting
   XdsNameResolver(String name, ServiceConfigParser serviceConfigParser,
       SynchronizationContext syncContext, ScheduledExecutorService scheduler,
       XdsClientPoolFactory xdsClientPoolFactory, ThreadSafeRandom random,
-      FilterRegistry filterRegistry) {
+      FilterRegistry filterRegistry, @Nullable Map<String, ?> bootstrapOverride) {
     authority = GrpcUtil.checkAuthority(checkNotNull(name, "name"));
     this.serviceConfigParser = checkNotNull(serviceConfigParser, "serviceConfigParser");
     this.syncContext = checkNotNull(syncContext, "syncContext");
     this.scheduler = checkNotNull(scheduler, "scheduler");
-    this.xdsClientPoolFactory = checkNotNull(xdsClientPoolFactory, "xdsClientPoolFactory");
+    this.xdsClientPoolFactory = bootstrapOverride == null ? checkNotNull(xdsClientPoolFactory,
+            "xdsClientPoolFactory") : new SharedXdsClientPoolProvider();
+    this.xdsClientPoolFactory.setBootstrapOverride(bootstrapOverride);
     this.random = checkNotNull(random, "random");
     this.filterRegistry = checkNotNull(filterRegistry, "filterRegistry");
     logId = InternalLogId.allocate("xds-resolver", name);

xds/src/main/java/io/grpc/xds/XdsNameResolverProvider.java

@@ -42,10 +42,31 @@ import javax.annotation.Nullable;
 public final class XdsNameResolverProvider extends NameResolverProvider {
 
   private static final String SCHEME = "xds";
+  private final String scheme;
+  private final Map<String, ?> bootstrapOverride;
+
+  public XdsNameResolverProvider() {
+    this(SCHEME, null);
+  }
+
+  private XdsNameResolverProvider(String scheme,
+                                 @Nullable Map<String, ?> bootstrapOverride) {
+    this.scheme = checkNotNull(scheme, "scheme");
+    this.bootstrapOverride = bootstrapOverride;
+  }
+
+  /**
+   * A convenient method to allow creating a {@link XdsNameResolverProvider} with custom scheme
+   * and bootstrap.
+   */
+  public static XdsNameResolverProvider createForTest(String scheme,
+                                                      @Nullable Map<String, ?> bootstrapOverride) {
+    return new XdsNameResolverProvider(scheme, bootstrapOverride);
+  }
 
   @Override
   public XdsNameResolver newNameResolver(URI targetUri, Args args) {
-    if (SCHEME.equals(targetUri.getScheme())) {
+    if (scheme.equals(targetUri.getScheme())) {
       String targetPath = checkNotNull(targetUri.getPath(), "targetPath");
       Preconditions.checkArgument(
           targetPath.startsWith("/"),
@@ -54,14 +75,15 @@ public final class XdsNameResolverProvider extends NameResolverProvider {
           targetUri);
       String name = targetPath.substring(1);
       return new XdsNameResolver(name, args.getServiceConfigParser(),
-          args.getSynchronizationContext(), args.getScheduledExecutorService());
+              args.getSynchronizationContext(), args.getScheduledExecutorService(),
+              bootstrapOverride);
     }
     return null;
   }
 
   @Override
   public String getDefaultScheme() {
-    return SCHEME;
+    return scheme;
   }
 
   @Override

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

@@ -20,15 +20,20 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 
+import com.google.common.collect.ImmutableMap;
 import io.grpc.ChannelLogger;
 import io.grpc.InternalServiceProviders;
 import io.grpc.NameResolver;
 import io.grpc.NameResolver.ServiceConfigParser;
 import io.grpc.NameResolverProvider;
+import io.grpc.NameResolverRegistry;
 import io.grpc.SynchronizationContext;
 import io.grpc.internal.FakeClock;
 import io.grpc.internal.GrpcUtil;
 import java.net.URI;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -114,4 +119,54 @@ public class XdsNameResolverProviderTest {
       // Expected
     }
   }
+
+  @Test
+  public void newProvider_multipleScheme() {
+    NameResolverRegistry registry = NameResolverRegistry.getDefaultRegistry();
+    XdsNameResolverProvider provider0 = XdsNameResolverProvider.createForTest("no-scheme", null);
+    registry.register(provider0);
+    XdsNameResolverProvider provider1 = XdsNameResolverProvider.createForTest("new-xds-scheme",
+            new HashMap<String, String>());
+    registry.register(provider1);
+    assertThat(registry.asFactory()
+            .newNameResolver(URI.create("xds:///localhost"), args)).isNotNull();
+    assertThat(registry.asFactory()
+            .newNameResolver(URI.create("new-xds-scheme:///localhost"), args)).isNotNull();
+    assertThat(registry.asFactory()
+            .newNameResolver(URI.create("no-scheme:///localhost"), args)).isNotNull();
+    registry.deregister(provider1);
+    assertThat(registry.asFactory()
+            .newNameResolver(URI.create("new-xds-scheme:///localhost"), args)).isNull();
+    registry.deregister(provider0);
+    assertThat(registry.asFactory()
+            .newNameResolver(URI.create("xds:///localhost"), args)).isNotNull();
+  }
+
+  @Test
+  public void newProvider_overrideBootstrap() {
+    Map<String, ?> b = ImmutableMap.of(
+            "node", ImmutableMap.of(
+                    "id", "ENVOY_NODE_ID",
+                    "cluster", "ENVOY_CLUSTER"),
+            "xds_servers", Collections.singletonList(
+                    ImmutableMap.of(
+                            "server_uri", "trafficdirector.googleapis.com:443",
+                            "channel_creds", Collections.singletonList(
+                                    ImmutableMap.of("type", "insecure")
+                            )
+                    )
+            )
+    );
+    NameResolverRegistry registry = new NameResolverRegistry();
+    XdsNameResolverProvider provider = XdsNameResolverProvider.createForTest("no-scheme", b);
+    registry.register(provider);
+    NameResolver resolver = registry.asFactory()
+            .newNameResolver(URI.create("no-scheme:///localhost"), args);
+    resolver.start(mock(NameResolver.Listener2.class));
+    assertThat(resolver).isInstanceOf(XdsNameResolver.class);
+    assertThat(((XdsNameResolver)resolver).getXdsClient().getBootstrapInfo().getNode().getId())
+            .isEqualTo("ENVOY_NODE_ID");
+    resolver.shutdown();
+    registry.deregister(provider);
+  }
 }

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

@@ -153,7 +153,7 @@ public class XdsNameResolverTest {
         new FaultFilter(mockRandom, new AtomicLong()),
         RouterFilter.INSTANCE);
     resolver = new XdsNameResolver(AUTHORITY, serviceConfigParser, syncContext, scheduler,
-        xdsClientPoolFactory, mockRandom, filterRegistry);
+        xdsClientPoolFactory, mockRandom, filterRegistry, null);
   }
 
   @After
@@ -172,7 +172,6 @@ public class XdsNameResolverTest {
     XdsClientPoolFactory xdsClientPoolFactory = new XdsClientPoolFactory() {
       @Override
       public void setBootstrapOverride(Map<String, ?> bootstrap) {
-        throw new UnsupportedOperationException("Should not be called");
       }
 
       @Override
@@ -187,7 +186,7 @@ public class XdsNameResolverTest {
       }
     };
     resolver = new XdsNameResolver(AUTHORITY, serviceConfigParser, syncContext, scheduler,
-        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry());
+        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry(), null);
     resolver.start(mockListener);
     verify(mockListener).onError(errorCaptor.capture());
     Status error = errorCaptor.getValue();
@@ -437,7 +436,7 @@ public class XdsNameResolverTest {
     ServiceConfigParser realParser = new ScParser(
         true, 5, 5, new AutoConfiguredLoadBalancerFactory("pick-first"));
     resolver = new XdsNameResolver(AUTHORITY, realParser, syncContext, scheduler,
-        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry());
+        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry(), null);
     resolver.start(mockListener);
     FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient();
     RetryPolicy retryPolicy = RetryPolicy.create(
@@ -640,7 +639,7 @@ public class XdsNameResolverTest {
     resolver.shutdown();
     reset(mockListener);
     resolver = new XdsNameResolver(AUTHORITY, serviceConfigParser, syncContext, scheduler,
-        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry());
+        xdsClientPoolFactory, mockRandom, FilterRegistry.getDefaultRegistry(), null);
     resolver.start(mockListener);
     xdsClient = (FakeXdsClient) resolver.getXdsClient();
     xdsClient.deliverLdsUpdate(
@@ -1701,10 +1700,8 @@ public class XdsNameResolverTest {
   }
 
   private final class FakeXdsClientPoolFactory implements XdsClientPoolFactory {
-
     @Override
     public void setBootstrapOverride(Map<String, ?> bootstrap) {
-      throw new UnsupportedOperationException("Should not be called");
     }
 
     @Override