Release connection between empty getMore responses (#1928)

- Extract getMoreLoop to loop on empty batches while releasing the connection between each getMore, preventing pool starvation from idle tailable cursors (e.g., change streams)
- Add test to verify connection is released and re-acquired between consecutive empty getMores

JAVA-6142

Author: vbabanin

driver-core/src/main/com/mongodb/internal/operation/AsyncCommandBatchCursor.java

@@ -121,7 +121,7 @@ public void next(final SingleResultCallback<List<T>> callback) {
                 commandCursorResult = withEmptyResults(commandCursorResult);
                 funcCallback.onResult(batchResults, null);
             } else {
-                getMore(localServerCursor, funcCallback);
+                getMoreLoop(localServerCursor, funcCallback);
             }
         }, callback);
     }
@@ -183,10 +183,10 @@ void checkTimeoutModeAndResetTimeoutContextIfIteration() {
 
     private void getMore(final ServerCursor cursor, final SingleResultCallback<List<T>> callback) {
         resourceManager.executeWithConnection((connection, wrappedCallback) ->
-                getMoreLoop(assertNotNull(connection), cursor, wrappedCallback), callback);
+                executeGetMoreCommand(assertNotNull(connection), cursor, wrappedCallback), callback);
     }
 
-    private void getMoreLoop(final AsyncConnection connection, final ServerCursor serverCursor,
+    private void executeGetMoreCommand(final AsyncConnection connection, final ServerCursor serverCursor,
             final SingleResultCallback<List<T>> callback) {
         connection.commandAsync(namespace.getDatabaseName(),
                 getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace, batchSize, comment),
@@ -206,19 +206,22 @@ private void getMoreLoop(final AsyncConnection connection, final ServerCursor se
                             connection.getDescription().getServerAddress(), NEXT_BATCH, assertNotNull(commandResult));
                     ServerCursor nextServerCursor = commandCursorResult.getServerCursor();
                     resourceManager.setServerCursor(nextServerCursor);
-                    List<T> nextBatch = commandCursorResult.getResults();
-                    if (nextServerCursor == null || !nextBatch.isEmpty()) {
-                        commandCursorResult = withEmptyResults(commandCursorResult);
-                        callback.onResult(nextBatch, null);
-                        return;
-                    }
-
-                    if (!resourceManager.operable()) {
-                        callback.onResult(emptyList(), null);
-                        return;
-                    }
+                    callback.onResult(commandCursorResult.getResults(), null);
+        });
+    }
 
-                    getMoreLoop(connection, nextServerCursor, callback);
+    private void getMoreLoop(final ServerCursor localServerCursor, final SingleResultCallback<List<T>> funcCallback) {
+        getMore(localServerCursor, (nextBatch, t) -> {
+            if (t != null) {
+                funcCallback.onResult(null, t);
+            } else if (resourceManager.getServerCursor() == null || (nextBatch != null && !nextBatch.isEmpty())) {
+                commandCursorResult = withEmptyResults(commandCursorResult);
+                funcCallback.onResult(nextBatch, null);
+            } else if (!resourceManager.operable()) {
+                funcCallback.onResult(emptyList(), null);
+            } else {
+                getMoreLoop(assertNotNull(resourceManager.getServerCursor()), funcCallback);
+            }
         });
     }
 

driver-core/src/test/functional/com/mongodb/internal/operation/AsyncCommandBatchCursorTest.java

@@ -47,10 +47,14 @@
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static com.mongodb.internal.operation.OperationUnitSpecification.getMaxWireVersionForServerVersion;
 import static com.mongodb.internal.thread.InterruptionUtil.interruptAndCreateMongoInterruptedException;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.argThat;
@@ -119,17 +123,17 @@ void shouldSkipKillsCursorsCommandWhenNetworkErrorOccurs() {
             return null;
         }).when(mockConnection).commandAsync(eq(NAMESPACE.getDatabaseName()), any(), any(), any(), any(), any(), any());
         when(serverDescription.getType()).thenReturn(ServerType.LOAD_BALANCER);
-        AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0);
 
         //when
-        commandBatchCursor.next((result, t) -> {
-            Assertions.assertNull(result);
-            Assertions.assertNotNull(t);
-            Assertions.assertEquals(MongoSocketException.class, t.getClass());
-        });
+        try (AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0)) {
+            commandBatchCursor.next((result, t) -> {
+                Assertions.assertNull(result);
+                Assertions.assertNotNull(t);
+                Assertions.assertEquals(MongoSocketException.class, t.getClass());
+            });
+        }
 
         //then
-        commandBatchCursor.close();
         verify(mockConnection, times(1)).commandAsync(eq(NAMESPACE.getDatabaseName()), any(), any(), any(), any(), any(), any());
     }
 
@@ -144,17 +148,14 @@ void shouldNotSkipKillsCursorsCommandWhenTimeoutExceptionDoesNotHaveNetworkError
         }).when(mockConnection).commandAsync(eq(NAMESPACE.getDatabaseName()), any(), any(), any(), any(), any(), any());
         when(serverDescription.getType()).thenReturn(ServerType.LOAD_BALANCER);
 
-        AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0);
-
         //when
-        commandBatchCursor.next((result, t) -> {
-            Assertions.assertNull(result);
-            Assertions.assertNotNull(t);
-            Assertions.assertEquals(MongoOperationTimeoutException.class, t.getClass());
-        });
-
-        commandBatchCursor.close();
-
+        try (AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0)) {
+            commandBatchCursor.next((result, t) -> {
+                Assertions.assertNull(result);
+                Assertions.assertNotNull(t);
+                Assertions.assertEquals(MongoOperationTimeoutException.class, t.getClass());
+            });
+        }
 
         //then
         verify(mockConnection, times(2)).commandAsync(any(),
@@ -175,16 +176,14 @@ void shouldSkipKillsCursorsCommandWhenTimeoutExceptionHaveNetworkErrorCause() {
         }).when(mockConnection).commandAsync(eq(NAMESPACE.getDatabaseName()), any(), any(), any(), any(), any(), any());
         when(serverDescription.getType()).thenReturn(ServerType.LOAD_BALANCER);
 
