Support per-MongoClient DNS lookup configuration (#1104)

* Add InetAddressResolver interface
* Add InetAddressResolverProvider interface
* Add configuration of InetAddressResolver and DnsClient in MongoClientSettings and ConnectionString
* Use configured InetAddressResolver and DnsClient to resolve all SRV, TXT, A, and AAAA records

JAVA-4911

Author: jyemin

driver-core/src/main/com/mongodb/ConnectionString.java

@@ -23,6 +23,7 @@
 import com.mongodb.internal.diagnostics.logging.Loggers;
 import com.mongodb.internal.dns.DefaultDnsResolver;
 import com.mongodb.lang.Nullable;
+import com.mongodb.spi.dns.DnsClient;
 import org.bson.UuidRepresentation;
 
 import java.io.UnsupportedEncodingException;
@@ -296,6 +297,21 @@ public class ConnectionString {
      * @since 3.0
      */
     public ConnectionString(final String connectionString) {
+        this(connectionString, null);
+    }
+
+    /**
+     * Creates a ConnectionString from the given string with the given {@link DnsClient}.
+     *
+     * <p>If setting {@link MongoClientSettings#getDnsClient()} explicitly, care should be taken to call this constructor with the same
+     * {@link DnsClient}.
+     *
+     * @param connectionString the connection string
+     * @param dnsClient        the DNS client with which to resolve TXT record for the mongodb+srv protocol
+     * @since 4.10
+     * @see MongoClientSettings#getDnsClient()
+     */
+    public ConnectionString(final String connectionString, @Nullable final DnsClient dnsClient) {
         this.connectionString = connectionString;
         boolean isMongoDBProtocol = connectionString.startsWith(MONGODB_PREFIX);
         isSrvProtocol = connectionString.startsWith(MONGODB_SRV_PREFIX);
@@ -394,7 +410,7 @@ public ConnectionString(final String connectionString) {
         }
 
         String txtRecordsQueryParameters = isSrvProtocol
-                ? new DefaultDnsResolver().resolveAdditionalQueryParametersFromTxtRecords(unresolvedHosts.get(0)) : "";
+                ? new DefaultDnsResolver(dnsClient).resolveAdditionalQueryParametersFromTxtRecords(unresolvedHosts.get(0)) : "";
         String connectionStringQueryParameters = unprocessedConnectionString;
 
         Map<String, List<String>> connectionStringOptionsMap = parseOptions(connectionStringQueryParameters);

driver-core/src/main/com/mongodb/MongoClientSettings.java

@@ -29,6 +29,8 @@
 import com.mongodb.connection.StreamFactoryFactory;
 import com.mongodb.event.CommandListener;
 import com.mongodb.lang.Nullable;
+import com.mongodb.spi.dns.DnsClient;
+import com.mongodb.spi.dns.InetAddressResolver;
 import org.bson.UuidRepresentation;
 import org.bson.codecs.BsonCodecProvider;
 import org.bson.codecs.BsonValueCodecProvider;
@@ -107,6 +109,8 @@ public final class MongoClientSettings {
     private final boolean heartbeatConnectTimeoutSetExplicitly;
 
     private final ContextProvider contextProvider;
+    private final DnsClient dnsClient;
+    private final InetAddressResolver inetAddressResolver;
 
     /**
      * Gets the default codec registry.  It includes the following providers:
@@ -160,6 +164,39 @@ public static Builder builder(final MongoClientSettings settings) {
         return new Builder(settings);
     }
 
+    /**
+     * Gets the {@link DnsClient} to use for resolving DNS queries.
+     *
+     * <p>If set, it will be used to resolve SRV and TXT records for mongodb+srv connections. Otherwise,
+     * implementations of {@link com.mongodb.spi.dns.DnsClientProvider} will be discovered via {@link java.util.ServiceLoader}.
+     * If no implementations are discovered, then {@code com.sun.jndi.dns.DnsContextFactory} will be used to resolve these records.
+     *
+     * <p>If applying a connection string to these settings, care must be taken to also pass the same {@link DnsClient} as an argument to
+     * the {@link ConnectionString} constructor.
+     *
+     * @return the DNS client
+     * @since 4.10
+     * @see ConnectionString#ConnectionString(String, DnsClient)
+     */
+    @Nullable
+    public DnsClient getDnsClient() {
+        return dnsClient;
+    }
+
+    /**
+     * Gets the {@link InetAddressResolver} to use for looking up the {@link java.net.InetAddress} instances for each host.
+     *
+     * <p>If set, it will be used to look up the {@link java.net.InetAddress} for each host, via
+     * {@link InetAddressResolver#lookupByName(String)}. Otherwise, {@link java.net.InetAddress#getAllByName(String)} will be used.
+     *
+     * @return the {@link java.net.InetAddress} resolver
+     * @since 4.10
+     */
+    @Nullable
+    public InetAddressResolver getInetAddressResolver() {
+        return inetAddressResolver;
+    }
+
     /**
      * A builder for {@code MongoClientSettings} so that {@code MongoClientSettings} can be immutable, and to support easier construction
      * through chaining.
@@ -193,6 +230,8 @@ public static final class Builder {
         private int heartbeatSocketTimeoutMS;
 
         private ContextProvider contextProvider;
+        private DnsClient dnsClient;
+        private InetAddressResolver inetAddressResolver;
 
         private Builder() {
         }
@@ -211,6 +250,8 @@ private Builder(final MongoClientSettings settings) {
             credential = settings.getCredential();
             uuidRepresentation = settings.getUuidRepresentation();
             serverApi = settings.getServerApi();
+            dnsClient = settings.getDnsClient();
+            inetAddressResolver = settings.getInetAddressResolver();
             streamFactoryFactory = settings.getStreamFactoryFactory();
             autoEncryptionSettings = settings.getAutoEncryptionSettings();
             contextProvider = settings.getContextProvider();
@@ -220,6 +261,7 @@ private Builder(final MongoClientSettings settings) {
             socketSettingsBuilder.applySettings(settings.getSocketSettings());
             connectionPoolSettingsBuilder.applySettings(settings.getConnectionPoolSettings());
             sslSettingsBuilder.applySettings(settings.getSslSettings());
+
             if (settings.heartbeatConnectTimeoutSetExplicitly) {
                 heartbeatConnectTimeoutMS = settings.heartbeatSocketSettings.getConnectTimeout(MILLISECONDS);
             }
@@ -273,7 +315,7 @@ public Builder applyConnectionString(final ConnectionString connectionString) {
         /**
          * Applies the {@link LoggerSettings.Builder} block and then sets the loggerSettings.
          *
-         * @param block the block to apply to the LoggerSettins.
+         * @param block the block to apply to the LoggerSettings.
          * @return this
          * @see MongoClientSettings#getLoggerSettings()
          * @since 4.9
@@ -580,6 +622,45 @@ public Builder contextProvider(@Nullable final ContextProvider contextProvider)
             return this;
         }
 
+        /**
+         * Sets the {@link DnsClient} to use for resolving DNS queries.
+         *
+         * <p> If set, it will be used to resolve SRV and TXT records for mongodb+srv connections. Otherwise,
+         * implementation of {@link com.mongodb.spi.dns.DnsClientProvider} will be discovered via {@link java.util.ServiceLoader}
+         * and used to create an instance of {@link DnsClient}. If no implementation is discovered, then
+         * {@code com.sun.jndi.dns.DnsContextFactory} will be used to resolve these records.
+         *
+         * <p>If {@linkplain #applyConnectionString(ConnectionString) applying a connection string to these settings}, care must be
+         * taken to also pass the same {@link DnsClient} as an argument to the {@link ConnectionString} constructor.
+         *
+         * @param dnsClient the DNS client
+         * @return the DNS client
+         * @since 4.10
+         * @see ConnectionString#ConnectionString(String, DnsClient)
+         */
+        public Builder dnsClient(@Nullable final DnsClient dnsClient) {
+            this.dnsClient = dnsClient;
+            return this;
+        }
+
+        /**
+         * Sets the {@link InetAddressResolver} to use for looking up the {@link java.net.InetAddress} instances for each host.
+         *
+         * <p>If set, it will be used to look up the {@link java.net.InetAddress} for each host, via
+         * {@link InetAddressResolver#lookupByName(String)}. Otherwise,
+         * an implementation of {@link com.mongodb.spi.dns.InetAddressResolverProvider} will be discovered via
+         * {@link java.util.ServiceLoader} and used to create an instance of {@link InetAddressResolver}.  If no implementation is
+         * discovered, {@link java.net.InetAddress#getAllByName(String)} will be used to lookup the {@link java.net.InetAddress}
+         * instances for a host.
+         *
+         * @param inetAddressResolver the InetAddress provider
+         * @return the {@link java.net.InetAddress} resolver
+         * @since 4.10
+         */
+        public Builder inetAddressResolver(@Nullable final InetAddressResolver inetAddressResolver) {
+            this.inetAddressResolver = inetAddressResolver;
+            return this;
+        }
 
         // Package-private to provide interop with MongoClientOptions
         Builder heartbeatConnectTimeoutMS(final int heartbeatConnectTimeoutMS) {
@@ -905,6 +986,8 @@ public boolean equals(final Object o) {
                 && uuidRepresentation == that.uuidRepresentation
                 && Objects.equals(serverApi, that.serverApi)
                 && Objects.equals(autoEncryptionSettings, that.autoEncryptionSettings)
+                && Objects.equals(dnsClient, that.dnsClient)
+                && Objects.equals(inetAddressResolver, that.inetAddressResolver)
                 && Objects.equals(contextProvider, that.contextProvider);
     }
 
@@ -913,7 +996,8 @@ public int hashCode() {
         return Objects.hash(readPreference, writeConcern, retryWrites, retryReads, readConcern, credential, streamFactoryFactory,
                 commandListeners, codecRegistry, loggerSettings, clusterSettings, socketSettings, heartbeatSocketSettings,
                 connectionPoolSettings, serverSettings, sslSettings, applicationName, compressorList, uuidRepresentation, serverApi,
-                autoEncryptionSettings, heartbeatSocketTimeoutSetExplicitly, heartbeatConnectTimeoutSetExplicitly, contextProvider);
+                autoEncryptionSettings, heartbeatSocketTimeoutSetExplicitly, heartbeatConnectTimeoutSetExplicitly, dnsClient,
+                inetAddressResolver, contextProvider);
     }
 
     @Override
@@ -940,6 +1024,8 @@ public String toString() {
                 + ", uuidRepresentation=" + uuidRepresentation
                 + ", serverApi=" + serverApi
                 + ", autoEncryptionSettings=" + autoEncryptionSettings
+                + ", dnsClient=" + dnsClient
+                + ", inetAddressResolver=" + inetAddressResolver
                 + ", contextProvider=" + contextProvider
                 + '}';
     }
@@ -964,6 +1050,8 @@ private MongoClientSettings(final Builder builder) {
         compressorList = builder.compressorList;
         uuidRepresentation = builder.uuidRepresentation;
         serverApi = builder.serverApi;
+        dnsClient = builder.dnsClient;
+        inetAddressResolver = builder.inetAddressResolver;
         autoEncryptionSettings = builder.autoEncryptionSettings;
         heartbeatSocketSettings = SocketSettings.builder()
                 .readTimeout(builder.heartbeatSocketTimeoutMS == 0

driver-core/src/main/com/mongodb/MongoSocketException.java

@@ -33,7 +33,7 @@ public class MongoSocketException extends MongoException {
      * @param msg the message
      * @param e the cause
      */
-    MongoSocketException(final String msg, final ServerAddress serverAddress, final Throwable e) {
+    public MongoSocketException(final String msg, final ServerAddress serverAddress, final Throwable e) {
         super(-2, msg, e);
         this.serverAddress = serverAddress;
     }

driver-core/src/main/com/mongodb/internal/connection/DefaultClusterFactory.java

@@ -32,6 +32,8 @@
 import com.mongodb.event.ServerListener;
 import com.mongodb.event.ServerMonitorListener;
 import com.mongodb.lang.Nullable;
+import com.mongodb.spi.dns.DnsClient;
+import com.mongodb.spi.dns.InetAddressResolver;
 
 import java.util.List;
 
@@ -59,7 +61,8 @@ public Cluster createCluster(final ClusterSettings originalClusterSettings, fina
                                  @Nullable final CommandListener commandListener,
                                  @Nullable final String applicationName,
                                  @Nullable final MongoDriverInformation mongoDriverInformation,
-                                 final List<MongoCompressor> compressorList, @Nullable final ServerApi serverApi) {
+                                 final List<MongoCompressor> compressorList, @Nullable final ServerApi serverApi,
+                                 @Nullable final DnsClient dnsClient, @Nullable final InetAddressResolver inetAddressResolver) {
 
         ClusterId clusterId = new ClusterId(applicationName);
         ClusterSettings clusterSettings;
@@ -87,20 +90,20 @@ public Cluster createCluster(final ClusterSettings originalClusterSettings, fina
                     .build();
         }
 
-        DnsSrvRecordMonitorFactory dnsSrvRecordMonitorFactory = new DefaultDnsSrvRecordMonitorFactory(clusterId, serverSettings);
+        DnsSrvRecordMonitorFactory dnsSrvRecordMonitorFactory = new DefaultDnsSrvRecordMonitorFactory(clusterId, serverSettings, dnsClient);
 
         if (clusterSettings.getMode() == ClusterConnectionMode.LOAD_BALANCED) {
             ClusterableServerFactory serverFactory = new LoadBalancedClusterableServerFactory(serverSettings,
                     connectionPoolSettings, internalConnectionPoolSettings, streamFactory, credential, loggerSettings, commandListener,
                     applicationName, mongoDriverInformation != null ? mongoDriverInformation : MongoDriverInformation.builder().build(),
-                    compressorList, serverApi);
+                    compressorList, serverApi, inetAddressResolver);
             return new LoadBalancedCluster(clusterId, clusterSettings, serverFactory, dnsSrvRecordMonitorFactory);
         } else {
             ClusterableServerFactory serverFactory = new DefaultClusterableServerFactory(serverSettings,
                     connectionPoolSettings, internalConnectionPoolSettings,
                     streamFactory, heartbeatStreamFactory, credential, loggerSettings, commandListener, applicationName,
                     mongoDriverInformation != null ? mongoDriverInformation : MongoDriverInformation.builder().build(), compressorList,
-                    serverApi);
+                    serverApi, inetAddressResolver);
 
             if (clusterSettings.getMode() == ClusterConnectionMode.SINGLE) {
                 return new SingleServerCluster(clusterId, clusterSettings, serverFactory);

driver-core/src/main/com/mongodb/internal/connection/DefaultClusterableServerFactory.java

@@ -31,6 +31,7 @@
 import com.mongodb.event.ServerListener;
 import com.mongodb.internal.inject.SameObjectProvider;
 import com.mongodb.lang.Nullable;
+import com.mongodb.spi.dns.InetAddressResolver;
 
 import java.util.List;
 
@@ -54,6 +55,8 @@ public class DefaultClusterableServerFactory implements ClusterableServerFactory
     private final List<MongoCompressor> compressorList;
     @Nullable
     private final ServerApi serverApi;
+    @Nullable
+    private final InetAddressResolver inetAddressResolver;
 
     public DefaultClusterableServerFactory(
             final ServerSettings serverSettings, final ConnectionPoolSettings connectionPoolSettings,
@@ -63,7 +66,8 @@ public DefaultClusterableServerFactory(
             final LoggerSettings loggerSettings,
             @Nullable final CommandListener commandListener,
             @Nullable final String applicationName, @Nullable final MongoDriverInformation mongoDriverInformation,
-            final List<MongoCompressor> compressorList, @Nullable final ServerApi serverApi) {
+            final List<MongoCompressor> compressorList, @Nullable final ServerApi serverApi,
+            @Nullable final InetAddressResolver inetAddressResolver) {
         this.serverSettings = serverSettings;
         this.connectionPoolSettings = connectionPoolSettings;
         this.internalConnectionPoolSettings = internalConnectionPoolSettings;
@@ -76,6 +80,7 @@ public DefaultClusterableServerFactory(
         this.mongoDriverInformation = mongoDriverInformation;
         this.compressorList = compressorList;
         this.serverApi = serverApi;
+        this.inetAddressResolver = inetAddressResolver;
     }
 
     @Override
@@ -86,11 +91,11 @@ public ClusterableServer create(final Cluster cluster, final ServerAddress serve
         ServerMonitor serverMonitor = new DefaultServerMonitor(serverId, serverSettings, cluster.getClock(),
                 // no credentials, compressor list, or command listener for the server monitor factory
                 new InternalStreamConnectionFactory(clusterMode, true, heartbeatStreamFactory, null, applicationName,
-                        mongoDriverInformation, emptyList(), loggerSettings, null, serverApi),
+                        mongoDriverInformation, emptyList(), loggerSettings, null, serverApi, inetAddressResolver),
                 clusterMode, serverApi, sdamProvider);
         ConnectionPool connectionPool = new DefaultConnectionPool(serverId,
                 new InternalStreamConnectionFactory(clusterMode, streamFactory, credential, applicationName,
-                        mongoDriverInformation, compressorList, loggerSettings, commandListener, serverApi),
+                        mongoDriverInformation, compressorList, loggerSettings, commandListener, serverApi, inetAddressResolver),
                 connectionPoolSettings, internalConnectionPoolSettings, sdamProvider);
         ServerListener serverListener = singleServerListener(serverSettings);
         SdamServerDescriptionManager sdam = new DefaultSdamServerDescriptionManager(cluster, serverId, serverListener, serverMonitor,

driver-core/src/main/com/mongodb/internal/connection/DefaultDnsSrvRecordMonitorFactory.java

@@ -19,6 +19,8 @@
 import com.mongodb.connection.ClusterId;
 import com.mongodb.connection.ServerSettings;
 import com.mongodb.internal.dns.DefaultDnsResolver;
+import com.mongodb.lang.Nullable;
+import com.mongodb.spi.dns.DnsClient;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
@@ -32,15 +34,17 @@ public class DefaultDnsSrvRecordMonitorFactory implements DnsSrvRecordMonitorFac
 
     private final ClusterId clusterId;
     private final long noRecordsRescanFrequency;
+    private final DnsClient dnsClient;
 
-    public DefaultDnsSrvRecordMonitorFactory(final ClusterId clusterId, final ServerSettings serverSettings) {
+    public DefaultDnsSrvRecordMonitorFactory(final ClusterId clusterId, final ServerSettings serverSettings, @Nullable final DnsClient dnsClient) {
         this.clusterId = clusterId;
         this.noRecordsRescanFrequency = serverSettings.getHeartbeatFrequency(MILLISECONDS);
+        this.dnsClient = dnsClient;
     }
 
     @Override
     public DnsSrvRecordMonitor create(final String hostName, final String srvServiceName, final DnsSrvRecordInitializer dnsSrvRecordInitializer) {
         return new DefaultDnsSrvRecordMonitor(hostName, srvServiceName, DEFAULT_RESCAN_FREQUENCY_MILLIS, noRecordsRescanFrequency,
-                dnsSrvRecordInitializer, clusterId, new DefaultDnsResolver());
+                dnsSrvRecordInitializer, clusterId, new DefaultDnsResolver(dnsClient));
     }
 }