Introduce custom NameResolver.Args

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

@@ -263,6 +263,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

@@ -633,6 +633,23 @@
     throw new UnsupportedOperationException();
   }
 
+  /**
+   * Provides a "custom" 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 custom argument will be silently ignored.
+   *
+   * <p>See {@link NameResolver.Args#getArg(NameResolver.Args.Key)} 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;
 
 /**
@@ -276,7 +278,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 "custom" args that aren't so widely
+   * useful or that would be inappropriate dependencies for this low level API. See {@link
+   * Args#getArg} for more.
+   *
+   * <p>Note this class overrides neither {@code equals()} nor {@code hashCode()}.
    *
    * @since 1.21.0
    */
@@ -291,26 +298,20 @@
     @Nullable private final Executor executor;
     @Nullable private final String overrideAuthority;
     @Nullable private final MetricRecorder metricRecorder;
+    @Nullable private final IdentityHashMap<Key<?>, Object> customArgs;
 
-    private Args(
-        Integer defaultPort,
-        ProxyDetector proxyDetector,
-        SynchronizationContext syncContext,
-        ServiceConfigParser serviceConfigParser,
-        @Nullable ScheduledExecutorService scheduledExecutorService,
-        @Nullable ChannelLogger channelLogger,
-        @Nullable Executor executor,
-        @Nullable String overrideAuthority,
-        @Nullable MetricRecorder metricRecorder) {
-      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;
-      this.metricRecorder = metricRecorder;
+    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.scheduledExecutorService = builder.scheduledExecutorService;
+      this.channelLogger = builder.channelLogger;
+      this.executor = builder.executor;
+      this.overrideAuthority = builder.overrideAuthority;
+      this.metricRecorder = builder.metricRecorder;
+      this.customArgs = cloneCustomArgs(builder.customArgs);
     }
 
     /**
@@ -319,6 +320,7 @@
      *
      * @since 1.21.0
      */
+    // <p>TODO: Only meaningful for InetSocketAddress producers. Make this a custom arg?
     public int getDefaultPort() {
       return defaultPort;
     }
@@ -371,6 +373,30 @@
       return serviceConfigParser;
     }
 
+    /**
+     * Returns the value of a custom arg named 'key', or {@code null} if it's not set.
+     *
+     * <p>While ordinary {@link Args} should be universally useful and meaningful, custom arguments
+     * 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.
+     * Custom args are identified by an instance of {@link Args.Key} which should be a constant
+     * defined in a java package and class appropriate for 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.
+     * Custom args, in contrast, can hold arbitrary Java types, making them a useful work around in
+     * these cases.
+     *
+     * <p>Custom args can also be used simply to avoid adding inappropriate deps to the low level
+     * io.grpc package.
+     */
+    @SuppressWarnings("unchecked") // Cast is safe because all put()s go through the setArg() API.
+    @Nullable
+    public <T> T getArg(Key<T> key) {
+      return customArgs != null ? (T) customArgs.get(key) : null;
+    }
+
     /**
      * Returns the {@link ChannelLogger} for the Channel served by this NameResolver.
      *
@@ -424,6 +450,7 @@
           .add("proxyDetector", proxyDetector)
           .add("syncContext", syncContext)
           .add("serviceConfigParser", serviceConfigParser)
+          .add("customArgs", customArgs)
           .add("scheduledExecutorService", scheduledExecutorService)
           .add("channelLogger", channelLogger)
           .add("executor", executor)
@@ -448,6 +475,7 @@
       builder.setOffloadExecutor(executor);
       builder.setOverrideAuthority(overrideAuthority);
       builder.setMetricRecorder(metricRecorder);
+      builder.customArgs = cloneCustomArgs(customArgs);
       return builder;
     }
 
@@ -475,6 +503,7 @@
       private Executor executor;
       private String overrideAuthority;
       private MetricRecorder metricRecorder;
+      private IdentityHashMap<Key<?>, Object> customArgs;
 
       Builder() {
       }
@@ -561,6 +590,17 @@
         return this;
       }
 
+      /** See {@link Args#getArg(Key)}. */
+      public <T> Builder setArg(Key<T> key, T value) {
+        checkNotNull(key, "key");
+        checkNotNull(value, "value");
+        if (customArgs == null) {
+          customArgs = new IdentityHashMap<>();
+        }
+        customArgs.put(key, value);
+        return this;
+      }
+
       /**
        * See {@link Args#getMetricRecorder()}. This is an optional field.
        */