-        AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0);
-
         //when
-        commandBatchCursor.next((result, t) -> {
-            Assertions.assertNull(result);
-            Assertions.assertNotNull(t);
-            Assertions.assertEquals(MongoOperationTimeoutException.class, t.getClass());
-        });
-
-        commandBatchCursor.close();
+        try (AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0)) {
+            commandBatchCursor.next((result, t) -> {
+                Assertions.assertNull(result);
+                Assertions.assertNotNull(t);
+                Assertions.assertEquals(MongoOperationTimeoutException.class, t.getClass());
+            });
+        }
 
         //then
         verify(mockConnection, times(1)).commandAsync(any(),
@@ -203,8 +202,6 @@ void closeShouldResetTimeoutContextToDefaultMaxTime() {
         try (AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(maxTimeMS)) {
             // verify that the `maxTimeMS` override was applied
             timeoutContext.runMaxTimeMS(remainingMillis -> assertTrue(remainingMillis <= maxTimeMS));
-        } catch (Exception e) {
-            throw new RuntimeException(e);
         }
         timeoutContext.runMaxTimeMS(remainingMillis -> {
             // verify that the `maxTimeMS` override was reset
@@ -236,8 +233,6 @@ void closeShouldNotResetOriginalTimeout(final boolean disableTimeoutResetWhenClo
                 Thread.sleep(thirdOfTimeout.toMillis());
                 return null;
             });
-        } catch (Exception e) {
-            throw new RuntimeException(e);
         }
         verify(mockConnection, times(1)).release();
         // at this point at least (2 * thirdOfTimeout) have passed
@@ -252,6 +247,34 @@ void closeShouldNotResetOriginalTimeout(final boolean disableTimeoutResetWhenClo
                 Assertions::fail);
     }
 
+    @Test
+    void shouldReleaseConnectionBetweenEmptyGetMoreResponses() {
+        AtomicInteger callCount = new AtomicInteger();
+        doAnswer(invocation -> {
+            SingleResultCallback<BsonDocument> cb = invocation.getArgument(6);
+            cb.onResult(new BsonDocument("cursor",
+                    new BsonDocument("ns", new BsonString(NAMESPACE.getFullName()))
+                            .append("id", new BsonInt64(callCount.incrementAndGet() < 3 ? 1 : 0))
+                            .append("nextBatch", new BsonArrayWrapper<>(new BsonArray()))), null);
+            return null;
+        }).when(mockConnection).commandAsync(eq(NAMESPACE.getDatabaseName()),
+                argThat(doc -> doc.containsKey("getMore")), any(), any(), any(), any(), any());
+
+        when(serverDescription.getType()).thenReturn(ServerType.STANDALONE);
+        try (AsyncCommandBatchCursor<Document> commandBatchCursor = createBatchCursor(0)) {
+            commandBatchCursor.next((result, t) -> {
+                assertNotNull(result);
+                assertTrue(result.isEmpty());
+                assertNull(t);
+            });
+        }
+
+        // 2 empty-batch getMores + 1 exhausted getMore = 3 getMores, but the 3rd
+        // exhausts the cursor (id=0), which makes the cursor break the loop and return an empty result.
+        verify(mockConnection, times(3)).release();
+        verify(connectionSource, times(3)).getConnection(any());
+        assertEquals(3, callCount.get());
+    }
 
     private AsyncCommandBatchCursor<Document> createBatchCursor(final long maxTimeMS) {
         return new AsyncCommandBatchCursor<Document>(

driver-core/src/test/unit/com/mongodb/internal/operation/AsyncCommandBatchCursorSpecification.groovy

@@ -162,8 +162,9 @@ class AsyncCommandBatchCursorSpecification extends Specification {
     def 'should handle getMore when there are empty results but there is a cursor'() {
         given:
         def initialConnection = referenceCountedAsyncConnection()
-        def connection = referenceCountedAsyncConnection()
-        def connectionSource = getAsyncConnectionSource(connection)
+        def connectionA = referenceCountedAsyncConnection()
+        def connectionB = referenceCountedAsyncConnection()
+        def connectionSource = getAsyncConnectionSource(connectionA, connectionB)
 
         when:
         def firstBatch = createCommandResult([], CURSOR_ID)
@@ -172,14 +173,15 @@ class AsyncCommandBatchCursorSpecification extends Specification {
         def batch = nextBatch(cursor)
 
         then:
-        1 * connection.commandAsync(*_) >> {
-            connection.getCount() == 1
+        1 * connectionA.commandAsync(*_) >> {
+            connectionA.getCount() == 1
             connectionSource.getCount() == 1
             it.last().onResult(response, null)
         }
 
-        1 * connection.commandAsync(*_) >> {
-            connection.getCount() == 1
+        then:
+        1 * connectionB.commandAsync(*_) >> {
+            connectionB.getCount() == 1
             connectionSource.getCount() == 1
             it.last().onResult(response2, null)
         }
@@ -191,7 +193,10 @@ class AsyncCommandBatchCursorSpecification extends Specification {
         cursor.close()
 
         then:
-        0 * connection._
+        0 * connectionA._
+        0 * connectionB._
+        connectionA.getCount() == 0
+        connectionB.getCount() == 0
         initialConnection.getCount() == 0
         connectionSource.getCount() == 0