Introduce custom NameResolver.Args

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

@@ -18,6 +18,7 @@
 
 import com.google.common.base.MoreObjects;
 import com.google.errorprone.annotations.DoNotCall;
+import io.grpc.NameResolver;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -263,6 +264,12 @@ protected T addMetricSink(MetricSink metricSink) {
     return thisT();
   }
 
+  @Override
+  public <X> T setNameResolverArg(NameResolver.Args.Key<X> key, X value) {
+    delegate().setNameResolverArg(key, value);
+    return thisT();
+  }
+
   /**
    * Returns the {@link ManagedChannel} built by the delegate by default. Overriding method can
    * return different value.

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

@@ -17,6 +17,7 @@
 package io.grpc;
 
 import com.google.common.base.Preconditions;
+import io.grpc.NameResolver;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -633,6 +634,23 @@
     throw new UnsupportedOperationException();
   }
 
+  /**
+   * Provides an "extended" argument for the {@link NameResolver}, if applicable, replacing any
+   * 'value' previously provided for 'key'.
+   *
+   * <p>NB: If the selected {@link NameResolver} does not understand 'key', or target URI resolution
+   * isn't needed at all, your extended argument will be silently ignored.
+   *
+   * <p>See {@link NameResolver.Args.Extensions} for more.
+   *
+   * @param key identifies the argument in a type-safe manner
+   * @param value the argument itself
+   * @return this
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1770")
+  public <X> T setNameResolverArg(NameResolver.Args.Key<X> key, X value) {
+    throw new UnsupportedOperationException();
+  }
 
   /**
    * Builds a channel using the given parameters.

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

@@ -28,11 +28,13 @@
 import java.lang.annotation.RetentionPolicy;
 import java.net.URI;
 import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import javax.annotation.Nullable;
+import javax.annotation.concurrent.Immutable;
 import javax.annotation.concurrent.ThreadSafe;
 
 /**
@@ -274,7 +276,12 @@
   /**
    * Information that a {@link Factory} uses to create a {@link NameResolver}.
    *
-   * <p>Note this class doesn't override neither {@code equals()} nor {@code hashCode()}.
+   * <p>Args applicable to all {@link NameResolver}s are defined here using ordinary setters and
+   * getters. This container can also hold externally-defined "extended" args that aren't so widely
+   * useful or that would be inappropriate dependencies for this low level API. See {@link
+   * Args.Extensions} for more.
+   *
+   * <p>Note this class overrides neither {@code equals()} nor {@code hashCode()}.
    *
    * @since 1.21.0
    */
@@ -284,34 +291,31 @@
     private final ProxyDetector proxyDetector;
     private final SynchronizationContext syncContext;
     private final ServiceConfigParser serviceConfigParser;
+    private final Extensions extensions;
     @Nullable private final ScheduledExecutorService scheduledExecutorService;
     @Nullable private final ChannelLogger channelLogger;
     @Nullable private final Executor executor;
     @Nullable private final String overrideAuthority;
 
