Remove limit and count from CommandBatchCursors and clarify initial connection in tests

Author: rozza

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

@@ -224,16 +224,16 @@ private CommandReadTransformer<BsonDocument, CommandBatchCursor<T>> transformer(
         return (result, source, connection) -> {
             // TODO (CSOT) JAVA-4058
             long maxAwaitTimeMS = timeoutSettings.getMaxAwaitTimeMS();
-            return new CommandBatchCursor<>(result, 0, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder,
-                    comment, source, connection);
+            return new CommandBatchCursor<>(result, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder,
+                                            comment, source, connection);
         };
     }
 
     private CommandReadTransformerAsync<BsonDocument, AsyncBatchCursor<T>> asyncTransformer() {
         return (result, source, connection) -> {
             // TODO (CSOT) JAVA-4058
             long maxAwaitTimeMS = timeoutSettings.getMaxAwaitTimeMS();
-            return new AsyncCommandBatchCursor<>(result, 0, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder,
+            return new AsyncCommandBatchCursor<>(result, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder,
                     comment, source, connection);
         };
     }

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

@@ -44,7 +44,6 @@
 
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.assertions.Assertions.assertTrue;
@@ -63,7 +62,6 @@
 class AsyncCommandBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T> {
 
     private final MongoNamespace namespace;
-    private final int limit;
     private final long maxTimeMS;
     private final Decoder<T> decoder;
     @Nullable
@@ -72,21 +70,19 @@ class AsyncCommandBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T>
     private final boolean firstBatchEmpty;
     private final ResourceManager resourceManager;
     private final CommandCursorResult<T> commandCursorResult;
-    private final AtomicInteger count = new AtomicInteger();
     private final AtomicBoolean processedInitial = new AtomicBoolean();
     private int batchSize;
 
     AsyncCommandBatchCursor(
             final BsonDocument commandCursorDocument,
-            final int limit, final int batchSize, final long maxTimeMS,
+            final int batchSize, final long maxTimeMS,
             final Decoder<T> decoder,
             @Nullable final BsonValue comment,
             final AsyncConnectionSource connectionSource,
             final AsyncConnection connection) {
         ConnectionDescription connectionDescription = connection.getDescription();
         this.commandCursorResult = toCommandCursorResult(connectionDescription.getServerAddress(), FIRST_BATCH, commandCursorDocument);
         this.namespace = commandCursorResult.getNamespace();
-        this.limit = limit;
         this.batchSize = batchSize;
         this.maxTimeMS = maxTimeMS;
         this.decoder = notNull("decoder", decoder);
@@ -95,17 +91,11 @@ class AsyncCommandBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T>
         this.firstBatchEmpty = commandCursorResult.getResults().isEmpty();
 
         AsyncConnection connectionToPin = null;
-        boolean releaseServerAndResources = false;
-        if (limitReached()) {
-            releaseServerAndResources = true;
-        } else if (connectionDescription.getServerType() == ServerType.LOAD_BALANCER) {
+        if (connectionDescription.getServerType() == ServerType.LOAD_BALANCER) {
             connectionToPin = connection;
         }
 
         resourceManager = new ResourceManager(namespace, connectionSource, connectionToPin, commandCursorResult.getServerCursor());
-        if (releaseServerAndResources) {
-            resourceManager.releaseServerAndClientResources(connection);
-        }
     }
 
     @Override
@@ -124,15 +114,12 @@ public void next(final SingleResultCallback<List<T>> callback) {
             }
 
             if (serverCursorIsNull || !batchResults.isEmpty()) {
-                if (serverCursorIsNull) {
-                    close();
-                }
                 funcCallback.onResult(batchResults, null);
             } else {
                 getMore(localServerCursor, funcCallback);
             }
         }).get((r, t) -> {
-            if (limitReached()) {
+            if (resourceManager.getServerCursor() == null) {
                 close();
             }
             callback.onResult(r, t);
@@ -196,8 +183,7 @@ private void getMore(final ServerCursor cursor, final SingleResultCallback<List<
     private void getMoreLoop(final AsyncConnection connection, final ServerCursor serverCursor,
             final SingleResultCallback<List<T>> callback) {
         connection.commandAsync(namespace.getDatabaseName(),
-                getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace,
-                        limit, batchSize, count.get(), maxTimeMS, comment),
+                getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace, batchSize, maxTimeMS, comment),
                 NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(),
                 CommandResultDocumentCodec.create(decoder, NEXT_BATCH),
                 assertNotNull(resourceManager.getConnectionSource()).getOperationContext(),
@@ -214,35 +200,29 @@ private void getMoreLoop(final AsyncConnection connection, final ServerCursor se
                             connection.getDescription().getServerAddress(), NEXT_BATCH, assertNotNull(commandResult));
                     resourceManager.setServerCursor(commandCursorResult.getServerCursor());
 
+                    List<T> nextBatch = commandCursorResult.getResults();
+                    if (commandCursorResult.getServerCursor() == null || !nextBatch.isEmpty()) {
+                        callback.onResult(nextBatch, null);
+                        return;
+                    }
+
                     if (!resourceManager.operable()) {
-                        // The cursor was closed
-                        resourceManager.releaseServerAndClientResources(connection);
                         callback.onResult(emptyList(), null);
                         return;
                     }
 
-                    List<T> nextBatch = commandCursorResult.getResults();
-                    if (nextBatch.isEmpty() && commandCursorResult.getServerCursor() != null) {
-                        getMoreLoop(connection, commandCursorResult.getServerCursor(), callback);
-                    } else {
-                        callback.onResult(nextBatch, null);
-                    }
-                });
+                    getMoreLoop(connection, commandCursorResult.getServerCursor(), callback);
+        });
     }
 
     private CommandCursorResult<T> toCommandCursorResult(final ServerAddress serverAddress, final String fieldNameContainingBatch,
             final BsonDocument commandCursorDocument) {
         CommandCursorResult<T> commandCursorResult = new CommandCursorResult<>(serverAddress, fieldNameContainingBatch,
                 commandCursorDocument);
         logCommandCursorResult(commandCursorResult);
-        this.count.addAndGet(commandCursorResult.getResults().size());
         return commandCursorResult;
     }
 
