Allow users to specify DNS servers to use

src/main/java/com/spotify/dns/DnsSrvResolvers.java

@@ -16,15 +16,17 @@
 
 package com.spotify.dns;
 
-import com.spotify.dns.statistics.DnsReporter;
-
-import org.xbill.DNS.Lookup;
-
 import static com.google.common.primitives.Ints.checkedCast;
 import static java.util.concurrent.TimeUnit.HOURS;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
+import com.spotify.dns.statistics.DnsReporter;
+import java.net.UnknownHostException;
+import java.util.List;
+import org.xbill.DNS.ExtendedResolver;
+import org.xbill.DNS.Resolver;
+
 /**
  * Provides builders for configuring and instantiating {@link DnsSrvResolver}s.
  */
@@ -44,39 +46,51 @@ public static final class DnsSrvResolverBuilder {
     private final boolean cacheLookups;
     private final long dnsLookupTimeoutMillis;
     private final long retentionDurationMillis;
+    private final List<String> servers;
 
     private DnsSrvResolverBuilder() {
       this(null,
            false,
            false,
            SECONDS.toMillis(DEFAULT_DNS_TIMEOUT_SECONDS),
-           HOURS.toMillis(DEFAULT_RETENTION_DURATION_HOURS));
+           HOURS.toMillis(DEFAULT_RETENTION_DURATION_HOURS),
+           null);
     }
 
     private DnsSrvResolverBuilder(
         DnsReporter reporter,
         boolean retainData,
         boolean cacheLookups,
         long dnsLookupTimeoutMillis,
-        long retentionDurationMillis) {
+        long retentionDurationMillis,
+        List<String> servers) {
       this.reporter = reporter;
       this.retainData = retainData;
       this.cacheLookups = cacheLookups;
       this.dnsLookupTimeoutMillis = dnsLookupTimeoutMillis;
       this.retentionDurationMillis = retentionDurationMillis;
+      this.servers = servers;
     }
 
     public DnsSrvResolver build() {
-      // NOTE: this sucks, but is the only reasonably sane way to set a timeout in dnsjava...
-      // the effect of doing this is to set a global timeout for all Lookup instances - except
-      // those that potentially get a new Resolver assigned via the setResolver method... Since
-      // Lookup instances are mostly encapsulated in this library, we should be fine.
+      Resolver resolver;
+      try {
+        // If the user specified DNS servers, create a new ExtendedResolver which uses them.
+        // Otherwise, use the default constructor. That will use the servers in ResolverConfig,
+        // or if that's empty, localhost.
+        resolver = servers == null ?
+                   new ExtendedResolver() :
+                   new ExtendedResolver(servers.toArray(new String[servers.size()]));
+      } catch (UnknownHostException e) {
+        throw new RuntimeException(e);
+      }
+
+      // Configure the Resolver to use our timeouts.
       int timeoutSecs = checkedCast(MILLISECONDS.toSeconds(dnsLookupTimeoutMillis));
       int millisRemainder = checkedCast(dnsLookupTimeoutMillis - SECONDS.toMillis(timeoutSecs));
+      resolver.setTimeout(timeoutSecs, millisRemainder);
 
-      Lookup.getDefaultResolver().setTimeout(timeoutSecs, millisRemainder);
-
-      LookupFactory lookupFactory = new SimpleLookupFactory();
+      LookupFactory lookupFactory = new SimpleLookupFactory(resolver);
 
       if (cacheLookups) {
         lookupFactory = new CachingLookupFactory(lookupFactory);
@@ -97,27 +111,41 @@ public DnsSrvResolver build() {
 
     public DnsSrvResolverBuilder metered(DnsReporter reporter) {
       return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
-                                       retentionDurationMillis);
+                                       retentionDurationMillis, servers);
     }
 
     public DnsSrvResolverBuilder retainingDataOnFailures(boolean retainData) {
       return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
-                                       retentionDurationMillis);
+                                       retentionDurationMillis, servers);
     }
 
     public DnsSrvResolverBuilder cachingLookups(boolean cacheLookups) {
       return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
-                                       retentionDurationMillis);
+                                       retentionDurationMillis, servers);
     }
 
     public DnsSrvResolverBuilder dnsLookupTimeoutMillis(long dnsLookupTimeoutMillis) {
       return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
-                                       retentionDurationMillis);
+                                       retentionDurationMillis, servers);
     }
 
     public DnsSrvResolverBuilder retentionDurationMillis(long retentionDurationMillis) {
       return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
-                                       retentionDurationMillis);
+                                       retentionDurationMillis, servers);
+    }
+
+    /**
+     * Allows the user to specify which DNS servers should be used to perform DNS lookups. Servers
+     * can be specified using either hostname or IP address. If not specified, the underlying DNS
+     * library will determine which servers to use according to the steps documented in
+     * <a href="https://github.com/dnsjava/dnsjava/blob/master/org/xbill/DNS/ResolverConfig.java">
+     * ResolverConfig.java</a>
+     * @param servers the DNS servers to use
+     * @return this builder
+     */
+    public DnsSrvResolverBuilder servers(List<String> servers) {
+      return new DnsSrvResolverBuilder(reporter, retainData, cacheLookups, dnsLookupTimeoutMillis,
+                                       retentionDurationMillis, servers);
     }
   }
 