@@ -575,11 +615,40 @@
        * @since 1.21.0
        */
       public Args build() {
-        return
-            new Args(
-                defaultPort, proxyDetector, syncContext, serviceConfigParser,
-                scheduledExecutorService, channelLogger, executor, overrideAuthority,
-                metricRecorder);
+        return new Args(this);
+      }
+    }
+
+    /**
+     * Identifies an externally-defined custom 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);
       }
     }
   }
@@ -877,4 +946,10 @@
       }
     }
   }
+
+  @Nullable
+  private static IdentityHashMap<Args.Key<?>, Object> cloneCustomArgs(
+          @Nullable IdentityHashMap<Args.Key<?>, Object> customArgs) {
+    return customArgs != null ? new IdentityHashMap<>(customArgs) : null;
+  }
 }

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

@@ -47,9 +47,13 @@ public class NameResolverTest {
   private static final List<EquivalentAddressGroup> ADDRESSES =
       Collections.singletonList(
           new EquivalentAddressGroup(new FakeSocketAddress("fake-address-1"), Attributes.EMPTY));
-  private static final Attributes.Key<String> YOLO_KEY = Attributes.Key.create("yolo");
-  private static Attributes ATTRIBUTES = Attributes.newBuilder()
-      .set(YOLO_KEY, "To be, or not to be?").build();
+  private static final Attributes.Key<String> YOLO_ATTR_KEY = Attributes.Key.create("yolo");
+  private static Attributes ATTRIBUTES =
+      Attributes.newBuilder().set(YOLO_ATTR_KEY, "To be, or not to be?").build();
+  private static final NameResolver.Args.Key<Integer> FOO_ARG_KEY =
+      NameResolver.Args.Key.create("foo");
+  private static final NameResolver.Args.Key<Integer> BAR_ARG_KEY =
+          NameResolver.Args.Key.create("bar");
   private static ConfigOrError CONFIG = ConfigOrError.fromConfig("foo");
 
   @Rule
@@ -65,6 +69,7 @@ public class NameResolverTest {
   private final Executor executor = Executors.newSingleThreadExecutor();
   private final String overrideAuthority = "grpc.io";
   private final MetricRecorder metricRecorder = new MetricRecorder() {};
+  private final int customArgValue = 42;
   @Mock NameResolver.Listener mockListener;
 
   @Test
@@ -79,6 +84,8 @@ public void args() {
     assertThat(args.getOffloadExecutor()).isSameInstanceAs(executor);
     assertThat(args.getOverrideAuthority()).isSameInstanceAs(overrideAuthority);
     assertThat(args.getMetricRecorder()).isSameInstanceAs(metricRecorder);
+    assertThat(args.getArg(FOO_ARG_KEY)).isEqualTo(customArgValue);
+    assertThat(args.getArg(BAR_ARG_KEY)).isNull();
 
     NameResolver.Args args2 = args.toBuilder().build();
     assertThat(args2.getDefaultPort()).isEqualTo(defaultPort);
@@ -90,6 +97,8 @@ public void args() {
     assertThat(args2.getOffloadExecutor()).isSameInstanceAs(executor);
     assertThat(args2.getOverrideAuthority()).isSameInstanceAs(overrideAuthority);
     assertThat(args.getMetricRecorder()).isSameInstanceAs(metricRecorder);
+    assertThat(args.getArg(FOO_ARG_KEY)).isEqualTo(customArgValue);
+    assertThat(args.getArg(BAR_ARG_KEY)).isNull();
 
     assertThat(args2).isNotSameInstanceAs(args);
     assertThat(args2).isNotEqualTo(args);
@@ -106,6 +115,7 @@ private NameResolver.Args createArgs() {
         .setOffloadExecutor(executor)
         .setOverrideAuthority(overrideAuthority)
         .setMetricRecorder(metricRecorder)
+        .setArg(FOO_ARG_KEY, customArgValue)
         .build();
   }
 

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

@@ -591,8 +591,7 @@ ClientStream newSubstream(
     this.authorityOverride = builder.authorityOverride;
     this.metricRecorder = new MetricRecorderImpl(builder.metricSinks,
         MetricInstrumentRegistry.getDefaultRegistry());
-    this.nameResolverArgs =
-        NameResolver.Args.newBuilder()
+    NameResolver.Args.Builder nameResolverArgsBuilder = NameResolver.Args.newBuilder()
             .setDefaultPort(builder.getDefaultPort())
             .setProxyDetector(proxyDetector)
             .setSynchronizationContext(syncContext)
@@ -601,8 +600,9 @@ ClientStream newSubstream(
             .setChannelLogger(channelLogger)
             .setOffloadExecutor(this.offloadExecutorHolder)
             .setOverrideAuthority(this.authorityOverride)
-            .setMetricRecorder(this.metricRecorder)
-            .build();
+            .setMetricRecorder(this.metricRecorder);
+    builder.copyAllNameResolverCustomArgsTo(nameResolverArgsBuilder);
+    this.nameResolverArgs = nameResolverArgsBuilder.build();
     this.nameResolver = getNameResolver(
         targetUri, authorityOverride, nameResolverProvider, nameResolverArgs);
     this.balancerRpcExecutorPool = checkNotNull(balancerRpcExecutorPool, "balancerRpcExecutorPool");

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

@@ -55,6 +55,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -159,6 +160,8 @@ public static ManagedChannelBuilder<?> forTarget(String target) {
   final ChannelCredentials channelCredentials;
   @Nullable
   final CallCredentials callCredentials;
+  @Nullable
+  IdentityHashMap<NameResolver.Args.Key<?>, Object> nameResolverCustomArgs;
 
   @Nullable
   private final SocketAddress directServerAddress;
@@ -613,6 +616,24 @@ private static List<?> checkListEntryTypes(List<?> list) {
     return Collections.unmodifiableList(parsedList);
   }
 
+  @Override
+  public <X> ManagedChannelImplBuilder setNameResolverArg(NameResolver.Args.Key<X> key, X value) {
+    if (nameResolverCustomArgs == null) {
+      nameResolverCustomArgs = new IdentityHashMap<>();
+    }
+    nameResolverCustomArgs.put(key, value);
+    return this;
+  }
+
+  @SuppressWarnings("unchecked") // This cast is safe because of setNameResolverArg()'s signature.
+  void copyAllNameResolverCustomArgsTo(NameResolver.Args.Builder dest) {
+    if (nameResolverCustomArgs != null) {
+      for (Map.Entry<NameResolver.Args.Key<?>, Object> entry : nameResolverCustomArgs.entrySet()) {
+        dest.setArg((NameResolver.Args.Key<Object>) entry.getKey(), entry.getValue());
+      }
+    }
+  }
+
   @Override
   public ManagedChannelImplBuilder disableServiceConfigLookUp() {
     this.lookUpServiceConfig = false;

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.nameResolverCustomArgs)
+        .isNull();
+
+    NameResolver.Args.Key<Integer> testKey = NameResolver.Args.Key.create("test-key");
+    builder.setNameResolverArg(testKey, 42);
+    assertThat(builder.nameResolverCustomArgs.get(testKey)).isEqualTo(42);
+  }
+
   @Test
   public void metricSinks() {
     MetricSink mocksink = mock(MetricSink.class);