JAVA-2666: Allow an application to register a server selector that is applied after the suitable servers are selected and before servers within the latency window are selected

jyemin · jyemin · commit 8d661c030ecd · 2017-11-17T08:53:18.000-05:00
diff --git a/docs/reference/content/whats-new.md b/docs/reference/content/whats-new.md
@@ -11,27 +11,32 @@ title = "What's New"
 
 Key new features of the 3.6 Java driver release:
 
-### Compression
-
-The 3.6 release adds support for compression of messages to and from appropriately configured MongoDB servers:
-
-* [Compression Tutorial]({{<ref "driver/tutorials/compression.md">}})
-* [Compression Tutorial (Async)]({{<ref "driver-async/tutorials/compression.md">}})
-
 ### Change Stream support
 
 The 3.6 release adds support for [change streams](https://docs.mongodb.com/manual/operator/aggregation/changeStream).
 
 * [Change Stream Quick Start]({{<ref "driver/tutorials/change-streams.md">}}) 
 * [Change Stream Quick Start (Async)]({{<ref "driver-async/tutorials/change-streams.md">}})
 
+### Retryable writes
+
+The 3.6 release adds support for retryable writes using the `retryWrites` option in 
+[`MongoClientOptions`]({{<apiref "com/mongodb/MongoClientOptions">}}).
+
+### Compression
+
+The 3.6 release adds support for compression of messages to and from appropriately configured MongoDB servers:
+
+* [Compression Tutorial]({{<ref "driver/tutorials/compression.md">}})
+* [Compression Tutorial (Async)]({{<ref "driver-async/tutorials/compression.md">}})
+
 ### Causal consistency
 
 The 3.6 release adds support for causally consistency.
 
-### Retryable writes
+### Application-configured server selection
 
-The 3.6 release adds support for retryable writes using the `retryWrites` option in 
+The 3.6 release adds support for application-configured control over server selection, using the `serverSelector` option in
 [`MongoClientOptions`]({{<apiref "com/mongodb/MongoClientOptions">}}).
 
 ## What's New in 3.5
diff --git a/driver-core/src/main/com/mongodb/selector/ServerSelector.java b/driver-core/src/main/com/mongodb/selector/ServerSelector.java
@@ -23,7 +23,7 @@
 import java.util.List;
 
 /**
- * <p>An interface for selecting a server from a cluster according some preference.</p>
+ * <p>An interface for selecting a server from a cluster according to some preference.</p>
  *
  * <p>Implementations of this interface should ensure that their equals and hashCode methods compare equal preferences as equal, as users of
  * this interface may rely on that behavior to efficiently consolidate handling of multiple requests waiting on a server that can satisfy
diff --git a/driver/src/main/com/mongodb/Mongo.java b/driver/src/main/com/mongodb/Mongo.java
@@ -44,12 +44,14 @@
 import com.mongodb.operation.ListDatabasesOperation;
 import com.mongodb.operation.ReadOperation;
 import com.mongodb.operation.WriteOperation;
+import com.mongodb.selector.CompositeServerSelector;
 import com.mongodb.selector.LatencyMinimizingServerSelector;
 import com.mongodb.selector.ServerSelector;
 import com.mongodb.session.ClientSession;
 import org.bson.BsonBoolean;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -757,7 +759,7 @@ private static ClusterSettings getClusterSettings(final List<ServerAddress> seed
                 .mode(clusterConnectionMode)
                 .requiredReplicaSetName(options.getRequiredReplicaSetName())
                 .serverSelectionTimeout(options.getServerSelectionTimeout(), MILLISECONDS)
-                .serverSelector(new LatencyMinimizingServerSelector(options.getLocalThreshold(), MILLISECONDS))
+                .serverSelector(getServerSelector(options))
                 .description(options.getDescription())
                 .maxWaitQueueSize(options.getConnectionPoolSettings().getMaxWaitQueueSize());
         for (ClusterListener clusterListener: options.getClusterListeners()) {
@@ -766,6 +768,17 @@ private static ClusterSettings getClusterSettings(final List<ServerAddress> seed
         return builder.build();
     }
 
+    private static ServerSelector getServerSelector(final MongoClientOptions options) {
+        LatencyMinimizingServerSelector latencyMinimizingServerSelector =
+                new LatencyMinimizingServerSelector(options.getLocalThreshold(), MILLISECONDS);
+
+        if (options.getServerSelector() == null) {
+            return latencyMinimizingServerSelector;
+        }
+
+        return new CompositeServerSelector(Arrays.asList(options.getServerSelector(), latencyMinimizingServerSelector));
+    }
+
     Cluster getCluster() {
         return cluster;
     }
diff --git a/driver/src/main/com/mongodb/MongoClientOptions.java b/driver/src/main/com/mongodb/MongoClientOptions.java
@@ -27,6 +27,7 @@
 import com.mongodb.event.ConnectionPoolListener;
 import com.mongodb.event.ServerListener;
 import com.mongodb.event.ServerMonitorListener;
+import com.mongodb.selector.ServerSelector;
 import org.bson.codecs.configuration.CodecRegistry;
 
 import javax.net.SocketFactory;
@@ -65,6 +66,7 @@ public class MongoClientOptions {
     private final boolean retryWrites;
     private final ReadConcern readConcern;
     private final CodecRegistry codecRegistry;
+    private final ServerSelector serverSelector;
 
     private final int minConnectionsPerHost;
     private final int maxConnectionsPerHost;
@@ -121,6 +123,7 @@ private MongoClientOptions(final Builder builder) {
         retryWrites = builder.retryWrites;
         readConcern = builder.readConcern;
         codecRegistry = builder.codecRegistry;
+        serverSelector = builder.serverSelector;
         sslEnabled = builder.sslEnabled;
         sslInvalidHostNameAllowed = builder.sslInvalidHostNameAllowed;
         sslContext = builder.sslContext;
@@ -564,6 +567,34 @@ public CodecRegistry getCodecRegistry() {
         return codecRegistry;
     }
 
+    /**
+     * Gets the server selector.
+     *
+     * <p>The server selector augments the normal server selection rules applied by the driver when determining
+     * which server to send an operation to.  At the point that it's called by the driver, the
+     * {@link com.mongodb.connection.ClusterDescription} which is passed to it contains a list of
+     * {@link com.mongodb.connection.ServerDescription} instances which satisfy either the configured {@link ReadPreference} for any
+     * read operation or ones that can take writes (e.g. a standalone, mongos, or replica set primary).</p>
+     * <p>The server selector can then filter the {@code ServerDescription} list using whatever criteria that is required by the
+     * application.</p>
+     * <p>After this selector executes, two additional selectors are applied by the driver:</p>
+     * <ul>
+     * <li>select from within the latency window</li>
+     * <li>select a random server from those remaining</li>
+     * </ul>
+     * <p>To skip the latency window selector, an application can:</p>
+     * <ul>
+     * <li>configure the local threshold to a sufficiently high value so that it doesn't exclude any servers</li>
+     * <li>return a list containing a single server from this selector (which will also make the random member selector a no-op)</li>
+     * </ul>
+     *
+     * @return the server selector, which may be null
+     * @since 3.6
+     */
+    public ServerSelector getServerSelector() {
+        return serverSelector;
+    }
+
     /**
      * Gets the list of added {@code ClusterListener}. The default is an empty list.
      *
@@ -799,6 +830,9 @@ public boolean equals(final Object o) {
         if (!codecRegistry.equals(that.codecRegistry)) {
             return false;
         }
+        if (serverSelector != null ? !serverSelector.equals(that.serverSelector) : that.serverSelector != null) {
+            return false;
+        }
         if (!clusterListeners.equals(that.clusterListeners)) {
             return false;
         }
@@ -828,6 +862,7 @@ public int hashCode() {
         result = 31 * result + (retryWrites ? 1 : 0);
         result = 31 * result + (readConcern != null ? readConcern.hashCode() : 0);
         result = 31 * result + codecRegistry.hashCode();
+        result = 31 * result + (serverSelector != null ? serverSelector.hashCode() : 0);
         result = 31 * result + clusterListeners.hashCode();
         result = 31 * result + commandListeners.hashCode();
         result = 31 * result + minConnectionsPerHost;
@@ -869,6 +904,7 @@ public String toString() {
                + ", retryWrites=" + retryWrites
                + ", readConcern=" + readConcern
                + ", codecRegistry=" + codecRegistry
+               + ", serverSelector=" + serverSelector
                + ", clusterListeners=" + clusterListeners
                + ", commandListeners=" + commandListeners
                + ", minConnectionsPerHost=" + minConnectionsPerHost
@@ -923,7 +959,7 @@ public static class Builder {
         private boolean retryWrites = false;
         private ReadConcern readConcern = ReadConcern.DEFAULT;
         private CodecRegistry codecRegistry = MongoClient.getDefaultCodecRegistry();
-
+        private ServerSelector serverSelector;
         private int minConnectionsPerHost;
         private int maxConnectionsPerHost = 100;
         private int threadsAllowedToBlockForConnectionMultiplier = 5;
@@ -987,6 +1023,7 @@ public Builder(final MongoClientOptions options) {
             retryWrites = options.getRetryWrites();
             readConcern = options.getReadConcern();
             codecRegistry = options.getCodecRegistry();
+            serverSelector = options.getServerSelector();
             sslEnabled = options.isSslEnabled();
             sslInvalidHostNameAllowed = options.isSslInvalidHostNameAllowed();
             sslContext = options.getSslContext();
@@ -1313,6 +1350,19 @@ public Builder codecRegistry(final CodecRegistry codecRegistry) {
             return this;
         }
 
+        /**
+         * Sets a server selector that augments the normal server selection rules applied by the driver when determining
+         * which server to send an operation to.  See {@link #getServerSelector()} for further details.
+         *
+         * @param serverSelector the server selector
+         * @return this
+         * @since 3.6
+         * @see #getServerSelector()
+         */
+        public Builder serverSelector(final ServerSelector serverSelector) {
+            this.serverSelector = serverSelector;
+            return this;
+        }
         /**
          * Adds the given command listener.
          *
diff --git a/driver/src/test/functional/com/mongodb/MongoClientsSpecification.groovy b/driver/src/test/functional/com/mongodb/MongoClientsSpecification.groovy
@@ -18,11 +18,19 @@
 package com.mongodb
 
 import com.mongodb.client.MongoDriverInformation
+import com.mongodb.connection.ClusterDescription
+import com.mongodb.connection.ServerDescription
+import com.mongodb.event.CommandFailedEvent
+import com.mongodb.event.CommandListener
+import com.mongodb.event.CommandStartedEvent
+import com.mongodb.event.CommandSucceededEvent
 import org.bson.Document
 import spock.lang.IgnoreIf
 
+import static com.mongodb.ClusterFixture.isDiscoverableReplicaSet
 import static com.mongodb.ClusterFixture.isStandalone
 import static com.mongodb.ClusterFixture.serverVersionAtLeast
+import static com.mongodb.Fixture.getDefaultDatabaseName
 import static com.mongodb.Fixture.getMongoClientURI
 
 class MongoClientsSpecification extends FunctionalSpecification {
@@ -52,4 +60,48 @@ class MongoClientsSpecification extends FunctionalSpecification {
         profileCollection?.drop()
         client?.close()
     }
+
+    @IgnoreIf({ !isDiscoverableReplicaSet() })
+    def 'should use server selector from MongoClientOptions'() {
+        given:
+        def expectedWinner
+        def actualWinningAddresses = [] as Set
+        def optionsBuilder = MongoClientOptions.builder()
+        // select the suitable server with the highest port number
+                .serverSelector { ClusterDescription clusterDescription ->
+            for (ServerDescription cur : clusterDescription.getServerDescriptions()) {
+                if (expectedWinner == null || cur.address.port > expectedWinner.address.port) {
+                    expectedWinner = cur
+                }
+            }
+            expectedWinner == null ? [] : [expectedWinner]
+        }.addCommandListener(new CommandListener() {
+            // record each address actually used
+            @Override
+            void commandStarted(final CommandStartedEvent event) {
+                actualWinningAddresses.add(event.connectionDescription.connectionId.serverId.address)
+            }
+
+            @Override
+            void commandSucceeded(final CommandSucceededEvent event) {
+            }
+
+            @Override
+            void commandFailed(final CommandFailedEvent event) {
+            }
+        })
+
+        def client = new MongoClient(getMongoClientURI(optionsBuilder))
+        def collection = client.getDatabase(getDefaultDatabaseName()).getCollection(getCollectionName())
+                .withReadPreference(ReadPreference.nearest())
+
+        when:
+        for (int i = 0; i < 10; i++) {
+            collection.count()
+        }
+
+        then:
+        actualWinningAddresses.size() == 1
+        actualWinningAddresses.contains(expectedWinner.address)
+    }
 }
diff --git a/driver/src/test/unit/com/mongodb/MongoClientOptionsSpecification.groovy b/driver/src/test/unit/com/mongodb/MongoClientOptionsSpecification.groovy
@@ -25,6 +25,7 @@ import com.mongodb.event.CommandListener
 import com.mongodb.event.ConnectionPoolListener
 import com.mongodb.event.ServerListener
 import com.mongodb.event.ServerMonitorListener
+import com.mongodb.selector.ServerSelector
 import spock.lang.IgnoreIf
 import spock.lang.Specification
 
@@ -56,6 +57,7 @@ class MongoClientOptionsSpecification extends Specification {
         options.getConnectionsPerHost() == 100
         options.getConnectTimeout() == 10000
         options.getReadPreference() == ReadPreference.primary()
+        options.getServerSelector() == null;
         options.getThreadsAllowedToBlockForConnectionMultiplier() == 5
         options.isSocketKeepAlive()
         !options.isSslEnabled()
@@ -157,6 +159,7 @@ class MongoClientOptionsSpecification extends Specification {
         given:
         def encoderFactory = new MyDBEncoderFactory()
         def socketFactory = SSLSocketFactory.getDefault()
+        def serverSelector = Mock(ServerSelector)
         def options = MongoClientOptions.builder()
                                         .description('test')
                                         .applicationName('appName')
@@ -167,6 +170,7 @@ class MongoClientOptionsSpecification extends Specification {
                                         .connectionsPerHost(500)
                                         .connectTimeout(100)
                                         .socketTimeout(700)
+                                        .serverSelector(serverSelector)
                                         .serverSelectionTimeout(150)
                                         .maxWaitTime(200)
                                         .maxConnectionIdleTime(300)
@@ -194,6 +198,7 @@ class MongoClientOptionsSpecification extends Specification {
         options.getApplicationName() == 'appName'
         options.getReadPreference() == ReadPreference.secondary()
         options.getWriteConcern() == WriteConcern.JOURNALED
+        options.getServerSelector() == serverSelector
         options.getRetryWrites()
         options.getServerSelectionTimeout() == 150
         options.getMaxWaitTime() == 200
@@ -656,8 +661,8 @@ class MongoClientOptionsSpecification extends Specification {
                         'description', 'heartbeatConnectTimeout', 'heartbeatFrequency', 'heartbeatSocketTimeout', 'localThreshold',
                         'maxConnectionIdleTime', 'maxConnectionLifeTime', 'maxConnectionsPerHost', 'maxWaitTime', 'minConnectionsPerHost',
                         'minHeartbeatFrequency', 'readConcern', 'readPreference', 'requiredReplicaSetName', 'retryWrites',
-                        'serverListeners', 'serverMonitorListeners', 'serverSelectionTimeout', 'socketFactory', 'socketKeepAlive',
-                        'socketTimeout', 'sslContext', 'sslEnabled', 'sslInvalidHostNameAllowed',
+                        'serverListeners', 'serverMonitorListeners', 'serverSelectionTimeout', 'serverSelector', 'socketFactory',
+                        'socketKeepAlive', 'socketTimeout', 'sslContext', 'sslEnabled', 'sslInvalidHostNameAllowed',
                         'threadsAllowedToBlockForConnectionMultiplier', 'writeConcern']
 
         then:

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@`
`23`	`23`	`import java.util.List;`
`24`	`24`
`25`	`25`	`/**`
`26`		`- * <p>An interface for selecting a server from a cluster according some preference.</p>`
	`26`	`+ * <p>An interface for selecting a server from a cluster according to some preference.</p>`
`27`	`27`	`*`
`28`	`28`	`* <p>Implementations of this interface should ensure that their equals and hashCode methods compare equal preferences as equal, as users of`
`29`	`29`	`* this interface may rely on that behavior to efficiently consolidate handling of multiple requests waiting on a server that can satisfy`