-    private boolean limitReached() {
-        return Math.abs(limit) != 0 && count.get() >= Math.abs(limit);
-    }
-
     @ThreadSafe
     private static final class ResourceManager extends CursorResourceManager<AsyncConnectionSource, AsyncConnection> {
 
@@ -282,7 +262,7 @@ void doClose() {
             if (getServerCursor() != null) {
                 getConnection((connection, t) -> {
                     if (connection != null) {
-                        releaseServerAndClientResources(connection, (r, t1) -> connection.release());
+                        releaseServerAndClientResources(connection);
                     } else {
                         unsetServerCursor();
                         releaseClientResources();
@@ -320,10 +300,6 @@ private void getConnection(final SingleResultCallback<AsyncConnection> callback)
         }
 
         private void releaseServerAndClientResources(final AsyncConnection connection) {
-            releaseServerAndClientResources(connection, (r, t) -> { /* Do nothing */ });
-        }
-
-        private void releaseServerAndClientResources(final AsyncConnection connection, final SingleResultCallback<Void> callback) {
             AsyncCallbackSupplier<Void> callbackSupplier = funcCallback -> {
                 ServerCursor localServerCursor = getServerCursor();
                 if (localServerCursor != null) {
@@ -333,7 +309,7 @@ private void releaseServerAndClientResources(final AsyncConnection connection, f
             callbackSupplier.whenComplete(() -> {
                 unsetServerCursor();
                 releaseClientResources();
-            }).get(callback);
+            }).whenComplete(connection::release).get((r,t) -> { /* do nothing */ });
         }
 
         private void killServerCursor(final MongoNamespace namespace, final ServerCursor localServerCursor,

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

@@ -339,8 +339,7 @@ static <T> CommandReadTransformerAsync<BsonDocument, AsyncBatchCursor<T>> asyncS
 
     static <T> AsyncBatchCursor<T> cursorDocumentToAsyncBatchCursor(final BsonDocument cursorDocument, final Decoder<T> decoder,
             final BsonValue comment, final AsyncConnectionSource source, final AsyncConnection connection, final int batchSize) {
-        return new AsyncCommandBatchCursor<>(cursorDocument, 0, batchSize, 0, decoder,
-                comment, source, connection);
+        return new AsyncCommandBatchCursor<>(cursorDocument, batchSize, 0, decoder, comment, source, connection);
     }
 
     static <T> SingleResultCallback<T> releasingCallback(final SingleResultCallback<T> wrapped, final AsyncConnection connection) {

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

@@ -59,7 +59,6 @@
 class CommandBatchCursor<T> implements AggregateResponseBatchCursor<T> {
 
     private final MongoNamespace namespace;
-    private final int limit;
     private final long maxTimeMS;
     private final Decoder<T> decoder;
     @Nullable
@@ -70,21 +69,19 @@ class CommandBatchCursor<T> implements AggregateResponseBatchCursor<T> {
 
     private int batchSize;
     private CommandCursorResult<T> commandCursorResult;
-    private int count = 0;
     @Nullable
     private List<T> nextBatch;
 
     CommandBatchCursor(
             final BsonDocument commandCursorDocument,
-            final int limit, final int batchSize, final long maxTimeMS,
+            final int batchSize, final long maxTimeMS,
             final Decoder<T> decoder,
             @Nullable final BsonValue comment,
             final ConnectionSource connectionSource,
             final Connection connection) {
         ConnectionDescription connectionDescription = connection.getDescription();
         this.commandCursorResult = toCommandCursorResult(connectionDescription.getServerAddress(), FIRST_BATCH, commandCursorDocument);
         this.namespace = commandCursorResult.getNamespace();
-        this.limit = limit;
         this.batchSize = batchSize;
         this.maxTimeMS = maxTimeMS;
         this.decoder = notNull("decoder", decoder);
@@ -93,17 +90,11 @@ class CommandBatchCursor<T> implements AggregateResponseBatchCursor<T> {
         this.firstBatchEmpty = commandCursorResult.getResults().isEmpty();
 
         Connection connectionToPin = null;
-        boolean releaseServerAndResources = false;
-        if (limitReached()) {
-            releaseServerAndResources = true;
-        } else if (connectionDescription.getServerType() == ServerType.LOAD_BALANCER) {
+        if (connectionDescription.getServerType() == ServerType.LOAD_BALANCER) {
             connectionToPin = connection;
         }
 
         resourceManager = new ResourceManager(namespace, connectionSource, connectionToPin, commandCursorResult.getServerCursor());
-        if (releaseServerAndResources) {
-            resourceManager.releaseServerAndClientResources(connection);
-        }
     }
 
     @Override
@@ -116,10 +107,6 @@ private boolean doHasNext() {
             return true;
         }
 
-        if (limitReached()) {
-            return false;
-        }
-
         while (resourceManager.getServerCursor() != null) {
             getMore();
             if (!resourceManager.operable()) {
@@ -195,10 +182,6 @@ private boolean tryHasNext() {
             return true;
         }
 
-        if (limitReached()) {
-            return false;
-        }
-
         if (resourceManager.getServerCursor() != null) {
             getMore();
         }
@@ -252,8 +235,8 @@ private void getMore() {
                 this.commandCursorResult = toCommandCursorResult(connection.getDescription().getServerAddress(), NEXT_BATCH,
                         assertNotNull(
                             connection.command(namespace.getDatabaseName(),
-                                 getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace,
-                                     limit, batchSize, count, maxTimeMS, comment),
+                                 getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace, batchSize,
+                                         maxTimeMS, comment),
                                  NO_OP_FIELD_NAME_VALIDATOR,
                                  ReadPreference.primary(),
                                  CommandResultDocumentCodec.create(decoder, NEXT_BATCH),
@@ -263,9 +246,6 @@ private void getMore() {
                 throw translateCommandException(e, serverCursor);
             }
             resourceManager.setServerCursor(nextServerCursor);
-            if (limitReached() || !resourceManager.operable()) {
-                resourceManager.releaseServerAndClientResources(connection);
-            }
         });
     }
 
@@ -275,14 +255,9 @@ private CommandCursorResult<T> toCommandCursorResult(final ServerAddress serverA
                 commandCursorDocument);
         logCommandCursorResult(commandCursorResult);
         this.nextBatch = commandCursorResult.getResults().isEmpty() ? null : commandCursorResult.getResults();
-        this.count += commandCursorResult.getResults().size();
         return commandCursorResult;
     }
 
-    private boolean limitReached() {
-        return Math.abs(limit) != 0 && count >= Math.abs(limit);
-    }
-
     @ThreadSafe
     private static final class ResourceManager extends CursorResourceManager<ConnectionSource, Connection> {
 
@@ -362,14 +337,6 @@ private Connection getConnection() {
             }
         }
 
-        private void releaseServerAndClientResources(final Connection connection) {
-            try {
-                releaseServerResources(assertNotNull(connection));
-            } finally {
-                releaseClientResources();
-            }
-        }
-
         private void releaseServerResources(final Connection connection) {
             try {
                 ServerCursor localServerCursor = getServerCursor();

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

@@ -51,14 +51,13 @@ final class CommandBatchCursorHelper {
             + "supported";
 
     static BsonDocument getMoreCommandDocument(
-            final long cursorId, final ConnectionDescription connectionDescription, final MongoNamespace namespace, final int limit,
-            final int batchSize, final int count, final long maxTimeMS, @Nullable final BsonValue comment) {
+            final long cursorId, final ConnectionDescription connectionDescription, final MongoNamespace namespace, final int batchSize,
+            final long maxTimeMS, @Nullable final BsonValue comment) {
         BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursorId))
                 .append("collection", new BsonString(namespace.getCollectionName()));
 
-        int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, batchSize, count));
-        if (batchSizeForGetMoreCommand != 0) {
-            document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));
+        if (batchSize != 0) {
+            document.append("batchSize", new BsonInt32(batchSize));
         }
         if (maxTimeMS != 0) {
             document.append("maxTimeMS", new BsonInt64(maxTimeMS));

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

@@ -456,7 +456,7 @@ private boolean isAwaitData() {
 
     private CommandReadTransformer<BsonDocument, CommandBatchCursor<T>> transformer() {
         return (result, source, connection) ->
-                new CommandBatchCursor<>(result, limit, batchSize, getMaxTimeForCursor(), decoder, comment, source, connection);
+                new CommandBatchCursor<>(result, batchSize, getMaxTimeForCursor(), decoder, comment, source, connection);
     }
 
     // TODO - CSOT JAVA-4058
@@ -466,6 +466,6 @@ private long getMaxTimeForCursor() {
 
     private CommandReadTransformerAsync<BsonDocument, AsyncBatchCursor<T>> asyncTransformer() {
         return (result, source, connection) ->
-            new AsyncCommandBatchCursor<>(result, limit, batchSize, getMaxTimeForCursor(), decoder, comment, source, connection);
+            new AsyncCommandBatchCursor<>(result, batchSize, getMaxTimeForCursor(), decoder, comment, source, connection);
     }
 }

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

@@ -331,7 +331,7 @@ static <T> CommandReadTransformer<BsonDocument, BatchCursor<T>> singleBatchCurso
 
     static <T> BatchCursor<T> cursorDocumentToBatchCursor(final BsonDocument cursorDocument, final Decoder<T> decoder,
             final BsonValue comment, final ConnectionSource source, final Connection connection, final int batchSize) {
-        return new CommandBatchCursor<>(cursorDocument, 0, batchSize, 0, decoder, comment, source, connection);
+        return new CommandBatchCursor<>(cursorDocument, batchSize, 0, decoder, comment, source, connection);
     }
 
     private SyncOperationHelper() {

driver-core/src/test/functional/com/mongodb/ClusterFixture.java

@@ -775,7 +775,7 @@ public static int getReferenceCountAfterTimeout(final ReferenceCounted reference
         int count = referenceCounted.getCount();
         while (count > target) {
             try {
-                if (System.currentTimeMillis() > startTime + 5000) {
+                if (System.currentTimeMillis() > startTime + TIMEOUT_DURATION.toMillis()) {
                     return count;
                 }
                 Thread.sleep(10);