-    private Args(
-        Integer defaultPort,
-        ProxyDetector proxyDetector,
-        SynchronizationContext syncContext,
-        ServiceConfigParser serviceConfigParser,
-        @Nullable ScheduledExecutorService scheduledExecutorService,
-        @Nullable ChannelLogger channelLogger,
-        @Nullable Executor executor,
-        @Nullable String overrideAuthority) {
-      this.defaultPort = checkNotNull(defaultPort, "defaultPort not set");
-      this.proxyDetector = checkNotNull(proxyDetector, "proxyDetector not set");
-      this.syncContext = checkNotNull(syncContext, "syncContext not set");
-      this.serviceConfigParser = checkNotNull(serviceConfigParser, "serviceConfigParser not set");
-      this.scheduledExecutorService = scheduledExecutorService;
-      this.channelLogger = channelLogger;
-      this.executor = executor;
-      this.overrideAuthority = overrideAuthority;
+    private Args(Builder builder) {
+      this.defaultPort = checkNotNull(builder.defaultPort, "defaultPort not set");
+      this.proxyDetector = checkNotNull(builder.proxyDetector, "proxyDetector not set");
+      this.syncContext = checkNotNull(builder.syncContext, "syncContext not set");
+      this.serviceConfigParser =
+          checkNotNull(builder.serviceConfigParser, "serviceConfigParser not set");
+      this.extensions = checkNotNull(builder.extensions, "extensions not set");
+      this.scheduledExecutorService = builder.scheduledExecutorService;
+      this.channelLogger = builder.channelLogger;
+      this.executor = builder.executor;
+      this.overrideAuthority = builder.overrideAuthority;
     }
 
     /**
      * The port number used in case the target or the underlying naming system doesn't provide a
      * port number.
      *
+     * <p>TODO: Only meaningful for InetSocketAddress producers. Move to {@link Extensions}?
+     *
      * @since 1.21.0
      */
     public int getDefaultPort() {
@@ -366,6 +370,15 @@
       return serviceConfigParser;
     }
 
+    /**
+     * Gets the value of an extended arg by key, or {@code null} if it's not set.
+     */
+    @SuppressWarnings("unchecked")
+    @Nullable
+    public <T> T getExtension(Key<T> key) {
+      return extensions.get(key);
+    }
+
     /**
      * Returns the {@link ChannelLogger} for the Channel served by this NameResolver.
      *
@@ -411,6 +424,7 @@
           .add("proxyDetector", proxyDetector)
           .add("syncContext", syncContext)
           .add("serviceConfigParser", serviceConfigParser)
+          .add("extensions", extensions)
           .add("scheduledExecutorService", scheduledExecutorService)
           .add("channelLogger", channelLogger)
           .add("executor", executor)
@@ -429,6 +443,7 @@
       builder.setProxyDetector(proxyDetector);
       builder.setSynchronizationContext(syncContext);
       builder.setServiceConfigParser(serviceConfigParser);
+      builder.setExtensions(extensions);
       builder.setScheduledExecutorService(scheduledExecutorService);
       builder.setChannelLogger(channelLogger);
       builder.setOffloadExecutor(executor);
@@ -459,6 +474,7 @@
       private ChannelLogger channelLogger;
       private Executor executor;
       private String overrideAuthority;
+      private Extensions extensions = Extensions.EMPTY;
 
       Builder() {
       }
@@ -545,16 +561,169 @@
         return this;
       }
 
+      /**
+       * See {@link Args#getExtension(Key)}.
+       *
+       * <p>Optional, {@link Extensions#EMPTY} by default.
+       */
+      public Builder setExtensions(Extensions extensions) {
+        this.extensions = extensions;
+        return this;
+      }
+
       /**
        * Builds an {@link Args}.
        *
        * @since 1.21.0
        */
       public Args build() {
-        return
-            new Args(
-                defaultPort, proxyDetector, syncContext, serviceConfigParser,
-                scheduledExecutorService, channelLogger, executor, overrideAuthority);
+        return new Args(this);
+      }
+    }
+
+    /**
+     * Identifies an externally-defined extension argument that can be stored in {@link Args}.
+     *
+     * <p>Uses reference equality so keys should be defined as global constants.
+     *
+     * @param <T> type of values that can be stored under this key
+     */
+    @Immutable
+    @SuppressWarnings("UnusedTypeParameter")
+    public static final class Key<T> {
+      private final String debugString;
+
+      private Key(String debugString) {
+        this.debugString = debugString;
+      }
+
+      @Override
+      public String toString() {
+        return debugString;
+      }
+
+      /**
+       * Creates a new instance of {@link Key}.
+       *
+       * @param debugString a string used to describe the key, used for debugging.
+       * @param <T> Key type
+       * @return a new instance of Key
+       */
+      public static <T> Key<T> create(String debugString) {
+        return new Key<>(debugString);
+      }
+    }
+
+    /**
+     * An immutable type-safe container of externally-defined {@link NameResolver} arguments.
+     *
+     * <p>While ordinary {@link Args} should be universally useful and meaningful, argument {@link
+     * Extensions} can apply just to resolvers of a certain URI scheme, just to resolvers producing
+     * a particular type of {@link java.net.SocketAddress}, or even an individual {@link
+     * NameResolver} subclass. Extended args are identified by an instance of {@link Args.Key} which
+     * should be defined in a java package and class appropriate to the argument's scope.
+     *
+     * <p>{@link Args} are normally reserved for information in *support* of name resolution, not
+     * the name to be resolved itself. However, there are rare cases where all or part of the target
+     * name can't be represented by any standard URI scheme or can't be encoded as a String at all.
+     * Extensions, in contrast, can be an arbitrary Java type, making them a useful work around in
+     * these cases.
+     *
+     * <p>Extensions can also be used simply to avoid adding inappropriate deps to the low level
+     * io.grpc package.
+     *
+     * <p>NB: This class overrides neither {@code equals()} nor {@code hashCode()}.
+     */
+    @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1770")
+    @Immutable
+    public static final class Extensions {
+      private static final IdentityHashMap<Key<?>, Object> EMPTY_MAP = new IdentityHashMap<>();
+
+      /**
+       * The canonical empty instance.
+       */
+      public static final Extensions EMPTY = new Extensions(EMPTY_MAP);
+
+      private final IdentityHashMap<Key<?>, Object> data;
+
+      private Extensions(IdentityHashMap<Key<?>, Object> data) {
+        assert data != null;
+        this.data = data;
+      }
+
+      /**
+       * Gets the value for the key, or {@code null} if it's not present.
+       */
+      @SuppressWarnings("unchecked")
+      @Nullable
+      public <T> T get(Key<T> key) {
+        return (T) data.get(key);
+      }
+
+      /**
+       * Creates a new builder.
+       */
+      public static Builder newBuilder() {
+        return new Builder(EMPTY);
+      }
+
+      /**
+       * Creates a new builder that is pre-populated with the content of this container.
+       * @return a new builder.
+       */
+      public Builder toBuilder() {
+        return new Builder(this);
+      }
+
+      @Override
+      public String toString() {
+        return data.toString();
+      }
+
+      /**
+       * Fluently builds instances of {@link Extensions}.
+       */
+      public static final class Builder {
+        private Extensions base;
+        private IdentityHashMap<Key<?>, Object> newdata;
+
+        private Builder(Extensions base) {
+          assert base != null;
+          this.base = base;
+        }
+
+        private IdentityHashMap<Key<?>, Object> data(int size) {
+          if (newdata == null) {
+            newdata = new IdentityHashMap<>(size);
+          }
+          return newdata;
+        }
+
+        /**
+         * Associates 'value' with 'key', replacing any previously associated value.
+         *
+         * @return this
+         */
+        public <T> Builder set(Key<T> key, T value) {
+          data(1).put(key, value);
+          return this;
+        }
+
+        /**
+         * Builds the extensions.
+         */
+        public Extensions build() {
+          if (newdata != null) {
+            for (Map.Entry<Key<?>, Object> entry : base.data.entrySet()) {
+              if (!newdata.containsKey(entry.getKey())) {
+                newdata.put(entry.getKey(), entry.getValue());
+              }
+            }
+            base = new Extensions(newdata);
+            newdata = null;
+          }
+          return base;
+        }
       }
     }
   }

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

