Allow users to specify DNS servers to use

This PR modifies DnsSrvResolver to that each instance can be configured to use
specific DNS servers. This PR also makes it so that timeouts are configured per
instance, instead of globally.

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,32 @@ 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);
+    }
+
+    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().getAddress().getHostAddress();
+    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.