JAVA-2564: Prune server sessions on release

Author: jyemin

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

@@ -38,4 +38,11 @@ public interface ServerSession {
      * @return the next transaction number
      */
     long advanceTransactionNumber();
+
+    /**
+     * Whether the server session is closed.
+     *
+     * @return true if the session has been closed
+     */
+    boolean isClosed();
 }

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

@@ -81,6 +81,7 @@ ServerSession get() {
 
     void release(final ServerSession serverSession) {
         serverSessionPool.release((ServerSessionImpl) serverSession);
+        serverSessionPool.prune();
     }
 
     void close() {
@@ -94,6 +95,7 @@ void close() {
     }
 
     private void closeSession(final ServerSessionImpl serverSession) {
+        serverSession.close();
         // only track closed sessions when pool is in the process of closing
         if (!closing) {
             return;
@@ -139,11 +141,16 @@ final class ServerSessionImpl implements ServerSession {
         private final BsonDocument identifier;
         private int transactionNumber;
         private volatile long lastUsedAtMillis = clock.millis();
+        private volatile boolean closed;
 
         ServerSessionImpl(final BsonBinary identifier) {
             this.identifier = new BsonDocument("id", identifier);
         }
 
+        void close() {
+            closed = true;
+        }
+
         long getLastUsedAtMillis() {
             return lastUsedAtMillis;
         }
@@ -162,6 +169,11 @@ public BsonDocument getIdentifier() {
         public long advanceTransactionNumber() {
             return transactionNumber++;
         }
+
+        @Override
+        public boolean isClosed() {
+            return closed;
+        }
     }
 
     private final class ServerSessionItemFactory implements ConcurrentPool.ItemFactory<ServerSessionImpl> {
@@ -177,7 +189,7 @@ public void close(final ServerSessionImpl serverSession) {
 
         @Override
         public Prune shouldPrune(final ServerSessionImpl serverSession) {
-            return Prune.STOP;
+            return ServerSessionPool.this.shouldPrune(serverSession) ? Prune.YES : Prune.STOP;
         }
 
         private BsonBinary createNewServerSessionIdentifier() {

driver/src/test/unit/com/mongodb/ServerSessionPoolSpecification.groovy

@@ -108,33 +108,76 @@ class ServerSessionPoolSpecification extends Specification {
         session == pooledSession
     }
 
-    def 'should prune session'() {
+    def 'should prune sessions on release'() {
         given:
         def cluster = Mock(Cluster) {
             getDescription() >> connectedDescription
         }
         def clock = Stub(ServerSessionPool.Clock) {
-            millis() >>> [0,
-                          1,
-                          MINUTES.toMillis(29),
-                          MINUTES.toMillis(29) + 1]
+            millis() >>> [0, 0,                                // first get
+                          1, 1,                                // second get
+                          2, 2,                                // third get
+                          3,                                   // first release
+                          MINUTES.toMillis(29),       // second release
+                          MINUTES.toMillis(29) + 2,   // third release
+                          MINUTES.toMillis(29) + 2,
+                          MINUTES.toMillis(29) + 2
+            ]
         }
         def pool = new ServerSessionPool(cluster, clock)
-        def session = pool.get()
+        def sessionOne = pool.get()
+        def sessionTwo = pool.get()
+        def sessionThree = pool.get()
 
         when:
-        pool.release(session)
-        def newSession = pool.get()
+        pool.release(sessionOne)
 
         then:
-        session == newSession
+        !sessionOne.closed
+
+        when:
+        pool.release(sessionTwo)
+
+        then:
+        !sessionOne.closed
+        !sessionTwo.closed
+
+        when:
+        pool.release(sessionThree)
+
+        then:
+        sessionOne.closed
+        sessionTwo.closed
+        !sessionThree.closed
+        0 * cluster.selectServer(_)
+    }
+
+    def 'should prune sessions on get'() {
+        given:
+        def cluster = Mock(Cluster) {
+            getDescription() >> connectedDescription
+        }
+        def clock = Stub(ServerSessionPool.Clock) {
+            millis() >>> [0, 0,                               // first get
+                          0,                                  // first release
+                          MINUTES.toMillis(29) + 1,   // second get
+            ]
+        }
+        def pool = new ServerSessionPool(cluster, clock)
+        def sessionOne = pool.get()
+
+        when:
+        pool.release(sessionOne)
+
+        then:
+        !sessionOne.closed
 
         when:
-        pool.release(newSession)
-        newSession = pool.get()
+        def sessionTwo = pool.get()
 
         then:
-        session != newSession
+        sessionTwo != sessionOne
+        sessionOne.closed
         0 * cluster.selectServer(_)
     }