JAVA-2564: Implement server session pooling behavior from specification

Author: jyemin

driver/src/main/com/mongodb/Mongo.java

@@ -99,7 +99,7 @@ public class Mongo {
 
     private final ConcurrentLinkedQueue<ServerCursorAndNamespace> orphanedCursors = new ConcurrentLinkedQueue<ServerCursorAndNamespace>();
     private final ExecutorService cursorCleaningService;
-    private final ServerSessionPool serverSessionPool = new ServerSessionPool();
+    private final ServerSessionPool serverSessionPool;
 
     /**
      * Creates a Mongo instance based on a (single) mongodb node (localhost, default port)
@@ -314,6 +314,7 @@ public Mongo(
 
     Mongo(final Cluster cluster, final MongoClientOptions options, final List<MongoCredential> credentialsList) {
         this.cluster = cluster;
+        this.serverSessionPool = new ServerSessionPool(cluster);
         this.options = options;
         this.readPreference = options.getReadPreference() != null ? options.getReadPreference() : primary();
         this.writeConcern = options.getWriteConcern() != null ? options.getWriteConcern() : WriteConcern.UNACKNOWLEDGED;
@@ -531,8 +532,8 @@ public void dropDatabase(final String dbName) {
      * databases obtained from it can no longer be used.
      */
     public void close() {
-        cluster.close();
         serverSessionPool.close();
+        cluster.close();
         if (cursorCleaningService != null) {
             cursorCleaningService.shutdownNow();
         }

driver/src/main/com/mongodb/ServerSessionPool.java

@@ -16,43 +16,145 @@
 
 package com.mongodb;
 
+import com.mongodb.connection.Cluster;
+import com.mongodb.connection.Connection;
 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 org.bson.BsonArray;
 import org.bson.BsonBinary;
 import org.bson.BsonDocument;
 import org.bson.BsonDocumentWriter;
 import org.bson.UuidRepresentation;
+import org.bson.codecs.BsonDocumentCodec;
 import org.bson.codecs.EncoderContext;
 import org.bson.codecs.UuidCodec;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.UUID;
 
+import static com.mongodb.assertions.Assertions.isTrue;
+import static java.util.concurrent.TimeUnit.MINUTES;
+
 class ServerSessionPool {
-    private final ConcurrentPool<ServerSession> serverSessionPool =
-            new ConcurrentPool<ServerSession>(Integer.MAX_VALUE, new ServerSessionItemFactory());
+
+    private static final int END_SESSIONS_BATCH_SIZE = 10000;
+
+    private final ConcurrentPool<ServerSessionImpl> serverSessionPool =
+            new ConcurrentPool<ServerSessionImpl>(Integer.MAX_VALUE, new ServerSessionItemFactory());
+    private final Cluster cluster;
+    private final ServerSessionPool.Clock clock;
+    private volatile boolean closing;
+    private volatile boolean closed;
+    private final List<BsonDocument> closedSessionIdentifiers = new ArrayList<BsonDocument>();
+
+    interface Clock {
+        long millis();
+    }
+
+    ServerSessionPool(final Cluster cluster) {
+        this(cluster, new Clock() {
+            @Override
+            public long millis() {
+                return System.currentTimeMillis();
+            }
+        });
+    }
+
+    ServerSessionPool(final Cluster cluster, final Clock clock) {
+        this.cluster = cluster;
+        this.clock = clock;
+    }
 
     ServerSession get() {
-        return serverSessionPool.get();
+        isTrue("server session pool is open", !closed);
+        ServerSessionImpl serverSession = serverSessionPool.get();
+        while (shouldPrune(serverSession)) {
+            serverSessionPool.release(serverSession, true);
+            serverSession = serverSessionPool.get();
+        }
+        return serverSession;
     }
 
     void release(final ServerSession serverSession) {
-        serverSessionPool.release(serverSession);
+        serverSessionPool.release((ServerSessionImpl) serverSession);
     }
 
     void close() {
-        serverSessionPool.close();
+        try {
+            closing = true;
+            serverSessionPool.close();
+            endClosedSessions();
+        } finally {
+            closed = true;
+        }
+    }
+
+    private void closeSession(final ServerSessionImpl serverSession) {
+        // only track closed sessions when pool is in the process of closing
+        if (!closing) {
+            return;
+        }
+
+        closedSessionIdentifiers.add(serverSession.getIdentifier());
+        if (closedSessionIdentifiers.size() == END_SESSIONS_BATCH_SIZE) {
+            endClosedSessions();
+        }
+    }
+
+    private void endClosedSessions() {
+        if (closedSessionIdentifiers.isEmpty()) {
+            return;
+        }
+
+        Connection connection = cluster.selectServer(new ReadPreferenceServerSelector(ReadPreference.primaryPreferred())).getConnection();
+        try {
+            connection.command("admin",
+                    new BsonDocument("endSessions", new BsonArray(closedSessionIdentifiers)),
+                    ReadPreference.primaryPreferred(), new NoOpFieldNameValidator(),
+                    new BsonDocumentCodec(), NoOpSessionContext.INSTANCE);
+        } catch (MongoException e) {
+            // ignore exceptions
+        } finally {
+            closedSessionIdentifiers.clear();
+            connection.release();
+        }
     }
 
-    private static class ServerSessionImpl implements ServerSession {
+    private boolean shouldPrune(final ServerSessionImpl serverSession) {
+        Integer logicalSessionTimeoutMinutes = cluster.getDescription().getLogicalSessionTimeoutMinutes();
+        if (logicalSessionTimeoutMinutes == null) {
+            return false;
+        }
+        long currentTimeMillis = clock.millis();
+        final long timeSinceLastUse = currentTimeMillis - serverSession.lastUsedAtMillis;
+        final long oneMinuteFromTimeout = MINUTES.toMillis(logicalSessionTimeoutMinutes - 1);
+        return timeSinceLastUse > oneMinuteFromTimeout;
+    }
+
+    final class ServerSessionImpl implements ServerSession {
         private final BsonDocument identifier;
         private int transactionNumber;
+        private volatile long lastUsedAtMillis = clock.millis();
 
         ServerSessionImpl(final BsonBinary identifier) {
             this.identifier = new BsonDocument("id", identifier);
         }
 
+        long getLastUsedAtMillis() {
+            return lastUsedAtMillis;
+        }
+
+        int getTransactionNumber() {
+            return transactionNumber;
+        }
+
         @Override
         public BsonDocument getIdentifier() {
+            lastUsedAtMillis = clock.millis();
             return identifier;
         }
 
@@ -62,19 +164,19 @@ public long advanceTransactionNumber() {
         }
     }
 
-    private static final class ServerSessionItemFactory implements ConcurrentPool.ItemFactory<ServerSession> {
+    private final class ServerSessionItemFactory implements ConcurrentPool.ItemFactory<ServerSessionImpl> {
         @Override
-        public ServerSession create(final boolean initialize) {
+        public ServerSessionImpl create(final boolean initialize) {
             return new ServerSessionImpl(createNewServerSessionIdentifier());
         }
 
         @Override
-        public void close(final ServerSession serverSession) {
-            // TODO: pruning
+        public void close(final ServerSessionImpl serverSession) {
+            closeSession(serverSession);
         }
 
         @Override
-        public Prune shouldPrune(final ServerSession serverSession) {
+        public Prune shouldPrune(final ServerSessionImpl serverSession) {
             return Prune.STOP;
         }
 

driver/src/test/functional/com/mongodb/MongoClientSessionSpecification.groovy

@@ -266,7 +266,6 @@ class MongoClientSessionSpecification extends FunctionalSpecification {
         then:
         def pingCommandStartedEvent = commandListener.events.get(0)
         !(pingCommandStartedEvent as CommandStartedEvent).command.containsKey('lsid')
-
         cleanup:
         client?.close()
     }