src/main/java/com/spotify/dns/SimpleLookupFactory.java

@@ -18,17 +18,33 @@
 
 import org.xbill.DNS.DClass;
 import org.xbill.DNS.Lookup;
+import org.xbill.DNS.Resolver;
 import org.xbill.DNS.TextParseException;
 import org.xbill.DNS.Type;
 
 /**
  * A LookupFactory that always returns new instances.
  */
 public class SimpleLookupFactory implements LookupFactory {
+
+  private final Resolver resolver;
+
+  public SimpleLookupFactory() {
+    this(null);
+  }
+
+  public SimpleLookupFactory(Resolver resolver) {
+    this.resolver = resolver;
+  }
+
   @Override
   public Lookup forName(String fqdn) {
     try {
-      return new Lookup(fqdn, Type.SRV, DClass.IN);
+      final Lookup lookup = new Lookup(fqdn, Type.SRV, DClass.IN);
+      if (resolver != null) {
+        lookup.setResolver(resolver);
+      }
+      return lookup;
     } catch (TextParseException e) {
       throw new DnsException("unable to create lookup for name: " + fqdn, e);
     }

src/test/java/com/spotify/dns/DnsSrvResolversIT.java

@@ -21,6 +21,7 @@
 import com.jayway.awaitility.Awaitility;
 import com.spotify.dns.statistics.DnsReporter;
 import com.spotify.dns.statistics.DnsTimingContext;
+import java.util.Arrays;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -29,6 +30,9 @@
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
+import org.xbill.DNS.ExtendedResolver;
+import org.xbill.DNS.Lookup;
+import org.xbill.DNS.SimpleResolver;
 
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.CoreMatchers.is;
@@ -113,6 +117,16 @@ public void shouldFailForBadHostNames() throws Exception {
     }
   }
 
+  @Test
+  public void shouldReturnResultsUsingSpecifiedServers() throws Exception {
+    final String server = new SimpleResolver().getAddress().getHostName();
+    final DnsSrvResolver resolver = DnsSrvResolvers
+        .newBuilder()
+        .servers(ImmutableList.of(server))
+        .build();
+    assertThat(resolver.resolve("_spotify-client._tcp.spotify.com").isEmpty(), is(false));
+  }
+
   @Test
   public void shouldSucceedCreatingRetainingDnsResolver() throws Exception {
     try {
@@ -125,7 +139,6 @@ public void shouldSucceedCreatingRetainingDnsResolver() throws Exception {
       assertTrue("Illegal argument exception should not be thrown", false);
     }
   }
-
   // TODO: it would be nice to be able to also test things like intermittent DNS failures, etc.,
   // but that takes a lot of work setting up a DNS infrastructure that can be made to fail in a
   // controlled way, so I'm skipping that.