JAVA-2859: Don't block selecting server when closing server session pool

Author: jyemin

driver-core/src/main/com/mongodb/internal/session/ServerSessionPool.java

@@ -19,12 +19,15 @@
 import com.mongodb.MongoException;
 import com.mongodb.ReadPreference;
 import com.mongodb.connection.Cluster;
+import com.mongodb.connection.ClusterDescription;
 import com.mongodb.connection.Connection;
+import com.mongodb.connection.ServerDescription;
 import com.mongodb.internal.connection.ConcurrentPool;
 import com.mongodb.internal.connection.ConcurrentPool.Prune;
 import com.mongodb.internal.connection.NoOpSessionContext;
 import com.mongodb.internal.validator.NoOpFieldNameValidator;
 import com.mongodb.selector.ReadPreferenceServerSelector;
+import com.mongodb.selector.ServerSelector;
 import com.mongodb.session.ServerSession;
 import org.bson.BsonArray;
 import org.bson.BsonBinary;
@@ -36,6 +39,7 @@
 import org.bson.codecs.UuidCodec;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.UUID;
 
@@ -118,7 +122,23 @@ private void endClosedSessions() {
             return;
         }
 
-        Connection connection = cluster.selectServer(new ReadPreferenceServerSelector(ReadPreference.primaryPreferred())).getConnection();
+        final List<ServerDescription> primaryPreferred = new ReadPreferenceServerSelector(ReadPreference.primaryPreferred())
+                .select(cluster.getCurrentDescription());
+        if (primaryPreferred.isEmpty()) {
+            return;
+        }
+
+        Connection connection = cluster.selectServer(new ServerSelector() {
+            @Override
+            public List<ServerDescription> select(final ClusterDescription clusterDescription) {
+                for (ServerDescription cur : clusterDescription.getServerDescriptions()) {
+                    if (cur.getAddress().equals(primaryPreferred.get(0).getAddress())) {
+                        return Collections.singletonList(cur);
+                    }
+                }
+                return Collections.emptyList();
+            }
+        }).getConnection();
         try {
             connection.command("admin",
                     new BsonDocument("endSessions", new BsonArray(closedSessionIdentifiers)), new NoOpFieldNameValidator(),

driver-core/src/test/unit/com/mongodb/internal/session/ServerSessionPoolSpecification.groovy

@@ -26,7 +26,6 @@ import com.mongodb.connection.ServerDescription
 import com.mongodb.connection.ServerSettings
 import com.mongodb.internal.connection.NoOpSessionContext
 import com.mongodb.internal.validator.NoOpFieldNameValidator
-import com.mongodb.selector.ReadPreferenceServerSelector
 import org.bson.BsonArray
 import org.bson.BsonBinarySubType
 import org.bson.BsonDocument
@@ -252,14 +251,14 @@ class ServerSessionPoolSpecification extends Specification {
 
         then:
         // first batch is the first 10K sessions, final batch is the last one
-        1 * cluster.selectServer { (it as ReadPreferenceServerSelector).readPreference == primaryPreferred() }  >> server
+        1 * cluster.selectServer(_)  >> server
         1 * connection.command('admin',
                 new BsonDocument('endSessions', new BsonArray(sessions*.getIdentifier())),
                 { it instanceof NoOpFieldNameValidator }, primaryPreferred(),
                 { it instanceof BsonDocumentCodec }, NoOpSessionContext.INSTANCE) >> new BsonDocument()
         1 * connection.release()
 
-        1 * cluster.selectServer { (it as ReadPreferenceServerSelector).readPreference == primaryPreferred() }  >> server
+        1 * cluster.selectServer(_)  >> server
         1 * connection.command('admin',
                 new BsonDocument('endSessions', new BsonArray([oneOverBatchSizeSession.getIdentifier()])),
                 { it instanceof NoOpFieldNameValidator }, primaryPreferred(),