expose experimental settings for dns

Author: jrhee17

core/src/main/java/com/linecorp/armeria/client/AbstractDnsResolverBuilder.java

@@ -38,10 +38,14 @@
 
 import io.micrometer.core.instrument.MeterRegistry;
 import io.netty.channel.EventLoopGroup;
+import io.netty.channel.socket.DatagramChannel;
+import io.netty.channel.socket.SocketChannel;
 import io.netty.handler.codec.dns.DnsRecord;
 import io.netty.resolver.HostsFileEntriesResolver;
 import io.netty.resolver.dns.BiDnsQueryLifecycleObserverFactory;
+import io.netty.resolver.dns.DnsNameResolver;
 import io.netty.resolver.dns.DnsNameResolverBuilder;
+import io.netty.resolver.dns.DnsNameResolverChannelStrategy;
 import io.netty.resolver.dns.DnsQueryLifecycleObserverFactory;
 import io.netty.resolver.dns.DnsServerAddressStream;
 import io.netty.resolver.dns.DnsServerAddressStreamProvider;
@@ -83,6 +87,11 @@ public abstract class AbstractDnsResolverBuilder<SELF extends AbstractDnsResolve
     private List<String> searchDomains = DnsUtil.defaultSearchDomains();
     private int ndots = DnsUtil.defaultNdots();
     private boolean decodeIdn = true;
+    @Nullable
+    private DnsNameResolverChannelStrategy datagramChannelStrategy;
+    @Nullable
+    private Class<? extends SocketChannel> socketChannelType;
+    private boolean retrySocketChannelOnTimeout;
 
     @Nullable
     private MeterRegistry meterRegistry;
@@ -487,6 +496,36 @@ protected final DnsCache maybeCreateDnsCache() {
         }
     }
 
+    /**
+     * Set the strategy that is used to determine how a {@link DatagramChannel} is used by the resolver for
+     * sending queries over UDP protocol.
+     */
+    @UnstableApi
+    public SELF datagramChannelStrategy(DnsNameResolverChannelStrategy datagramChannelStrategy) {
+        requireNonNull(datagramChannelStrategy, "datagramChannelStrategy");
+        this.datagramChannelStrategy = datagramChannelStrategy;
+        return self();
+    }
+
+    /**
+     * Enables <a href="https://tools.ietf.org/html/rfc7766">TCP fallback</a> using the specified
+     * {@link SocketChannel} type.
+     * <p>
+     * TCP fallback is disabled by default.
+     * @param socketChannelType the type of {@link SocketChannel} to use for TCP fallback.
+     * @param retrySocketChannelOnTimeout if {@code true} the {@link DnsNameResolver} will also fallback to
+     *                                    TCP if a timeout was detected, if {@code false} it will only try to
+     *                                    use TCP if the response was marked as truncated.
+     */
+    @UnstableApi
+    public SELF socketChannelType(Class<? extends SocketChannel> socketChannelType,
+                                  boolean retrySocketChannelOnTimeout) {
+        requireNonNull(socketChannelType, "channelType");
+        this.socketChannelType = socketChannelType;
+        this.retrySocketChannelOnTimeout = retrySocketChannelOnTimeout;
+        return self();
+    }
+
     /**
      * Builds a configurator that configures a {@link DnsNameResolverBuilder} with the properties set.
      */
@@ -561,6 +600,13 @@ protected final Consumer<DnsNameResolverBuilder> buildConfigurator(EventLoopGrou
             if (observerFactory != null) {
                 builder.dnsQueryLifecycleObserverFactory(observerFactory);
             }
+            final DnsNameResolverChannelStrategy datagramChannelStrategy = this.datagramChannelStrategy;
+            if (datagramChannelStrategy != null) {
+                builder.datagramChannelStrategy(datagramChannelStrategy);
+            }
+            final Class<? extends SocketChannel> socketChannelType = this.socketChannelType;
+            final boolean retrySocketChannelOnTimeout = this.retrySocketChannelOnTimeout;
+            builder.socketChannelType(socketChannelType, retrySocketChannelOnTimeout);
         };
     }
 }

core/src/test/java/com/linecorp/armeria/client/endpoint/dns/DnsAddressEndpointGroupTest.java

@@ -50,13 +50,15 @@
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.channel.socket.nio.NioSocketChannel;
 import io.netty.handler.codec.dns.DatagramDnsQuery;
 import io.netty.handler.codec.dns.DefaultDnsQuestion;
 import io.netty.handler.codec.dns.DefaultDnsRawRecord;
 import io.netty.handler.codec.dns.DefaultDnsResponse;
 import io.netty.handler.codec.dns.DnsRecord;
 import io.netty.handler.codec.dns.DnsSection;
 import io.netty.resolver.ResolvedAddressTypes;
+import io.netty.resolver.dns.DnsNameResolverChannelStrategy;
 import io.netty.util.NetUtil;
 import io.netty.util.ReferenceCountUtil;
 
@@ -612,6 +614,40 @@ public void onFailure(List<DnsRecord> oldRecords,
         }
     }
 
+    @Test
+    void experimentalSettingsYieldCorrectResults() {
+        try (TestDnsServer server = new TestDnsServer(ImmutableMap.of(
+                new DefaultDnsQuestion("baz.com.", A),
+                new DefaultDnsResponse(0).addRecord(ANSWER, newAddressRecord("baz.com.", "1.1.1.1"))))) {
+            final AtomicBoolean success = new AtomicBoolean(false);
+            final DnsQueryListener listener = new DnsQueryListener() {
+                @Override
+                public void onSuccess(List<DnsRecord> oldRecords,
+                                      List<DnsRecord> newRecords, String logPrefix) {
+                    success.set(true);
+                }
+
+                @Override
+                public void onFailure(List<DnsRecord> oldRecords,
+                                      Throwable cause, String logPrefix, long delayMillis, int attemptsSoFar) {
+                }
+            };
+            try (DnsAddressEndpointGroup group =
+                         DnsAddressEndpointGroup.builder("baz.com")
+                                                .port(8080)
+                                                .serverAddresses(server.addr())
+                                                .dnsCache(NoopDnsCache.INSTANCE)
+                                                .addDnsQueryListeners(listener)
+                                                .datagramChannelStrategy(
+                                                        DnsNameResolverChannelStrategy.ChannelPerResolution)
+                                                .socketChannelType(NioSocketChannel.class, true)
+                                                .build()) {
+                await().timeout(Duration.ofSeconds(1))
+                       .untilAsserted(() -> assertThat(success.get()).isTrue());
+            }
+        }
+    }
+
     private static DnsRecord newCompatibleAddressRecord(String name, String ipV4Addr) {
         final ByteBuf content = Unpooled.buffer();
         content.writeZero(12);