@@ -599,6 +599,7 @@ ClientStream newSubstream(
             .setChannelLogger(channelLogger)
             .setOffloadExecutor(this.offloadExecutorHolder)
             .setOverrideAuthority(this.authorityOverride)
+            .setExtensions(builder.nameResolverArgsExtBuilder.build())
             .build();
     this.nameResolver = getNameResolver(
         targetUri, authorityOverride, nameResolverProvider, nameResolverArgs);

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

@@ -151,6 +151,8 @@ public static ManagedChannelBuilder<?> forTarget(String target) {
 
   private final List<ClientInterceptor> interceptors = new ArrayList<>();
   NameResolverRegistry nameResolverRegistry = NameResolverRegistry.getDefaultRegistry();
+  final NameResolver.Args.Extensions.Builder nameResolverArgsExtBuilder =
+      NameResolver.Args.Extensions.newBuilder();
 
   final List<ClientTransportFilter> transportFilters = new ArrayList<>();
 
@@ -554,6 +556,12 @@ public ManagedChannelImplBuilder defaultServiceConfig(@Nullable Map<String, ?> s
     return this;
   }
 
+  @Override
+  public <X> ManagedChannelImplBuilder setNameResolverArg(NameResolver.Args.Key<X> key, X value) {
+    nameResolverArgsExtBuilder.set(key, value);
+    return this;
+  }
+
   @Nullable
   private static Map<String, ?> checkMapEntryTypes(@Nullable Map<?, ?> map) {
     if (map == null) {

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

@@ -764,6 +764,16 @@ public void disableNameResolverServiceConfig() {
     assertThat(builder.lookUpServiceConfig).isFalse();
   }
 
+  @Test
+  public void setNameResolverExtArgs() {
+    assertThat(builder.nameResolverArgsExtBuilder.build())
+        .isSameInstanceAs(NameResolver.Args.Extensions.EMPTY);
+
+    NameResolver.Args.Key<Integer> testKey = NameResolver.Args.Key.create("test-key");
+    builder.setNameResolverArg(testKey, 42);
+    assertThat(builder.nameResolverArgsExtBuilder.build().get(testKey)).isEqualTo(42);
+  }
+
   @Test
   public void metricSinks() {
     MetricSink mocksink = mock(MetricSink.class);