diff --git a/.evergreen/run-load-balancer-tests.sh b/.evergreen/run-load-balancer-tests.sh index 871e9c2f003..5297097971c 100755 --- a/.evergreen/run-load-balancer-tests.sh +++ b/.evergreen/run-load-balancer-tests.sh @@ -79,7 +79,7 @@ echo $second -Dorg.mongodb.test.uri=${SINGLE_MONGOS_LB_URI} \ -Dorg.mongodb.test.multi.mongos.uri=${MULTI_MONGOS_LB_URI} \ ${GRADLE_EXTRA_VARS} --stacktrace --info --continue driver-core:test \ - --tests QueryBatchCursorFunctionalSpecification + --tests CommandBatchCursorSpecification third=$? echo $third diff --git a/config/codenarc/codenarc.xml b/config/codenarc/codenarc.xml index 2d11b03296a..2bab2315e97 100644 --- a/config/codenarc/codenarc.xml +++ b/config/codenarc/codenarc.xml @@ -34,12 +34,6 @@ - - - - - - diff --git a/driver-core/src/main/com/mongodb/internal/async/AsyncAggregateResponseBatchCursor.java b/driver-core/src/main/com/mongodb/internal/async/AsyncAggregateResponseBatchCursor.java index 2e9da84550c..ccfc9f7a956 100644 --- a/driver-core/src/main/com/mongodb/internal/async/AsyncAggregateResponseBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/async/AsyncAggregateResponseBatchCursor.java @@ -16,6 +16,7 @@ package com.mongodb.internal.async; +import com.mongodb.lang.Nullable; import org.bson.BsonDocument; import org.bson.BsonTimestamp; @@ -25,8 +26,10 @@ *

This class is not part of the public API and may be removed or changed at any time

*/ public interface AsyncAggregateResponseBatchCursor extends AsyncBatchCursor { + @Nullable BsonDocument getPostBatchResumeToken(); + @Nullable BsonTimestamp getOperationTime(); boolean isFirstBatchEmpty(); diff --git a/driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java b/driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java index 14cc3faa71b..89260ac7b52 100644 --- a/driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java @@ -16,6 +16,8 @@ package com.mongodb.internal.async; +import com.mongodb.internal.operation.BatchCursor; + import java.io.Closeable; import java.util.List; @@ -28,9 +30,9 @@ */ public interface AsyncBatchCursor extends Closeable { /** - * Returns the next batch of results. A tailable cursor will block until another batch exists. After the last batch, the next call - * to this method will execute the callback with a null result to indicate that there are no more batches available and the cursor - * has been closed. + * Returns the next batch of results. A tailable cursor will block until another batch exists. + * Unlike the {@link BatchCursor} this method will automatically mark the cursor as closed when there are no more expected results. + * Care should be taken to check {@link #isClosed()} between calls. * * @param callback callback to receive the next batch of results * @throws java.util.NoSuchElementException if no next batch exists diff --git a/driver-core/src/main/com/mongodb/internal/operation/AggregateOperationImpl.java b/driver-core/src/main/com/mongodb/internal/operation/AggregateOperationImpl.java index 4379845bdd1..d70c7cbac66 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/AggregateOperationImpl.java +++ b/driver-core/src/main/com/mongodb/internal/operation/AggregateOperationImpl.java @@ -18,13 +18,11 @@ import com.mongodb.MongoNamespace; import com.mongodb.client.model.Collation; -import com.mongodb.connection.ConnectionDescription; import com.mongodb.internal.async.AsyncBatchCursor; import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; import com.mongodb.internal.client.model.AggregationLevel; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.session.SessionContext; import com.mongodb.lang.Nullable; import org.bson.BsonArray; @@ -40,7 +38,6 @@ import java.util.List; import java.util.concurrent.TimeUnit; -import static com.mongodb.assertions.Assertions.assertNotNull; import static com.mongodb.assertions.Assertions.isTrueArgument; import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; @@ -48,7 +45,6 @@ import static com.mongodb.internal.operation.AsyncOperationHelper.executeRetryableReadAsync; import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator; import static com.mongodb.internal.operation.OperationHelper.LOGGER; -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult; import static com.mongodb.internal.operation.OperationReadConcernHelper.appendReadConcernToCommand; import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; import static com.mongodb.internal.operation.SyncOperationHelper.executeRetryableRead; @@ -239,25 +235,16 @@ BsonDocument getCommand(final SessionContext sessionContext, final int maxWireVe return commandDocument; } - private QueryResult createQueryResult(final BsonDocument result, final ConnectionDescription description) { - assertNotNull(result); - return cursorDocumentToQueryResult(result.getDocument(CURSOR), description.getServerAddress()); - } - - private CommandReadTransformer> transformer() { - return (result, source, connection) -> { - QueryResult queryResult = createQueryResult(result, connection.getDescription()); - return new QueryBatchCursor<>(queryResult, 0, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder, comment, - source, connection, result); - }; + private CommandReadTransformer> transformer() { + return (result, source, connection) -> + new CommandBatchCursor<>(connection.getDescription().getServerAddress(), result, 0, batchSize != null ? batchSize : 0, + maxAwaitTimeMS, decoder, comment, source, connection); } private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> { - QueryResult queryResult = createQueryResult(result, connection.getDescription()); - return new AsyncQueryBatchCursor<>(queryResult, 0, batchSize != null ? batchSize : 0, maxAwaitTimeMS, decoder, - comment, source, connection, result); - }; + return (result, source, connection) -> + new AsyncCommandBatchCursor<>(connection.getDescription().getServerAddress(), result, 0, batchSize != null ? batchSize : 0, + maxAwaitTimeMS, decoder, comment, source, connection); } interface AggregateTarget { diff --git a/driver-core/src/main/com/mongodb/internal/operation/AggregateResponseBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/AggregateResponseBatchCursor.java index 5ec7d00bb26..e12a2249123 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/AggregateResponseBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/AggregateResponseBatchCursor.java @@ -17,6 +17,7 @@ package com.mongodb.internal.operation; import com.mongodb.annotations.NotThreadSafe; +import com.mongodb.lang.Nullable; import org.bson.BsonDocument; import org.bson.BsonTimestamp; @@ -27,8 +28,10 @@ */ @NotThreadSafe public interface AggregateResponseBatchCursor extends BatchCursor { + @Nullable BsonDocument getPostBatchResumeToken(); + @Nullable BsonTimestamp getOperationTime(); boolean isFirstBatchEmpty(); diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncChangeStreamBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncChangeStreamBatchCursor.java index 9ccd2f17b0a..1d271107a9a 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/AsyncChangeStreamBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/AsyncChangeStreamBatchCursor.java @@ -60,8 +60,7 @@ final class AsyncChangeStreamBatchCursor implements AsyncAggregateResponseBat final int maxWireVersion) { this.changeStreamOperation = changeStreamOperation; this.wrapped = new AtomicReference<>(assertNotNull(wrapped)); - this.binding = binding; - binding.retain(); + this.binding = binding.retain(); this.resumeToken = resumeToken; this.maxWireVersion = maxWireVersion; isClosed = new AtomicBoolean(); @@ -164,8 +163,8 @@ public int getMaxWireVersion() { return maxWireVersion; } - private void cachePostBatchResumeToken(final AsyncAggregateResponseBatchCursor queryBatchCursor) { - BsonDocument resumeToken = queryBatchCursor.getPostBatchResumeToken(); + private void cachePostBatchResumeToken(final AsyncAggregateResponseBatchCursor commandBatchCursor) { + BsonDocument resumeToken = commandBatchCursor.getPostBatchResumeToken(); if (resumeToken != null) { this.resumeToken = resumeToken; } diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncCommandBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncCommandBatchCursor.java new file mode 100644 index 00000000000..d10a1f9705e --- /dev/null +++ b/driver-core/src/main/com/mongodb/internal/operation/AsyncCommandBatchCursor.java @@ -0,0 +1,339 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.MongoCommandException; +import com.mongodb.MongoException; +import com.mongodb.MongoNamespace; +import com.mongodb.MongoSocketException; +import com.mongodb.ReadPreference; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; +import com.mongodb.annotations.ThreadSafe; +import com.mongodb.connection.ServerType; +import com.mongodb.internal.VisibleForTesting; +import com.mongodb.internal.async.AsyncAggregateResponseBatchCursor; +import com.mongodb.internal.async.SingleResultCallback; +import com.mongodb.internal.async.function.AsyncCallbackSupplier; +import com.mongodb.internal.binding.AsyncConnectionSource; +import com.mongodb.internal.connection.AsyncConnection; +import com.mongodb.internal.connection.Connection; +import com.mongodb.lang.Nullable; +import org.bson.BsonDocument; +import org.bson.BsonTimestamp; +import org.bson.BsonValue; +import org.bson.codecs.BsonDocumentCodec; +import org.bson.codecs.Decoder; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +import static com.mongodb.assertions.Assertions.assertNotNull; +import static com.mongodb.assertions.Assertions.assertTrue; +import static com.mongodb.assertions.Assertions.isTrueArgument; +import static com.mongodb.assertions.Assertions.notNull; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.FIRST_BATCH; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CLOSED_AS_CURSOR; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CONCURRENT_OPERATION; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.NEXT_BATCH; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.NO_OP_FIELD_NAME_VALIDATOR; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.getMoreCommandDocument; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.translateCommandException; +import static com.mongodb.internal.operation.OperationHelper.LOGGER; +import static java.lang.String.format; +import static java.util.Collections.emptyList; + +class AsyncCommandBatchCursor implements AsyncAggregateResponseBatchCursor { + + private final MongoNamespace namespace; + private final int limit; + private final Decoder decoder; + private final long maxTimeMS; + @Nullable + private final BsonValue comment; + + private final AtomicInteger count = new AtomicInteger(); + private final AtomicBoolean processedInitial = new AtomicBoolean(); + private final int maxWireVersion; + private final boolean firstBatchEmpty; + private final ResourceManager resourceManager; + private volatile CommandCursorResult commandCursorResult; + private int batchSize; + + AsyncCommandBatchCursor( + final ServerAddress serverAddress, + final BsonDocument commandCursorDocument, + final int limit, final int batchSize, final long maxTimeMS, + final Decoder decoder, + @Nullable final BsonValue comment, + final AsyncConnectionSource connectionSource, + final AsyncConnection connection) { + isTrueArgument("maxTimeMS >= 0", maxTimeMS >= 0); + CommandCursorResult commandCursor = initFromCommandCursorDocument(serverAddress, FIRST_BATCH, commandCursorDocument); + this.namespace = commandCursor.getNamespace(); + this.limit = limit; + this.batchSize = batchSize; + this.maxTimeMS = maxTimeMS; + this.decoder = notNull("decoder", decoder); + this.comment = comment; + this.maxWireVersion = connection.getDescription().getMaxWireVersion(); + this.firstBatchEmpty = commandCursor.getResults().isEmpty(); + + AsyncConnection connectionToPin = null; + boolean releaseServerAndResources = false; + if (limitReached()) { + releaseServerAndResources = true; + } else if (connectionSource.getServerDescription().getType() == ServerType.LOAD_BALANCER) { + connectionToPin = connection; + } + + resourceManager = new ResourceManager(namespace, connectionSource, connectionToPin, commandCursor.getServerCursor()); + if (releaseServerAndResources) { + resourceManager.releaseServerAndClientResources(connection); + } + } + + @Override + public void next(final SingleResultCallback> callback) { + if (isClosed()) { + callback.onResult(null, new MongoException(MESSAGE_IF_CLOSED_AS_CURSOR)); + return; + } else if (!resourceManager.tryStartOperation()) { + callback.onResult(null, new MongoException(MESSAGE_IF_CONCURRENT_OPERATION)); + return; + } + + ServerCursor localServerCursor = resourceManager.getServerCursor(); + boolean cursorClosed = localServerCursor == null; + List batchResults = emptyList(); + if (!processedInitial.getAndSet(true) && !firstBatchEmpty) { + batchResults = commandCursorResult.getResults(); + } + + if (cursorClosed || !batchResults.isEmpty()) { + resourceManager.endOperation(); + if (cursorClosed) { + close(); + } + callback.onResult(batchResults, null); + } else { + getMore(localServerCursor, callback); + } + } + + @Override + public void setBatchSize(final int batchSize) { + this.batchSize = batchSize; + } + + @Override + public int getBatchSize() { + return batchSize; + } + + @Override + public boolean isClosed() { + return !resourceManager.operable(); + } + + @Override + public void close() { + resourceManager.close(); + } + + @Override + public BsonDocument getPostBatchResumeToken() { + return commandCursorResult.getPostBatchResumeToken(); + } + + @Override + public BsonTimestamp getOperationTime() { + return commandCursorResult.getOperationTime(); + } + + @Override + public boolean isFirstBatchEmpty() { + return firstBatchEmpty; + } + + @Override + public int getMaxWireVersion() { + return maxWireVersion; + } + + @Nullable + @VisibleForTesting(otherwise = VisibleForTesting.AccessModifier.PRIVATE) + ServerCursor getServerCursor() { + return resourceManager.getServerCursor(); + } + + private void getMore(final ServerCursor cursor, final SingleResultCallback> callback) { + resourceManager.executeWithConnection(connection -> getMore(connection, cursor, callback), callback); + } + + private void getMore(final AsyncConnection connection, final ServerCursor cursor, final SingleResultCallback> callback) { + connection.commandAsync(namespace.getDatabaseName(), + getMoreCommandDocument(cursor.getId(), connection.getDescription(), namespace, + limit, batchSize, count.get(), maxTimeMS, comment), + NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), + CommandResultDocumentCodec.create(decoder, NEXT_BATCH), assertNotNull(resourceManager.getConnectionSource()), + (commandResult, t) -> { + if (t != null) { + Throwable translatedException = + t instanceof MongoCommandException + ? translateCommandException((MongoCommandException) t, cursor) + : t; + + connection.release(); + resourceManager.endOperation(); + callback.onResult(null, translatedException); + return; + } + + CommandCursorResult commandCursor = + initFromCommandCursorDocument(connection.getDescription().getServerAddress(), NEXT_BATCH, assertNotNull(commandResult)); + resourceManager.setServerCursor(commandCursor.getServerCursor()); + + if (!resourceManager.operable()) { + resourceManager.releaseServerAndClientResources(connection); + connection.release(); + callback.onResult(emptyList(), null); + return; + } + + if (commandCursor.getResults().isEmpty() && commandCursor.getServerCursor() != null) { + connection.release(); + getMore(commandCursor.getServerCursor(), callback); + } else { + resourceManager.endOperation(); + if (limitReached()) { + resourceManager.releaseServerAndClientResources(connection); + close(); + } + connection.release(); + callback.onResult(commandCursor.getResults(), null); + } + }); + } + + private CommandCursorResult initFromCommandCursorDocument( + final ServerAddress serverAddress, + final String fieldNameContainingBatch, + final BsonDocument commandCursorDocument) { + CommandCursorResult cursorResult = new CommandCursorResult<>(serverAddress, fieldNameContainingBatch, commandCursorDocument); + count.addAndGet(cursorResult.getResults().size()); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug(format("Received batch of %d documents with cursorId %d from server %s", cursorResult.getResults().size(), + cursorResult.getCursorId(), cursorResult.getServerAddress())); + } + this.commandCursorResult = cursorResult; + return cursorResult; + } + + private boolean limitReached() { + return Math.abs(limit) != 0 && count.get() >= Math.abs(limit); + } + + @ThreadSafe + private static final class ResourceManager extends CursorResourceManager { + + ResourceManager( + final MongoNamespace namespace, + final AsyncConnectionSource connectionSource, + @Nullable final AsyncConnection connectionToPin, + @Nullable final ServerCursor serverCursor) { + super(namespace, connectionSource, connectionToPin, serverCursor); + } + + @Override + void markAsPinned(final AsyncConnection connectionToPin, final Connection.PinningMode pinningMode) { + connectionToPin.markAsPinned(pinningMode); + } + + @Override + void executeWithConnection(final Consumer action) { + throw new UnsupportedOperationException(); + } + + @Override + void executeWithConnection(final Consumer action, final SingleResultCallback callback) { + assertTrue(getState() != State.IDLE); + AsyncConnection pinnedConnection = getPinnedConnection(); + if (pinnedConnection != null) { + executeWithConnection(assertNotNull(pinnedConnection).retain(), null, action, callback); + } else { + assertNotNull(getConnectionSource()).getConnection((conn, t) -> executeWithConnection(conn, t, action, callback)); + } + } + + void executeWithConnection( + @Nullable final AsyncConnection connection, + @Nullable final Throwable t, + final Consumer function, + final SingleResultCallback callback) { + assertTrue(connection != null || t != null); + if (t != null) { + callback.onResult(null, t); + } else { + AsyncCallbackSupplier curriedFunction = c -> function.accept(connection); + curriedFunction.whenComplete(connection::release).get((result, error) -> { + if (error instanceof MongoSocketException) { + onCorruptedConnection(connection); + } + connection.release(); + callback.onResult(result, error); + }); + } + } + + @Override + void doClose() { + if (isSkipReleasingServerResourcesOnClose()) { + unsetServerCursor(); + } + + if (getServerCursor() != null) { + executeWithConnection(c -> { + releaseServerResources(c); + c.release(); + }, (r, t) -> { + // guarantee that regardless of exceptions, `serverCursor` is null and client resources are released + unsetServerCursor(); + releaseClientResources(); + }); + } else { + releaseClientResources(); + } + } + + @Override + void killServerCursor(final MongoNamespace namespace, final ServerCursor serverCursor, final AsyncConnection connection) { + connection.retain(); + AsyncConnectionSource connectionSource = assertNotNull(getConnectionSource()).retain(); + connection.commandAsync(namespace.getDatabaseName(), getKillCursorsCommand(namespace, serverCursor), + NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), + connectionSource, (r, t) -> { + connectionSource.release(); + connection.release(); + releaseClientResources(); + }); + } + } + +} diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java index 21b10cdff08..b670083be78 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java +++ b/driver-core/src/main/com/mongodb/internal/operation/AsyncOperationHelper.java @@ -18,9 +18,7 @@ import com.mongodb.Function; import com.mongodb.MongoException; -import com.mongodb.MongoNamespace; import com.mongodb.ReadPreference; -import com.mongodb.ServerAddress; import com.mongodb.assertions.Assertions; import com.mongodb.internal.async.AsyncBatchCursor; import com.mongodb.internal.async.SingleResultCallback; @@ -35,7 +33,6 @@ import com.mongodb.internal.binding.ReferenceCounted; import com.mongodb.internal.connection.AsyncConnection; import com.mongodb.internal.connection.OperationContext; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.operation.retry.AttachmentKeys; import com.mongodb.internal.validator.NoOpFieldNameValidator; import com.mongodb.lang.Nullable; @@ -56,7 +53,6 @@ import static com.mongodb.internal.operation.CommandOperationHelper.isRetryWritesEnabled; import static com.mongodb.internal.operation.CommandOperationHelper.logRetryExecute; import static com.mongodb.internal.operation.CommandOperationHelper.transformWriteException; -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult; import static com.mongodb.internal.operation.WriteConcernHelper.throwOnWriteConcernError; final class AsyncOperationHelper { @@ -309,15 +305,15 @@ static CommandWriteTransformerAsync writeConcernErrorTransfo }; } - static AsyncBatchCursor createEmptyAsyncBatchCursor(final MongoNamespace namespace, final ServerAddress serverAddress) { - return new AsyncSingleBatchQueryCursor<>(new QueryResult<>(namespace, Collections.emptyList(), 0L, serverAddress)); + static CommandReadTransformerAsync> asyncSingleBatchCursorTransformer(final String resultsFieldName) { + return (result, source, connection) -> + new AsyncSingleBatchCursor<>(BsonDocumentWrapperHelper.toList(result, resultsFieldName), 0); } static AsyncBatchCursor cursorDocumentToAsyncBatchCursor(final BsonDocument cursorDocument, final Decoder decoder, final BsonValue comment, final AsyncConnectionSource source, final AsyncConnection connection, final int batchSize) { - return new AsyncQueryBatchCursor<>(cursorDocumentToQueryResult(cursorDocument, - source.getServerDescription().getAddress()), - 0, batchSize, 0, decoder, comment, source, connection, cursorDocument); + return new AsyncCommandBatchCursor<>(connection.getDescription().getServerAddress(), cursorDocument, + 0, batchSize, 0, decoder, comment, source, connection); } static SingleResultCallback releasingCallback(final SingleResultCallback wrapped, final AsyncConnection connection) { diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncQueryBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncQueryBatchCursor.java deleted file mode 100644 index be079abec36..00000000000 --- a/driver-core/src/main/com/mongodb/internal/operation/AsyncQueryBatchCursor.java +++ /dev/null @@ -1,428 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation; - -import com.mongodb.MongoCommandException; -import com.mongodb.MongoException; -import com.mongodb.MongoNamespace; -import com.mongodb.ReadPreference; -import com.mongodb.ServerCursor; -import com.mongodb.connection.ConnectionDescription; -import com.mongodb.connection.ServerType; -import com.mongodb.internal.async.AsyncAggregateResponseBatchCursor; -import com.mongodb.internal.async.SingleResultCallback; -import com.mongodb.internal.binding.AsyncConnectionSource; -import com.mongodb.internal.connection.AsyncConnection; -import com.mongodb.internal.connection.Connection; -import com.mongodb.internal.connection.QueryResult; -import com.mongodb.internal.diagnostics.logging.Logger; -import com.mongodb.internal.diagnostics.logging.Loggers; -import com.mongodb.internal.validator.NoOpFieldNameValidator; -import com.mongodb.lang.Nullable; -import org.bson.BsonArray; -import org.bson.BsonDocument; -import org.bson.BsonInt32; -import org.bson.BsonInt64; -import org.bson.BsonString; -import org.bson.BsonTimestamp; -import org.bson.BsonValue; -import org.bson.FieldNameValidator; -import org.bson.codecs.BsonDocumentCodec; -import org.bson.codecs.Decoder; - -import java.util.List; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.ReentrantLock; - -import static com.mongodb.assertions.Assertions.assertFalse; -import static com.mongodb.assertions.Assertions.assertNotNull; -import static com.mongodb.assertions.Assertions.isTrueArgument; -import static com.mongodb.assertions.Assertions.notNull; -import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; -import static com.mongodb.internal.operation.CursorHelper.getNumberToReturn; -import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; -import static com.mongodb.internal.operation.SyncOperationHelper.getMoreCursorDocumentToQueryResult; -import static com.mongodb.internal.operation.QueryHelper.translateCommandException; -import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionFourDotFour; -import static java.lang.String.format; -import static java.util.Collections.singletonList; - -class AsyncQueryBatchCursor implements AsyncAggregateResponseBatchCursor { - private static final Logger LOGGER = Loggers.getLogger("operation"); - private static final FieldNameValidator NO_OP_FIELD_NAME_VALIDATOR = new NoOpFieldNameValidator(); - private static final String CURSOR = "cursor"; - private static final String POST_BATCH_RESUME_TOKEN = "postBatchResumeToken"; - private static final String OPERATION_TIME = "operationTime"; - - private final MongoNamespace namespace; - private final int limit; - private final Decoder decoder; - private final long maxTimeMS; - private volatile AsyncConnectionSource connectionSource; - private volatile AsyncConnection pinnedConnection; - private final AtomicReference cursor; - private volatile QueryResult firstBatch; - private volatile int batchSize; - private final AtomicInteger count = new AtomicInteger(); - private volatile BsonDocument postBatchResumeToken; - private final BsonTimestamp operationTime; - private final BsonValue comment; - private final boolean firstBatchEmpty; - private final int maxWireVersion; - - private final Lock lock = new ReentrantLock(); - /* protected by `lock` */ - private boolean isOperationInProgress = false; - private boolean isClosed = false; - /* protected by `lock` */ - private volatile boolean isClosePending = false; - - AsyncQueryBatchCursor(final QueryResult firstBatch, final int limit, final int batchSize, final long maxTimeMS, - final Decoder decoder, final BsonValue comment, final AsyncConnectionSource connectionSource, - final AsyncConnection connection) { - this(firstBatch, limit, batchSize, maxTimeMS, decoder, comment, connectionSource, connection, null); - } - - AsyncQueryBatchCursor(final QueryResult firstBatch, final int limit, final int batchSize, final long maxTimeMS, - final Decoder decoder, final BsonValue comment, final AsyncConnectionSource connectionSource, - @Nullable final AsyncConnection connection, @Nullable final BsonDocument result) { - isTrueArgument("maxTimeMS >= 0", maxTimeMS >= 0); - this.maxTimeMS = maxTimeMS; - this.namespace = firstBatch.getNamespace(); - this.firstBatch = firstBatch; - this.limit = limit; - this.batchSize = batchSize; - this.decoder = decoder; - this.comment = comment; - this.cursor = new AtomicReference<>(firstBatch.getCursor()); - this.count.addAndGet(firstBatch.getResults().size()); - if (result != null) { - this.operationTime = result.getTimestamp(OPERATION_TIME, null); - this.postBatchResumeToken = getPostBatchResumeTokenFromResponse(result); - } else { - this.operationTime = null; - } - - firstBatchEmpty = firstBatch.getResults().isEmpty(); - if (cursor.get() != null) { - this.connectionSource = notNull("connectionSource", connectionSource).retain(); - assertNotNull(connection); - if (limitReached()) { - killCursor(connection); - } else { - if (connectionSource.getServerDescription().getType() == ServerType.LOAD_BALANCER) { - this.pinnedConnection = connection.retain(); - this.pinnedConnection.markAsPinned(Connection.PinningMode.CURSOR); - } - } - } - this.maxWireVersion = connection == null ? 0 : connection.getDescription().getMaxWireVersion(); - logQueryResult(firstBatch); - } - - /** - * {@inheritDoc} - *

- * From the perspective of the code external to this class, this method is idempotent as required by its specification. - * However, if this method sets {@link #isClosePending}, - * then it must be called by {@code this} again to release resources. - * This behavior does not violate externally observable idempotence because this method is allowed to release resources "eventually". - */ - @Override - public void close() { - boolean doClose = false; - - lock.lock(); - try { - if (isOperationInProgress) { - isClosePending = true; - } else if (!isClosed) { - isClosed = true; - isClosePending = false; - doClose = true; - } - } finally { - lock.unlock(); - } - - if (doClose) { - killCursorOnClose(); - } - } - - @Override - public void next(final SingleResultCallback> callback) { - if (isClosed()) { - callback.onResult(null, new MongoException("next() called after the cursor was closed.")); - } else if (firstBatch != null && (!firstBatch.getResults().isEmpty())) { - // May be empty for a tailable cursor - List results = firstBatch.getResults(); - firstBatch = null; - if (getServerCursor() == null) { - close(); - } - callback.onResult(results, null); - } else { - ServerCursor localCursor = getServerCursor(); - if (localCursor == null) { - close(); - callback.onResult(null, null); - } else { - lock.lock(); - try { - if (isClosed()) { - callback.onResult(null, new MongoException("next() called after the cursor was closed.")); - return; - } - isOperationInProgress = true; - } finally { - lock.unlock(); - } - getMore(localCursor, callback); - } - } - } - - @Override - public void setBatchSize(final int batchSize) { - assertFalse(isClosed()); - this.batchSize = batchSize; - } - - @Override - public int getBatchSize() { - assertFalse(isClosed()); - return batchSize; - } - - @Override - public boolean isClosed() { - lock.lock(); - try { - return isClosed || isClosePending; - } finally { - lock.unlock(); - } - } - - @Override - public BsonDocument getPostBatchResumeToken() { - return postBatchResumeToken; - } - - @Override - public BsonTimestamp getOperationTime() { - return operationTime; - } - - @Override - public boolean isFirstBatchEmpty() { - return firstBatchEmpty; - } - - @Override - public int getMaxWireVersion() { - return maxWireVersion; - } - - private boolean limitReached() { - return Math.abs(limit) != 0 && count.get() >= Math.abs(limit); - } - - private void getMore(final ServerCursor cursor, final SingleResultCallback> callback) { - if (pinnedConnection != null) { - getMore(pinnedConnection.retain(), cursor, callback); - } else { - connectionSource.getConnection((connection, t) -> { - if (t != null) { - endOperationInProgress(); - callback.onResult(null, t); - } else { - getMore(assertNotNull(connection), cursor, callback); - } - }); - } - } - - private void getMore(final AsyncConnection connection, final ServerCursor cursor, final SingleResultCallback> callback) { - connection.commandAsync(namespace.getDatabaseName(), asGetMoreCommandDocument(cursor.getId(), connection.getDescription()), - NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), CommandResultDocumentCodec.create(decoder, "nextBatch"), - connectionSource, new CommandResultSingleResultCallback(connection, cursor, callback)); - } - - private BsonDocument asGetMoreCommandDocument(final long cursorId, final ConnectionDescription connectionDescription) { - BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursorId)) - .append("collection", new BsonString(namespace.getCollectionName())); - - int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count.get())); - if (batchSizeForGetMoreCommand != 0) { - document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand)); - } - if (maxTimeMS != 0) { - document.append("maxTimeMS", new BsonInt64(maxTimeMS)); - } - if (serverIsAtLeastVersionFourDotFour(connectionDescription)) { - putIfNotNull(document, "comment", comment); - } - return document; - } - - private void killCursorOnClose() { - ServerCursor localCursor = getServerCursor(); - if (localCursor != null) { - if (pinnedConnection != null) { - killCursorAsynchronouslyAndReleaseConnectionAndSource(pinnedConnection, localCursor); - } else { - connectionSource.getConnection((connection, t) -> { - if (t != null) { - connectionSource.release(); - } else { - killCursorAsynchronouslyAndReleaseConnectionAndSource(assertNotNull(connection), localCursor); - } - }); - } - } else if (pinnedConnection != null) { - pinnedConnection.release(); - } - } - - private void killCursor(final AsyncConnection connection) { - ServerCursor localCursor = cursor.getAndSet(null); - if (localCursor != null) { - killCursorAsynchronouslyAndReleaseConnectionAndSource(connection.retain(), localCursor); - } else { - connectionSource.release(); - } - } - - private void killCursorAsynchronouslyAndReleaseConnectionAndSource(final AsyncConnection connection, final ServerCursor localCursor) { - connection.commandAsync(namespace.getDatabaseName(), asKillCursorsCommandDocument(localCursor), NO_OP_FIELD_NAME_VALIDATOR, - ReadPreference.primary(), new BsonDocumentCodec(), connectionSource, (result, t) -> { - connection.release(); - connectionSource.release(); - }); - } - - private BsonDocument asKillCursorsCommandDocument(final ServerCursor localCursor) { - return new BsonDocument("killCursors", new BsonString(namespace.getCollectionName())) - .append("cursors", new BsonArray(singletonList(new BsonInt64(localCursor.getId())))); - } - - private void endOperationInProgress() { - boolean closePending; - lock.lock(); - try { - isOperationInProgress = false; - closePending = this.isClosePending; - } finally { - lock.unlock(); - } - if (closePending) { - close(); - } - } - - - private void handleGetMoreQueryResult(final AsyncConnection connection, final SingleResultCallback> callback, - final QueryResult result) { - logQueryResult(result); - cursor.set(result.getCursor()); - if (isClosePending) { - try { - connection.release(); - if (result.getCursor() == null) { - connectionSource.release(); - } - endOperationInProgress(); - } finally { - callback.onResult(null, null); - } - } else if (result.getResults().isEmpty() && result.getCursor() != null) { - getMore(connection, assertNotNull(result.getCursor()), callback); - } else { - count.addAndGet(result.getResults().size()); - if (limitReached()) { - killCursor(connection); - connection.release(); - } else { - connection.release(); - if (result.getCursor() == null) { - connectionSource.release(); - } - } - endOperationInProgress(); - - if (result.getResults().isEmpty()) { - callback.onResult(null, null); - } else { - callback.onResult(result.getResults(), null); - } - } - } - - private void logQueryResult(final QueryResult result) { - LOGGER.debug(format("Received batch of %d documents with cursorId %d from server %s", result.getResults().size(), - result.getCursorId(), result.getAddress())); - } - - private class CommandResultSingleResultCallback implements SingleResultCallback { - private final AsyncConnection connection; - private final ServerCursor cursor; - private final SingleResultCallback> callback; - - CommandResultSingleResultCallback(final AsyncConnection connection, final ServerCursor cursor, - final SingleResultCallback> callback) { - this.connection = connection; - this.cursor = cursor; - this.callback = errorHandlingCallback(callback, LOGGER); - } - - @Override - public void onResult(@Nullable final BsonDocument result, @Nullable final Throwable t) { - if (t != null) { - Throwable translatedException = t instanceof MongoCommandException - ? translateCommandException((MongoCommandException) t, cursor) - : t; - connection.release(); - endOperationInProgress(); - callback.onResult(null, translatedException); - } else { - assertNotNull(result); - QueryResult queryResult = getMoreCursorDocumentToQueryResult(result.getDocument(CURSOR), - connection.getDescription().getServerAddress()); - postBatchResumeToken = getPostBatchResumeTokenFromResponse(result); - handleGetMoreQueryResult(connection, callback, queryResult); - } - } - } - - @Nullable - ServerCursor getServerCursor() { - return cursor.get(); - } - - @Nullable - private BsonDocument getPostBatchResumeTokenFromResponse(final BsonDocument result) { - BsonDocument cursor = result.getDocument(CURSOR, null); - if (cursor != null) { - return cursor.getDocument(POST_BATCH_RESUME_TOKEN, null); - } - return null; - } -} diff --git a/driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchQueryCursor.java b/driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchCursor.java similarity index 63% rename from driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchQueryCursor.java rename to driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchCursor.java index f29cda04dae..57b20ff1711 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchQueryCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/AsyncSingleBatchCursor.java @@ -19,19 +19,26 @@ import com.mongodb.MongoException; import com.mongodb.internal.async.AsyncBatchCursor; import com.mongodb.internal.async.SingleResultCallback; -import com.mongodb.internal.connection.QueryResult; import java.util.List; -import static com.mongodb.assertions.Assertions.isTrue; +import static java.util.Collections.emptyList; -class AsyncSingleBatchQueryCursor implements AsyncBatchCursor { - private volatile QueryResult firstBatch; - private volatile boolean closed; +class AsyncSingleBatchCursor implements AsyncBatchCursor { - AsyncSingleBatchQueryCursor(final QueryResult firstBatch) { - this.firstBatch = firstBatch; - isTrue("Empty Cursor", firstBatch.getCursor() == null); + static AsyncSingleBatchCursor createEmptyAsyncSingleBatchCursor(final int batchSize) { + return new AsyncSingleBatchCursor<>(emptyList(), batchSize); + } + + private final List batch; + private final int batchSize; + + private volatile boolean hasNext = true; + private volatile boolean closed = false; + + AsyncSingleBatchCursor(final List batch, final int batchSize) { + this.batch = batch; + this.batchSize = batchSize; } @Override @@ -43,13 +50,12 @@ public void close() { public void next(final SingleResultCallback> callback) { if (closed) { callback.onResult(null, new MongoException("next() called after the cursor was closed.")); - } else if (firstBatch != null && !firstBatch.getResults().isEmpty()) { - List results = firstBatch.getResults(); - firstBatch = null; - callback.onResult(results, null); + } else if (hasNext && !batch.isEmpty()) { + hasNext = false; + callback.onResult(batch, null); } else { closed = true; - callback.onResult(null, null); + callback.onResult(emptyList(), null); } } @@ -60,7 +66,7 @@ public void setBatchSize(final int batchSize) { @Override public int getBatchSize() { - return 0; + return batchSize; } @Override diff --git a/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java index acf70090457..a5d36f86f55 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java @@ -64,23 +64,23 @@ AggregateResponseBatchCursor getWrapped() { @Override public boolean hasNext() { - return resumeableOperation(queryBatchCursor -> { + return resumeableOperation(commandBatchCursor -> { try { - return queryBatchCursor.hasNext(); + return commandBatchCursor.hasNext(); } finally { - cachePostBatchResumeToken(queryBatchCursor); + cachePostBatchResumeToken(commandBatchCursor); } }); } @Override public List next() { - return resumeableOperation(queryBatchCursor -> { + return resumeableOperation(commandBatchCursor -> { try { - return convertAndProduceLastId(queryBatchCursor.next(), changeStreamOperation.getDecoder(), + return convertAndProduceLastId(commandBatchCursor.next(), changeStreamOperation.getDecoder(), lastId -> resumeToken = lastId); } finally { - cachePostBatchResumeToken(queryBatchCursor); + cachePostBatchResumeToken(commandBatchCursor); } }); } @@ -92,12 +92,12 @@ public int available() { @Override public List tryNext() { - return resumeableOperation(queryBatchCursor -> { + return resumeableOperation(commandBatchCursor -> { try { - return convertAndProduceLastId(queryBatchCursor.tryNext(), changeStreamOperation.getDecoder(), + return convertAndProduceLastId(commandBatchCursor.tryNext(), changeStreamOperation.getDecoder(), lastId -> resumeToken = lastId); } finally { - cachePostBatchResumeToken(queryBatchCursor); + cachePostBatchResumeToken(commandBatchCursor); } }); } @@ -155,9 +155,9 @@ public int getMaxWireVersion() { return maxWireVersion; } - private void cachePostBatchResumeToken(final AggregateResponseBatchCursor queryBatchCursor) { - if (queryBatchCursor.getPostBatchResumeToken() != null) { - resumeToken = queryBatchCursor.getPostBatchResumeToken(); + private void cachePostBatchResumeToken(final AggregateResponseBatchCursor commandBatchCursor) { + if (commandBatchCursor.getPostBatchResumeToken() != null) { + resumeToken = commandBatchCursor.getPostBatchResumeToken(); } } @@ -170,7 +170,7 @@ static List convertAndProduceLastId(@Nullable final List final Decoder decoder, final Consumer lastIdConsumer) { List results = null; - if (rawDocuments != null) { + if (rawDocuments != null && !rawDocuments.isEmpty()) { results = new ArrayList<>(); for (RawBsonDocument rawDocument : rawDocuments) { if (!rawDocument.containsKey("_id")) { diff --git a/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursor.java new file mode 100644 index 00000000000..d48cc9906c7 --- /dev/null +++ b/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursor.java @@ -0,0 +1,383 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.MongoCommandException; +import com.mongodb.MongoException; +import com.mongodb.MongoNamespace; +import com.mongodb.MongoSocketException; +import com.mongodb.ReadPreference; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; +import com.mongodb.annotations.ThreadSafe; +import com.mongodb.connection.ServerType; +import com.mongodb.internal.VisibleForTesting; +import com.mongodb.internal.async.SingleResultCallback; +import com.mongodb.internal.binding.ConnectionSource; +import com.mongodb.internal.connection.Connection; +import com.mongodb.lang.Nullable; +import org.bson.BsonDocument; +import org.bson.BsonTimestamp; +import org.bson.BsonValue; +import org.bson.codecs.BsonDocumentCodec; +import org.bson.codecs.Decoder; + +import java.util.List; +import java.util.NoSuchElementException; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import static com.mongodb.assertions.Assertions.assertNotNull; +import static com.mongodb.assertions.Assertions.assertTrue; +import static com.mongodb.assertions.Assertions.notNull; +import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.FIRST_BATCH; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CLOSED_AS_CURSOR; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CLOSED_AS_ITERATOR; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.NEXT_BATCH; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.NO_OP_FIELD_NAME_VALIDATOR; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.getMoreCommandDocument; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.translateCommandException; +import static com.mongodb.internal.operation.OperationHelper.LOGGER; +import static java.lang.String.format; + +class CommandBatchCursor implements AggregateResponseBatchCursor { + + private final MongoNamespace namespace; + private final int limit; + private final Decoder decoder; + private final long maxTimeMS; + @Nullable + private final BsonValue comment; + + private CommandCursorResult commandCursor; + private int batchSize; + @Nullable + private List nextBatch; + private int count = 0; + + private final boolean firstBatchEmpty; + private final int maxWireVersion; + private final ResourceManager resourceManager; + + CommandBatchCursor( + final ServerAddress serverAddress, + final BsonDocument commandCursorDocument, + final int limit, final int batchSize, final long maxTimeMS, + final Decoder decoder, + @Nullable final BsonValue comment, + final ConnectionSource connectionSource, + final Connection connection) { + this.commandCursor = initFromCommandCursorDocument(serverAddress, FIRST_BATCH, commandCursorDocument); + this.namespace = commandCursor.getNamespace(); + this.limit = limit; + this.batchSize = batchSize; + this.maxTimeMS = maxTimeMS; + this.decoder = notNull("decoder", decoder); + this.comment = comment; + this.maxWireVersion = connection.getDescription().getMaxWireVersion(); + this.firstBatchEmpty = commandCursor.getResults().isEmpty(); + + Connection connectionToPin = null; + boolean releaseServerAndResources = false; + if (limitReached()) { + releaseServerAndResources = true; + } else if (connectionSource.getServerDescription().getType() == ServerType.LOAD_BALANCER) { + connectionToPin = connection; + } + + resourceManager = new ResourceManager(namespace, connectionSource, connectionToPin, commandCursor.getServerCursor()); + if (releaseServerAndResources) { + resourceManager.releaseServerAndClientResources(connection); + } + } + + @Override + public boolean hasNext() { + return assertNotNull(resourceManager.execute(MESSAGE_IF_CLOSED_AS_CURSOR, this::doHasNext)); + } + + private boolean doHasNext() { + if (nextBatch != null) { + return true; + } + + if (limitReached()) { + return false; + } + + while (resourceManager.getServerCursor() != null) { + getMore(); + if (!resourceManager.operable()) { + throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_CURSOR); + } + if (nextBatch != null) { + return true; + } + } + + return false; + } + + @Override + public List next() { + return assertNotNull(resourceManager.execute(MESSAGE_IF_CLOSED_AS_ITERATOR, this::doNext)); + } + + @Override + public int available() { + return !resourceManager.operable() || nextBatch == null ? 0 : nextBatch.size(); + } + + @Nullable + private List doNext() { + if (!doHasNext()) { + throw new NoSuchElementException(); + } + + List retVal = nextBatch; + nextBatch = null; + return retVal; + } + + + @VisibleForTesting(otherwise = PRIVATE) + boolean isClosed() { + return !resourceManager.operable(); + } + + @Override + public void setBatchSize(final int batchSize) { + this.batchSize = batchSize; + } + + @Override + public int getBatchSize() { + return batchSize; + } + + @Override + public void remove() { + throw new UnsupportedOperationException("Not implemented yet!"); + } + + @Override + public void close() { + resourceManager.close(); + } + + @Nullable + @Override + public List tryNext() { + return resourceManager.execute(MESSAGE_IF_CLOSED_AS_CURSOR, () -> { + if (!tryHasNext()) { + return null; + } + return doNext(); + }); + } + + private boolean tryHasNext() { + if (nextBatch != null) { + return true; + } + + if (limitReached()) { + return false; + } + + if (resourceManager.getServerCursor() != null) { + getMore(); + } + + return nextBatch != null; + } + + @Override + @Nullable + public ServerCursor getServerCursor() { + if (!resourceManager.operable()) { + throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_ITERATOR); + } + return resourceManager.getServerCursor(); + } + + @Override + public ServerAddress getServerAddress() { + if (!resourceManager.operable()) { + throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_ITERATOR); + } + + return commandCursor.getServerAddress(); + } + + @Override + public BsonDocument getPostBatchResumeToken() { + return commandCursor.getPostBatchResumeToken(); + } + + @Override + public BsonTimestamp getOperationTime() { + return commandCursor.getOperationTime(); + } + + @Override + public boolean isFirstBatchEmpty() { + return firstBatchEmpty; + } + + @Override + public int getMaxWireVersion() { + return maxWireVersion; + } + + private void getMore() { + ServerCursor serverCursor = assertNotNull(resourceManager.getServerCursor()); + resourceManager.executeWithConnection(connection -> { + ServerCursor nextServerCursor; + try { + commandCursor = initFromCommandCursorDocument(connection.getDescription().getServerAddress(), NEXT_BATCH, + assertNotNull( + connection.command(namespace.getDatabaseName(), + getMoreCommandDocument(serverCursor.getId(), connection.getDescription(), namespace, + limit, batchSize, count, maxTimeMS, comment), + NO_OP_FIELD_NAME_VALIDATOR, + ReadPreference.primary(), + CommandResultDocumentCodec.create(decoder, NEXT_BATCH), + assertNotNull(resourceManager.getConnectionSource())))); + nextServerCursor = commandCursor.getServerCursor(); + } catch (MongoCommandException e) { + throw translateCommandException(e, serverCursor); + } + + resourceManager.setServerCursor(nextServerCursor); + if (!resourceManager.operable() || limitReached() || nextServerCursor == null) { + resourceManager.releaseServerAndClientResources(connection); + } + }); + } + + private CommandCursorResult initFromCommandCursorDocument(final ServerAddress serverAddress, final String fieldNameContainingBatch, + final BsonDocument commandCursorDocument) { + CommandCursorResult commandCursorResult = new CommandCursorResult<>(serverAddress, fieldNameContainingBatch, + commandCursorDocument); + this.nextBatch = commandCursorResult.getResults().isEmpty() ? null : commandCursorResult.getResults(); + this.count += commandCursorResult.getResults().size(); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug(format("Received batch of %d documents with cursorId %d from server %s", commandCursorResult.getResults().size(), + commandCursorResult.getCursorId(), commandCursorResult.getServerAddress())); + } + return commandCursorResult; + } + + private boolean limitReached() { + return Math.abs(limit) != 0 && count >= Math.abs(limit); + } + + @ThreadSafe + private static final class ResourceManager extends CursorResourceManager { + + ResourceManager( + final MongoNamespace namespace, + final ConnectionSource connectionSource, + @Nullable final Connection connectionToPin, + @Nullable final ServerCursor serverCursor) { + super(namespace, connectionSource, connectionToPin, serverCursor); + } + + /** + * Thread-safe. + * Executes {@code operation} within the {@link #tryStartOperation()}/{@link #endOperation()} bounds. + * + * @throws IllegalStateException If {@linkplain CommandBatchCursor#close() closed}. + */ + @Nullable + R execute(final String exceptionMessageIfClosed, final Supplier operation) throws IllegalStateException { + if (!tryStartOperation()) { + throw new IllegalStateException(exceptionMessageIfClosed); + } + try { + return operation.get(); + } finally { + endOperation(); + } + } + + private Connection connection() { + assertTrue(getState() != State.IDLE); + Connection pinnedConnection = getPinnedConnection(); + if (pinnedConnection != null) { + return assertNotNull(pinnedConnection).retain(); + } else { + return assertNotNull(getConnectionSource()).getConnection(); + } + } + + @Override + void markAsPinned(final Connection connectionToPin, final Connection.PinningMode pinningMode) { + connectionToPin.markAsPinned(pinningMode); + } + + @Override + void executeWithConnection(final Consumer action) { + Connection connection = connection(); + try { + action.accept(connection); + } catch (MongoSocketException e) { + try { + onCorruptedConnection(connection); + } catch (Exception suppressed) { + e.addSuppressed(suppressed); + } + throw e; + } finally { + connection.release(); + } + } + + @Override + void executeWithConnection(final Consumer action, final SingleResultCallback callback) { + throw new UnsupportedOperationException(); + } + + @Override + void doClose() { + if (isSkipReleasingServerResourcesOnClose()) { + unsetServerCursor(); + } + + try { + if (getServerCursor() != null) { + executeWithConnection(this::releaseServerResources); + } + } catch (MongoException e) { + // ignore exceptions when releasing server resources + } finally { + // guarantee that regardless of exceptions, `serverCursor` is null and client resources are released + unsetServerCursor(); + releaseClientResources(); + } + } + + @Override + void killServerCursor(final MongoNamespace namespace, final ServerCursor localServerCursor, final Connection localConnection) { + localConnection.command(namespace.getDatabaseName(), getKillCursorsCommand(namespace, localServerCursor), + NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), assertNotNull(getConnectionSource())); + releaseClientResources(); + } + } +} diff --git a/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursorHelper.java b/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursorHelper.java new file mode 100644 index 00000000000..d7e125b5eb8 --- /dev/null +++ b/driver-core/src/main/com/mongodb/internal/operation/CommandBatchCursorHelper.java @@ -0,0 +1,78 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.MongoCommandException; +import com.mongodb.MongoCursorNotFoundException; +import com.mongodb.MongoNamespace; +import com.mongodb.MongoQueryException; +import com.mongodb.ServerCursor; +import com.mongodb.connection.ConnectionDescription; +import com.mongodb.internal.validator.NoOpFieldNameValidator; +import com.mongodb.lang.Nullable; +import org.bson.BsonDocument; +import org.bson.BsonInt32; +import org.bson.BsonInt64; +import org.bson.BsonString; +import org.bson.BsonValue; +import org.bson.FieldNameValidator; + +import static com.mongodb.internal.operation.CursorHelper.getNumberToReturn; +import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; +import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionFourDotFour; + +final class CommandBatchCursorHelper { + + static final String FIRST_BATCH = "firstBatch"; + static final String NEXT_BATCH = "nextBatch"; + static final FieldNameValidator NO_OP_FIELD_NAME_VALIDATOR = new NoOpFieldNameValidator(); + static final String MESSAGE_IF_CLOSED_AS_CURSOR = "Cursor has been closed"; + static final String MESSAGE_IF_CLOSED_AS_ITERATOR = "Iterator has been closed"; + + static final String MESSAGE_IF_CONCURRENT_OPERATION = "Another operation is currently in progress, concurrent operations are not " + + "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) { + 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 (maxTimeMS != 0) { + document.append("maxTimeMS", new BsonInt64(maxTimeMS)); + } + if (serverIsAtLeastVersionFourDotFour(connectionDescription)) { + putIfNotNull(document, "comment", comment); + } + return document; + } + + static MongoQueryException translateCommandException(final MongoCommandException commandException, final ServerCursor cursor) { + if (commandException.getErrorCode() == 43) { + return new MongoCursorNotFoundException(cursor.getId(), commandException.getResponse(), cursor.getAddress()); + } else { + return new MongoQueryException(commandException.getResponse(), commandException.getServerAddress()); + } + } + + private CommandBatchCursorHelper() { + } +} diff --git a/driver-core/src/main/com/mongodb/internal/connection/QueryResult.java b/driver-core/src/main/com/mongodb/internal/operation/CommandCursorResult.java similarity index 52% rename from driver-core/src/main/com/mongodb/internal/connection/QueryResult.java rename to driver-core/src/main/com/mongodb/internal/operation/CommandCursorResult.java index 52970ba7b94..8fa6c1bdde3 100644 --- a/driver-core/src/main/com/mongodb/internal/connection/QueryResult.java +++ b/driver-core/src/main/com/mongodb/internal/operation/CommandCursorResult.java @@ -14,40 +14,52 @@ * limitations under the License. */ -package com.mongodb.internal.connection; +package com.mongodb.internal.operation; import com.mongodb.MongoNamespace; import com.mongodb.ServerAddress; import com.mongodb.ServerCursor; import com.mongodb.lang.Nullable; +import org.bson.BsonDocument; +import org.bson.BsonTimestamp; import java.util.List; +import static com.mongodb.assertions.Assertions.isTrue; + /** - * A batch of query results. + * The command cursor result * *

This class is not part of the public API and may be removed or changed at any time

*/ -public class QueryResult { - private final MongoNamespace namespace; +public class CommandCursorResult { + + private static final String CURSOR = "cursor"; + private static final String POST_BATCH_RESUME_TOKEN = "postBatchResumeToken"; + private static final String OPERATION_TIME = "operationTime"; + + private final BsonDocument cursorDocument; + private final ServerAddress serverAddress; private final List results; + private final MongoNamespace namespace; private final long cursorId; - private final ServerAddress serverAddress; + @Nullable + private final BsonTimestamp operationTime; + @Nullable + private final BsonDocument postBatchResumeToken; - /** - * Construct an instance. - * - * @param namespace the namespace - * @param results the query results - * @param cursorId the cursor id - * @param serverAddress the server address - */ - public QueryResult(@Nullable final MongoNamespace namespace, final List results, final long cursorId, - final ServerAddress serverAddress) { - this.namespace = namespace; - this.results = results; - this.cursorId = cursorId; + public CommandCursorResult( + final ServerAddress serverAddress, + final String fieldNameContainingBatch, + final BsonDocument commandCursorDocument) { + isTrue("Contains cursor", commandCursorDocument.isDocument(CURSOR)); this.serverAddress = serverAddress; + this.cursorDocument = commandCursorDocument.getDocument(CURSOR); + this.results = BsonDocumentWrapperHelper.toList(cursorDocument, fieldNameContainingBatch); + this.namespace = new MongoNamespace(cursorDocument.getString("ns").getValue()); + this.cursorId = cursorDocument.getNumber("id").longValue(); + this.operationTime = cursorDocument.getTimestamp(OPERATION_TIME, null); + this.postBatchResumeToken = cursorDocument.getDocument(POST_BATCH_RESUME_TOKEN, null); } /** @@ -55,7 +67,6 @@ public QueryResult(@Nullable final MongoNamespace namespace, final List resul * * @return the namespace */ - @Nullable public MongoNamespace getNamespace() { return namespace; } @@ -66,7 +77,7 @@ public MongoNamespace getNamespace() { * @return the cursor, which may be null if it's been exhausted */ @Nullable - public ServerCursor getCursor() { + public ServerCursor getServerCursor() { return cursorId == 0 ? null : new ServerCursor(cursorId, serverAddress); } @@ -84,11 +95,21 @@ public List getResults() { * * @return the server address */ - public ServerAddress getAddress() { + public ServerAddress getServerAddress() { return serverAddress; } public long getCursorId() { return cursorId; } + + @Nullable + public BsonDocument getPostBatchResumeToken() { + return postBatchResumeToken; + } + + @Nullable + public BsonTimestamp getOperationTime() { + return operationTime; + } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/CursorResourceManager.java b/driver-core/src/main/com/mongodb/internal/operation/CursorResourceManager.java new file mode 100644 index 00000000000..dcf0a232658 --- /dev/null +++ b/driver-core/src/main/com/mongodb/internal/operation/CursorResourceManager.java @@ -0,0 +1,327 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.MongoNamespace; +import com.mongodb.ServerCursor; +import com.mongodb.annotations.ThreadSafe; +import com.mongodb.internal.async.SingleResultCallback; +import com.mongodb.internal.binding.ReferenceCounted; +import com.mongodb.internal.connection.Connection; +import com.mongodb.lang.Nullable; +import org.bson.BsonArray; +import org.bson.BsonDocument; +import org.bson.BsonInt64; +import org.bson.BsonString; + +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.StampedLock; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import static com.mongodb.assertions.Assertions.assertNotNull; +import static com.mongodb.assertions.Assertions.assertNull; +import static com.mongodb.assertions.Assertions.assertTrue; +import static com.mongodb.assertions.Assertions.fail; +import static com.mongodb.assertions.Assertions.notNull; +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CONCURRENT_OPERATION; +import static java.util.Collections.singletonList; + +/** + * This is the resource manager for {@link CommandBatchCursor} or {@link AsyncCommandBatchCursor} implementations. + *

+ * This class maintains all resources that must be released in {@link CommandBatchCursor#close()} / + * {@link AsyncCommandBatchCursor#close()}. The abstract {@linkplain #doClose() deferred close action} is such that it is totally + * ordered with other operations of {@link CommandBatchCursor} / {@link AsyncCommandBatchCursor} (methods {@link #tryStartOperation()}/ + * {@link #endOperation()} must be used properly to enforce the order) despite the method {@link CommandBatchCursor#close()} / + * {@link AsyncCommandBatchCursor#close()} being called concurrently with those operations. + *

+ * This total order induces the happens-before order. + *

+ * The deferred close action does not violate externally observable idempotence of {@link CommandBatchCursor#close()} / + * {@link AsyncCommandBatchCursor#close()}, because the close method is allowed to release resources "eventually". + *

+ * Only methods explicitly documented as thread-safe are thread-safe, + * others are not and rely on the total order mentioned above. + */ +@ThreadSafe +abstract class CursorResourceManager { + + private final MongoNamespace namespace; + private final Lock lock; + private volatile State state; + @Nullable + private volatile CS connectionSource; + @Nullable + private volatile C pinnedConnection; + @Nullable + private volatile ServerCursor serverCursor; + private volatile boolean skipReleasingServerResourcesOnClose; + + CursorResourceManager( + final MongoNamespace namespace, + final CS connectionSource, + @Nullable final C connectionToPin, + @Nullable final ServerCursor serverCursor) { + this.namespace = notNull("mongoNamespace", namespace); + this.lock = new StampedLock().asWriteLock(); + this.state = State.IDLE; + if (serverCursor != null) { + connectionSource.retain(); + this.connectionSource = notNull("connectionSource", connectionSource); + if (connectionToPin != null) { + connectionToPin.retain(); + markAsPinned(connectionToPin, Connection.PinningMode.CURSOR); + this.pinnedConnection = connectionToPin; + } + } + this.skipReleasingServerResourcesOnClose = false; + this.serverCursor = serverCursor; + } + + public State getState() { + return state; + } + + @Nullable + public CS getConnectionSource() { + return connectionSource; + } + + @Nullable + public C getPinnedConnection() { + return pinnedConnection; + } + + public boolean isSkipReleasingServerResourcesOnClose() { + return skipReleasingServerResourcesOnClose; + } + + abstract void markAsPinned(C connectionToPin, Connection.PinningMode pinningMode); + + abstract void executeWithConnection(Consumer action); + + abstract void executeWithConnection(Consumer action, SingleResultCallback callback); + + /** + * This method is never executed concurrently with either itself or other operations + * demarcated by {@link #tryStartOperation()}/{@link #endOperation()}. + */ + abstract void doClose(); + + abstract void killServerCursor(MongoNamespace namespace, ServerCursor serverCursor, C connection); + + /** + * Thread-safe. + */ + boolean operable() { + return state.operable(); + } + + /** + * Thread-safe. + * Returns {@code true} iff started an operation. + * If {@linkplain #operable() closed}, then returns false, otherwise completes abruptly. + * + * @throws IllegalStateException Iff another operation is in progress. + */ + boolean tryStartOperation() throws IllegalStateException { + return withLock(() -> { + State localState = state; + if (!localState.operable()) { + return false; + } else if (localState == State.IDLE) { + state = State.OPERATION_IN_PROGRESS; + return true; + } else if (localState == State.OPERATION_IN_PROGRESS) { + throw new IllegalStateException(MESSAGE_IF_CONCURRENT_OPERATION); + } else { + throw fail(state.toString()); + } + }); + } + + /** + * Thread-safe. + */ + void endOperation() { + boolean doClose = withLock(() -> { + State localState = state; + if (localState == State.OPERATION_IN_PROGRESS) { + state = State.IDLE; + } else if (localState == State.CLOSE_PENDING) { + state = State.CLOSED; + return true; + } else if (localState != State.CLOSED) { + fail(localState.toString()); + } + return false; + }); + + if (doClose) { + doClose(); + } + } + + /** + * Thread-safe. + */ + void close() { + boolean doClose = withLock(() -> { + State localState = state; + if (localState == State.OPERATION_IN_PROGRESS) { + state = State.CLOSE_PENDING; + } else if (localState != State.CLOSED) { + state = State.CLOSED; + return true; + } + return false; + }); + + if (doClose) { + doClose(); + } + } + + void onCorruptedConnection(@Nullable final C corruptedConnection) { + // if `pinnedConnection` is corrupted, then we cannot kill `serverCursor` via such a connection + C localPinnedConnection = pinnedConnection; + if (localPinnedConnection != null) { + assertTrue(corruptedConnection == localPinnedConnection); + skipReleasingServerResourcesOnClose = true; + } + } + + /** + * Thread-safe. + */ + @Nullable + ServerCursor getServerCursor() { + return serverCursor; + } + + void unsetServerCursor() { + withLock(() -> { + this.serverCursor = null; + }); + } + + void setServerCursor(@Nullable final ServerCursor serverCursor) { + assertTrue(state.inProgress()); + assertNotNull(this.serverCursor); + // without `connectionSource` we will not be able to kill `serverCursor` later + assertNotNull(connectionSource); + this.serverCursor = serverCursor; + if (serverCursor == null) { + releaseClientResources(); + } + } + + void releaseServerAndClientResources(final C connection) { + ServerCursor localServerCursor = withLockNullable(() -> { + ServerCursor local = serverCursor; + serverCursor = null; + return local; + }); + if (localServerCursor != null) { + killServerCursor(namespace, localServerCursor, connection); + } else { + releaseClientResources(); + } + } + + void releaseServerResources(final C connection) { + ServerCursor localServerCursor = serverCursor; + serverCursor = null; + if (localServerCursor != null) { + killServerCursor(namespace, localServerCursor, connection); + } + } + + BsonDocument getKillCursorsCommand(final MongoNamespace namespace, final ServerCursor serverCursor) { + return new BsonDocument("killCursors", new BsonString(namespace.getCollectionName())) + .append("cursors", new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))); + } + + void releaseClientResources() { + assertNull(serverCursor); + CS localConnectionSource = connectionSource; + if (localConnectionSource != null) { + localConnectionSource.release(); + connectionSource = null; + } + C localPinnedConnection = pinnedConnection; + if (localPinnedConnection != null) { + localPinnedConnection.release(); + pinnedConnection = null; + } + } + + private void withLock(final Runnable runnable) { + try { + lock.lock(); + runnable.run(); + } finally { + lock.unlock(); + } + } + + @Nullable + private T withLockNullable(final Supplier supplier) { + try { + lock.lock(); + return supplier.get(); + } finally { + lock.unlock(); + } + } + + private T withLock(final Supplier supplier) { + try { + lock.lock(); + return supplier.get(); + } finally { + lock.unlock(); + } + } + + enum State { + IDLE(true, false), + OPERATION_IN_PROGRESS(true, true), + /** + * Implies {@link #OPERATION_IN_PROGRESS}. + */ + CLOSE_PENDING(false, true), + CLOSED(false, false); + + private final boolean operable; + private final boolean inProgress; + + State(final boolean operable, final boolean inProgress) { + this.operable = operable; + this.inProgress = inProgress; + } + + boolean operable() { + return operable; + } + + boolean inProgress() { + return inProgress; + } + } +} diff --git a/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java b/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java index a64c4cbfadd..488e044333e 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/DistinctOperation.java @@ -23,7 +23,6 @@ import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.session.SessionContext; import com.mongodb.lang.Nullable; import org.bson.BsonDocument; @@ -36,14 +35,12 @@ import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; -import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.executeRetryableReadAsync; import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator; import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; import static com.mongodb.internal.operation.DocumentHelper.putIfNotZero; import static com.mongodb.internal.operation.OperationHelper.LOGGER; import static com.mongodb.internal.operation.OperationReadConcernHelper.appendReadConcernToCommand; -import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; import static com.mongodb.internal.operation.SyncOperationHelper.executeRetryableRead; /** @@ -119,38 +116,23 @@ public DistinctOperation comment(final BsonValue comment) { @Override public BatchCursor execute(final ReadBinding binding) { - return executeRetryableRead(binding, namespace.getDatabaseName(), getCommandCreator(binding.getSessionContext()), - createCommandDecoder(), transformer(), retryReads); + return executeRetryableRead(binding, namespace.getDatabaseName(), + getCommandCreator(binding.getSessionContext()), createCommandDecoder(), + SyncOperationHelper.singleBatchCursorTransformer(VALUES), retryReads); } @Override public void executeAsync(final AsyncReadBinding binding, final SingleResultCallback> callback) { - executeRetryableReadAsync(binding, namespace.getDatabaseName(), getCommandCreator(binding.getSessionContext()), - createCommandDecoder(), asyncTransformer(), retryReads, errorHandlingCallback(callback, LOGGER)); + executeRetryableReadAsync(binding, namespace.getDatabaseName(), + getCommandCreator(binding.getSessionContext()), createCommandDecoder(), + AsyncOperationHelper.asyncSingleBatchCursorTransformer(VALUES), retryReads, + errorHandlingCallback(callback, LOGGER)); } private Codec createCommandDecoder() { return CommandResultDocumentCodec.create(decoder, VALUES); } - private QueryResult createQueryResult(final BsonDocument result, final ConnectionDescription description) { - return new QueryResult<>(namespace, BsonDocumentWrapperHelper.toList(result, VALUES), 0L, - description.getServerAddress()); - } - - private CommandReadTransformer> transformer() { - return (result, source, connection) -> { - QueryResult queryResult = createQueryResult(result, connection.getDescription()); - return new QueryBatchCursor<>(queryResult, 0, 0, decoder, comment, source); - }; - } - - private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> { - QueryResult queryResult = createQueryResult(result, connection.getDescription()); - return new AsyncSingleBatchQueryCursor<>(queryResult); - }; - } private CommandCreator getCommandCreator(final SessionContext sessionContext) { return (serverDescription, connectionDescription) -> getCommand(sessionContext, connectionDescription); diff --git a/driver-core/src/main/com/mongodb/internal/operation/FindOperation.java b/driver-core/src/main/com/mongodb/internal/operation/FindOperation.java index dcb94211fcf..83441bdc5e2 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/FindOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/FindOperation.java @@ -29,7 +29,6 @@ import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; import com.mongodb.internal.connection.NoOpSessionContext; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.session.SessionContext; import com.mongodb.lang.Nullable; import org.bson.BsonBoolean; @@ -57,7 +56,6 @@ import static com.mongodb.internal.operation.ExplainHelper.asExplainCommand; import static com.mongodb.internal.operation.OperationHelper.LOGGER; import static com.mongodb.internal.operation.OperationHelper.canRetryRead; -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult; import static com.mongodb.internal.operation.OperationReadConcernHelper.appendReadConcernToCommand; import static com.mongodb.internal.operation.ServerVersionHelper.MIN_WIRE_VERSION; import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; @@ -471,13 +469,10 @@ private boolean isAwaitData() { return cursorType == CursorType.TailableAwait; } - private CommandReadTransformer> transformer() { - return (result, source, connection) -> { - QueryResult queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"), - connection.getDescription().getServerAddress()); - return new QueryBatchCursor<>(queryResult, limit, batchSize, getMaxTimeForCursor(), decoder, comment, source, connection, - result); - }; + private CommandReadTransformer> transformer() { + return (result, source, connection) -> + new CommandBatchCursor<>(connection.getDescription().getServerAddress(), result, limit, batchSize, + getMaxTimeForCursor(), decoder, comment, source, connection); } private long getMaxTimeForCursor() { @@ -485,11 +480,9 @@ private long getMaxTimeForCursor() { } private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> { - QueryResult queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"), - connection.getDescription().getServerAddress()); - return new AsyncQueryBatchCursor<>(queryResult, limit, batchSize, getMaxTimeForCursor(), decoder, comment, source, - connection, result); - }; + return (result, source, connection) -> + new AsyncCommandBatchCursor<>(connection.getDescription().getServerAddress(), result, limit, batchSize, + getMaxTimeForCursor(), decoder, comment, source, connection); + } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/ListCollectionsOperation.java b/driver-core/src/main/com/mongodb/internal/operation/ListCollectionsOperation.java index fa2a5dcd995..ea122fe0250 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ListCollectionsOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/ListCollectionsOperation.java @@ -17,12 +17,10 @@ package com.mongodb.internal.operation; import com.mongodb.MongoCommandException; -import com.mongodb.MongoNamespace; import com.mongodb.internal.async.AsyncBatchCursor; import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.async.function.AsyncCallbackSupplier; import com.mongodb.internal.async.function.RetryState; -import com.mongodb.internal.binding.AsyncConnectionSource; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; import com.mongodb.lang.Nullable; @@ -40,11 +38,11 @@ import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync; -import static com.mongodb.internal.operation.AsyncOperationHelper.createEmptyAsyncBatchCursor; import static com.mongodb.internal.operation.AsyncOperationHelper.createReadCommandAndExecuteAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.cursorDocumentToAsyncBatchCursor; import static com.mongodb.internal.operation.AsyncOperationHelper.decorateReadWithRetriesAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.withAsyncSourceAndConnection; +import static com.mongodb.internal.operation.AsyncSingleBatchCursor.createEmptyAsyncSingleBatchCursor; import static com.mongodb.internal.operation.CommandOperationHelper.initialRetryState; import static com.mongodb.internal.operation.CommandOperationHelper.isNamespaceError; import static com.mongodb.internal.operation.CommandOperationHelper.rethrowIfNotNamespaceError; @@ -52,7 +50,7 @@ import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; import static com.mongodb.internal.operation.OperationHelper.LOGGER; import static com.mongodb.internal.operation.OperationHelper.canRetryRead; -import static com.mongodb.internal.operation.OperationHelper.createEmptyBatchCursor; +import static com.mongodb.internal.operation.SingleBatchCursor.createEmptySingleBatchCursor; import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; import static com.mongodb.internal.operation.SyncOperationHelper.createReadCommandAndExecute; import static com.mongodb.internal.operation.SyncOperationHelper.cursorDocumentToBatchCursor; @@ -148,8 +146,8 @@ public BatchCursor execute(final ReadBinding binding) { return createReadCommandAndExecute(retryState, binding, source, databaseName, getCommandCreator(), createCommandDecoder(), commandTransformer(), connection); } catch (MongoCommandException e) { - return rethrowIfNotNamespaceError(e, createEmptyBatchCursor(createNamespace(), decoder, - source.getServerDescription().getAddress(), batchSize)); + return rethrowIfNotNamespaceError(e, + createEmptySingleBatchCursor(source.getServerDescription().getAddress(), batchSize)); } }) ); @@ -173,7 +171,9 @@ public void executeAsync(final AsyncReadBinding binding, final SingleResultCallb if (t != null && !isNamespaceError(t)) { releasingCallback.onResult(null, t); } else { - releasingCallback.onResult(result != null ? result : emptyAsyncCursor(source), null); + releasingCallback.onResult(result != null + ? result + : createEmptyAsyncSingleBatchCursor(batchSize), null); } }); }) @@ -181,20 +181,12 @@ public void executeAsync(final AsyncReadBinding binding, final SingleResultCallb asyncRead.get(errorHandlingCallback(callback, LOGGER)); } - private AsyncBatchCursor emptyAsyncCursor(final AsyncConnectionSource source) { - return createEmptyAsyncBatchCursor(createNamespace(), source.getServerDescription().getAddress()); - } - - private MongoNamespace createNamespace() { - return new MongoNamespace(databaseName, "$cmd.listCollections"); - } - private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, comment, source, connection, batchSize); + return (result, source, connection) -> cursorDocumentToAsyncBatchCursor(result, decoder, comment, source, connection, batchSize); } private CommandReadTransformer> commandTransformer() { - return (result, source, connection) -> cursorDocumentToBatchCursor(result.getDocument("cursor"), decoder, comment, source, connection, batchSize); + return (result, source, connection) -> cursorDocumentToBatchCursor(result, decoder, comment, source, connection, batchSize); } private CommandOperationHelper.CommandCreator getCommandCreator() { diff --git a/driver-core/src/main/com/mongodb/internal/operation/ListDatabasesOperation.java b/driver-core/src/main/com/mongodb/internal/operation/ListDatabasesOperation.java index bacf64601c9..f86ada76f10 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ListDatabasesOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/ListDatabasesOperation.java @@ -16,12 +16,10 @@ package com.mongodb.internal.operation; -import com.mongodb.connection.ConnectionDescription; import com.mongodb.internal.async.AsyncBatchCursor; import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.lang.Nullable; import org.bson.BsonBoolean; import org.bson.BsonDocument; @@ -34,12 +32,10 @@ import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; -import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.executeRetryableReadAsync; import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator; import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; import static com.mongodb.internal.operation.OperationHelper.LOGGER; -import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; import static com.mongodb.internal.operation.SyncOperationHelper.executeRetryableRead; @@ -122,28 +118,16 @@ public ListDatabasesOperation comment(@Nullable final BsonValue comment) { @Override public BatchCursor execute(final ReadBinding binding) { return executeRetryableRead(binding, "admin", getCommandCreator(), - CommandResultDocumentCodec.create(decoder, "databases"), transformer(), retryReads); + CommandResultDocumentCodec.create(decoder, "databases"), + SyncOperationHelper.singleBatchCursorTransformer("databases"), retryReads); } @Override public void executeAsync(final AsyncReadBinding binding, final SingleResultCallback> callback) { executeRetryableReadAsync(binding, "admin", getCommandCreator(), - CommandResultDocumentCodec.create(decoder, "databases"), asyncTransformer(), - retryReads, errorHandlingCallback(callback, LOGGER)); - } - - private CommandReadTransformer> transformer() { - return (result, source, connection) -> new QueryBatchCursor<>(createQueryResult(result, connection.getDescription()), 0, 0, decoder, comment, source); - } - - private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> new AsyncQueryBatchCursor<>(createQueryResult(result, connection.getDescription()), 0, 0, 0, decoder, - comment, source, connection, result); - } - - private QueryResult createQueryResult(final BsonDocument result, final ConnectionDescription description) { - return new QueryResult<>(null, BsonDocumentWrapperHelper.toList(result, "databases"), 0, - description.getServerAddress()); + CommandResultDocumentCodec.create(decoder, "databases"), + AsyncOperationHelper.asyncSingleBatchCursorTransformer("databases"), retryReads, + errorHandlingCallback(callback, LOGGER)); } private CommandCreator getCommandCreator() { diff --git a/driver-core/src/main/com/mongodb/internal/operation/ListIndexesOperation.java b/driver-core/src/main/com/mongodb/internal/operation/ListIndexesOperation.java index 62ecdc953bd..9ac0ea92458 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ListIndexesOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/ListIndexesOperation.java @@ -22,7 +22,6 @@ import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.async.function.AsyncCallbackSupplier; import com.mongodb.internal.async.function.RetryState; -import com.mongodb.internal.binding.AsyncConnectionSource; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; import com.mongodb.lang.Nullable; @@ -39,11 +38,11 @@ import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback; import static com.mongodb.internal.operation.AsyncOperationHelper.CommandReadTransformerAsync; -import static com.mongodb.internal.operation.AsyncOperationHelper.createEmptyAsyncBatchCursor; import static com.mongodb.internal.operation.AsyncOperationHelper.createReadCommandAndExecuteAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.cursorDocumentToAsyncBatchCursor; import static com.mongodb.internal.operation.AsyncOperationHelper.decorateReadWithRetriesAsync; import static com.mongodb.internal.operation.AsyncOperationHelper.withAsyncSourceAndConnection; +import static com.mongodb.internal.operation.AsyncSingleBatchCursor.createEmptyAsyncSingleBatchCursor; import static com.mongodb.internal.operation.CommandOperationHelper.CommandCreator; import static com.mongodb.internal.operation.CommandOperationHelper.initialRetryState; import static com.mongodb.internal.operation.CommandOperationHelper.isNamespaceError; @@ -52,7 +51,7 @@ import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; import static com.mongodb.internal.operation.OperationHelper.LOGGER; import static com.mongodb.internal.operation.OperationHelper.canRetryRead; -import static com.mongodb.internal.operation.OperationHelper.createEmptyBatchCursor; +import static com.mongodb.internal.operation.SingleBatchCursor.createEmptySingleBatchCursor; import static com.mongodb.internal.operation.SyncOperationHelper.CommandReadTransformer; import static com.mongodb.internal.operation.SyncOperationHelper.createReadCommandAndExecute; import static com.mongodb.internal.operation.SyncOperationHelper.cursorDocumentToBatchCursor; @@ -127,8 +126,8 @@ public BatchCursor execute(final ReadBinding binding) { return createReadCommandAndExecute(retryState, binding, source, namespace.getDatabaseName(), getCommandCreator(), createCommandDecoder(), transformer(), connection); } catch (MongoCommandException e) { - return rethrowIfNotNamespaceError(e, createEmptyBatchCursor(namespace, decoder, - source.getServerDescription().getAddress(), batchSize)); + return rethrowIfNotNamespaceError(e, + createEmptySingleBatchCursor(source.getServerDescription().getAddress(), batchSize)); } }) ); @@ -152,7 +151,10 @@ public void executeAsync(final AsyncReadBinding binding, final SingleResultCallb if (t != null && !isNamespaceError(t)) { releasingCallback.onResult(null, t); } else { - releasingCallback.onResult(result != null ? result : emptyAsyncCursor(source), null); + releasingCallback.onResult(result != null + ? result + : createEmptyAsyncSingleBatchCursor(batchSize), + null); } }); }) @@ -160,10 +162,6 @@ public void executeAsync(final AsyncReadBinding binding, final SingleResultCallb asyncRead.get(errorHandlingCallback(callback, LOGGER)); } - private AsyncBatchCursor emptyAsyncCursor(final AsyncConnectionSource source) { - return createEmptyAsyncBatchCursor(namespace, source.getServerDescription().getAddress()); - } - private CommandCreator getCommandCreator() { return (serverDescription, connectionDescription) -> getCommand(); } @@ -179,11 +177,11 @@ private BsonDocument getCommand() { } private CommandReadTransformer> transformer() { - return (result, source, connection) -> cursorDocumentToBatchCursor(result.getDocument("cursor"), decoder, comment, source, connection, batchSize); + return (result, source, connection) -> cursorDocumentToBatchCursor(result, decoder, comment, source, connection, batchSize); } private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> cursorDocumentToAsyncBatchCursor(result.getDocument("cursor"), decoder, comment, source, connection, batchSize); + return (result, source, connection) -> cursorDocumentToAsyncBatchCursor(result, decoder, comment, source, connection, batchSize); } private Codec createCommandDecoder() { diff --git a/driver-core/src/main/com/mongodb/internal/operation/ListSearchIndexesOperation.java b/driver-core/src/main/com/mongodb/internal/operation/ListSearchIndexesOperation.java index 4c471a16bd4..4d9069b9009 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ListSearchIndexesOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/ListSearchIndexesOperation.java @@ -34,9 +34,10 @@ import java.util.Collections; import java.util.concurrent.TimeUnit; -import static com.mongodb.internal.operation.AsyncOperationHelper.createEmptyAsyncBatchCursor; +import static com.mongodb.internal.operation.AsyncSingleBatchCursor.createEmptyAsyncSingleBatchCursor; import static com.mongodb.internal.operation.CommandOperationHelper.isNamespaceError; -import static com.mongodb.internal.operation.OperationHelper.createEmptyBatchCursor; +import static com.mongodb.internal.operation.SingleBatchCursor.createEmptySingleBatchCursor; + /** * An operation that lists Alas Search indexes with the help of {@value #STAGE_LIST_SEARCH_INDEXES} pipeline stage. @@ -90,7 +91,7 @@ public BatchCursor execute(final ReadBinding binding) { if (!isNamespaceError(exception)) { throw exception; } else { - return createEmptyBatchCursor(namespace, decoder, exception.getServerAddress(), cursorBatchSize); + return createEmptySingleBatchCursor(exception.getServerAddress(), cursorBatchSize); } } } @@ -101,9 +102,8 @@ public void executeAsync(final AsyncReadBinding binding, final SingleResultCallb if (exception != null && !isNamespaceError(exception)) { callback.onResult(null, exception); } else if (exception != null) { - MongoCommandException commandException = (MongoCommandException) exception; - AsyncBatchCursor emptyAsyncBatchCursor = createEmptyAsyncBatchCursor(namespace, commandException.getServerAddress()); - callback.onResult(emptyAsyncBatchCursor, null); + int cursorBatchSize = batchSize == null ? 0 : batchSize; + callback.onResult(createEmptyAsyncSingleBatchCursor(cursorBatchSize), null); } else { callback.onResult(cursor, null); } diff --git a/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsAsyncCursor.java b/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsAsyncCursor.java index 1da84755100..ebf331fe47b 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsAsyncCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsAsyncCursor.java @@ -16,18 +16,21 @@ package com.mongodb.internal.operation; -import com.mongodb.internal.connection.QueryResult; +import com.mongodb.internal.async.SingleResultCallback; + +import java.util.List; /** * Cursor representation of the results of an inline map-reduce operation. This allows users to iterate over the results that were returned * from the operation, and also provides access to the statistics returned in the results. */ -class MapReduceInlineResultsAsyncCursor extends AsyncSingleBatchQueryCursor implements MapReduceAsyncBatchCursor { +class MapReduceInlineResultsAsyncCursor implements MapReduceAsyncBatchCursor { + private final AsyncSingleBatchCursor delegate; private final MapReduceStatistics statistics; - MapReduceInlineResultsAsyncCursor(final QueryResult queryResult, final MapReduceStatistics statistics) { - super(queryResult); + MapReduceInlineResultsAsyncCursor(final AsyncSingleBatchCursor delegate, final MapReduceStatistics statistics) { + this.delegate = delegate; this.statistics = statistics; } @@ -35,4 +38,29 @@ class MapReduceInlineResultsAsyncCursor extends AsyncSingleBatchQueryCursor> callback) { + delegate.next(callback); + } + + @Override + public void setBatchSize(final int batchSize) { + delegate.setBatchSize(batchSize); + } + + @Override + public int getBatchSize() { + return delegate.getBatchSize(); + } + + @Override + public boolean isClosed() { + return delegate.isClosed(); + } + + @Override + public void close() { + delegate.close(); + } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsCursor.java b/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsCursor.java index caa2f7fd355..564eac4a8f0 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/MapReduceInlineResultsCursor.java @@ -16,20 +16,21 @@ package com.mongodb.internal.operation; -import com.mongodb.internal.binding.ConnectionSource; -import com.mongodb.internal.connection.QueryResult; -import org.bson.codecs.Decoder; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; + +import java.util.List; /** * Cursor representation of the results of an inline map-reduce operation. This allows users to iterate over the results that were returned * from the operation, and also provides access to the statistics returned in the results. */ -class MapReduceInlineResultsCursor extends QueryBatchCursor implements MapReduceBatchCursor { +class MapReduceInlineResultsCursor implements MapReduceBatchCursor { + private final BatchCursor delegate; private final MapReduceStatistics statistics; - MapReduceInlineResultsCursor(final QueryResult queryResult, final Decoder decoder, final ConnectionSource connectionSource, - final MapReduceStatistics statistics) { - super(queryResult, 0, 0, decoder, null, connectionSource); + MapReduceInlineResultsCursor(final BatchCursor delegate, final MapReduceStatistics statistics) { + this.delegate = delegate; this.statistics = statistics; } @@ -37,4 +38,49 @@ class MapReduceInlineResultsCursor extends QueryBatchCursor implements Map public MapReduceStatistics getStatistics() { return statistics; } + + @Override + public boolean hasNext() { + return delegate.hasNext(); + } + + @Override + public List next() { + return delegate.next(); + } + + @Override + public int available() { + return delegate.available(); + } + + @Override + public void setBatchSize(final int batchSize) { + delegate.setBatchSize(batchSize); + } + + @Override + public int getBatchSize() { + return delegate.getBatchSize(); + } + + @Override + public List tryNext() { + return delegate.tryNext(); + } + + @Override + public ServerCursor getServerCursor() { + return delegate.getServerCursor(); + } + + @Override + public ServerAddress getServerAddress() { + return delegate.getServerAddress(); + } + + @Override + public void close() { + delegate.close(); + } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/MapReduceWithInlineResultsOperation.java b/driver-core/src/main/com/mongodb/internal/operation/MapReduceWithInlineResultsOperation.java index 131591dd6e2..e51852db25c 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/MapReduceWithInlineResultsOperation.java +++ b/driver-core/src/main/com/mongodb/internal/operation/MapReduceWithInlineResultsOperation.java @@ -19,12 +19,10 @@ import com.mongodb.ExplainVerbosity; import com.mongodb.MongoNamespace; import com.mongodb.client.model.Collation; -import com.mongodb.connection.ConnectionDescription; import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.binding.AsyncReadBinding; import com.mongodb.internal.binding.ReadBinding; import com.mongodb.internal.connection.NoOpSessionContext; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.session.SessionContext; import com.mongodb.lang.Nullable; import org.bson.BsonDocument; @@ -215,13 +213,18 @@ private CommandReadOperation createExplainableOperation(final Expl } private CommandReadTransformer> transformer() { - return (result, source, connection) -> new MapReduceInlineResultsCursor<>(createQueryResult(result, connection.getDescription()), decoder, source, + return (result, source, connection) -> + new MapReduceInlineResultsCursor<>( + new SingleBatchCursor<>(BsonDocumentWrapperHelper.toList(result, "results"), 0, + connection.getDescription().getServerAddress()), MapReduceHelper.createStatistics(result)); } private CommandReadTransformerAsync> asyncTransformer() { - return (result, source, connection) -> new MapReduceInlineResultsAsyncCursor<>(createQueryResult(result, connection.getDescription()), - MapReduceHelper.createStatistics(result)); + return (result, source, connection) -> + new MapReduceInlineResultsAsyncCursor<>( + new AsyncSingleBatchCursor<>(BsonDocumentWrapperHelper.toList(result, "results"), 0), + MapReduceHelper.createStatistics(result)); } private CommandCreator getCommandCreator(final SessionContext sessionContext) { @@ -248,9 +251,4 @@ private BsonDocument getCommand(final SessionContext sessionContext, final int m } return commandDocument; } - - private QueryResult createQueryResult(final BsonDocument result, final ConnectionDescription description) { - return new QueryResult<>(namespace, BsonDocumentWrapperHelper.toList(result, "results"), 0, - description.getServerAddress()); - } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/OperationHelper.java b/driver-core/src/main/com/mongodb/internal/operation/OperationHelper.java index 387bb2f5da6..bfa1adbd97e 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/OperationHelper.java +++ b/driver-core/src/main/com/mongodb/internal/operation/OperationHelper.java @@ -17,8 +17,6 @@ package com.mongodb.internal.operation; import com.mongodb.MongoClientException; -import com.mongodb.MongoNamespace; -import com.mongodb.ServerAddress; import com.mongodb.WriteConcern; import com.mongodb.client.model.Collation; import com.mongodb.connection.ConnectionDescription; @@ -30,18 +28,14 @@ import com.mongodb.internal.bulk.DeleteRequest; import com.mongodb.internal.bulk.UpdateRequest; import com.mongodb.internal.bulk.WriteRequest; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.diagnostics.logging.Logger; import com.mongodb.internal.diagnostics.logging.Loggers; import com.mongodb.internal.session.SessionContext; import com.mongodb.lang.NonNull; import com.mongodb.lang.Nullable; import org.bson.BsonDocument; -import org.bson.BsonInt64; -import org.bson.codecs.Decoder; import org.bson.conversions.Bson; -import java.util.Collections; import java.util.List; import java.util.function.Function; import java.util.function.Supplier; @@ -200,26 +194,6 @@ static boolean canRetryRead(final ServerDescription serverDescription, final Ses return true; } - static QueryBatchCursor createEmptyBatchCursor(final MongoNamespace namespace, final Decoder decoder, - final ServerAddress serverAddress, final int batchSize) { - return new QueryBatchCursor<>(new QueryResult<>(namespace, Collections.emptyList(), 0L, - serverAddress), - 0, batchSize, decoder); - } - - static QueryResult cursorDocumentToQueryResult(final BsonDocument cursorDocument, final ServerAddress serverAddress) { - return cursorDocumentToQueryResult(cursorDocument, serverAddress, "firstBatch"); - } - - static QueryResult cursorDocumentToQueryResult(final BsonDocument cursorDocument, final ServerAddress serverAddress, - final String fieldNameContainingBatch) { - long cursorId = ((BsonInt64) cursorDocument.get("id")).getValue(); - MongoNamespace queryResultNamespace = new MongoNamespace(cursorDocument.getString("ns").getValue()); - return new QueryResult<>(queryResultNamespace, BsonDocumentWrapperHelper.toList(cursorDocument, fieldNameContainingBatch), - cursorId, serverAddress); - } - - /** * This internal exception is used to *

    diff --git a/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java deleted file mode 100644 index 139c3e6fd27..00000000000 --- a/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java +++ /dev/null @@ -1,632 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation; - -import com.mongodb.MongoCommandException; -import com.mongodb.MongoException; -import com.mongodb.MongoNamespace; -import com.mongodb.MongoSocketException; -import com.mongodb.ReadPreference; -import com.mongodb.ServerAddress; -import com.mongodb.ServerCursor; -import com.mongodb.annotations.ThreadSafe; -import com.mongodb.connection.ConnectionDescription; -import com.mongodb.connection.ServerType; -import com.mongodb.internal.binding.ConnectionSource; -import com.mongodb.internal.connection.Connection; -import com.mongodb.internal.connection.QueryResult; -import com.mongodb.internal.diagnostics.logging.Logger; -import com.mongodb.internal.diagnostics.logging.Loggers; -import com.mongodb.internal.validator.NoOpFieldNameValidator; -import com.mongodb.lang.Nullable; -import org.bson.BsonArray; -import org.bson.BsonDocument; -import org.bson.BsonInt32; -import org.bson.BsonInt64; -import org.bson.BsonString; -import org.bson.BsonTimestamp; -import org.bson.BsonValue; -import org.bson.FieldNameValidator; -import org.bson.codecs.BsonDocumentCodec; -import org.bson.codecs.Decoder; - -import java.util.List; -import java.util.NoSuchElementException; -import java.util.concurrent.locks.Lock; -import java.util.concurrent.locks.StampedLock; -import java.util.function.Consumer; -import java.util.function.Supplier; - -import static com.mongodb.assertions.Assertions.assertNotNull; -import static com.mongodb.assertions.Assertions.assertNull; -import static com.mongodb.assertions.Assertions.assertTrue; -import static com.mongodb.assertions.Assertions.fail; -import static com.mongodb.assertions.Assertions.isTrueArgument; -import static com.mongodb.assertions.Assertions.notNull; -import static com.mongodb.internal.operation.CursorHelper.getNumberToReturn; -import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; -import static com.mongodb.internal.operation.SyncOperationHelper.getMoreCursorDocumentToQueryResult; -import static com.mongodb.internal.operation.QueryHelper.translateCommandException; -import static com.mongodb.internal.operation.ServerVersionHelper.serverIsAtLeastVersionFourDotFour; -import static java.lang.String.format; -import static java.util.Collections.singletonList; - -class QueryBatchCursor implements AggregateResponseBatchCursor { - private static final Logger LOGGER = Loggers.getLogger("operation"); - private static final FieldNameValidator NO_OP_FIELD_NAME_VALIDATOR = new NoOpFieldNameValidator(); - private static final String CURSOR = "cursor"; - private static final String POST_BATCH_RESUME_TOKEN = "postBatchResumeToken"; - private static final String OPERATION_TIME = "operationTime"; - private static final String MESSAGE_IF_CLOSED_AS_CURSOR = "Cursor has been closed"; - private static final String MESSAGE_IF_CLOSED_AS_ITERATOR = "Iterator has been closed"; - - private final MongoNamespace namespace; - private final ServerAddress serverAddress; - private final int limit; - private final Decoder decoder; - private final long maxTimeMS; - private int batchSize; - private final BsonValue comment; - private List nextBatch; - private int count; - private BsonDocument postBatchResumeToken; - private BsonTimestamp operationTime; - private final boolean firstBatchEmpty; - private int maxWireVersion = 0; - private final ResourceManager resourceManager; - - QueryBatchCursor(final QueryResult firstQueryResult, final int limit, final int batchSize, final Decoder decoder) { - this(firstQueryResult, limit, batchSize, decoder, null, null); - } - - QueryBatchCursor(final QueryResult firstQueryResult, final int limit, final int batchSize, final Decoder decoder, - @Nullable final BsonValue comment, @Nullable final ConnectionSource connectionSource) { - this(firstQueryResult, limit, batchSize, 0, decoder, comment, connectionSource, null, null); - } - - QueryBatchCursor(final QueryResult firstQueryResult, final int limit, final int batchSize, final long maxTimeMS, - final Decoder decoder, @Nullable final BsonValue comment, @Nullable final ConnectionSource connectionSource, - @Nullable final Connection connection) { - this(firstQueryResult, limit, batchSize, maxTimeMS, decoder, comment, connectionSource, connection, null); - } - - QueryBatchCursor(final QueryResult firstQueryResult, final int limit, final int batchSize, final long maxTimeMS, - final Decoder decoder, @Nullable final BsonValue comment, @Nullable final ConnectionSource connectionSource, - @Nullable final Connection connection, @Nullable final BsonDocument result) { - isTrueArgument("maxTimeMS >= 0", maxTimeMS >= 0); - this.maxTimeMS = maxTimeMS; - this.namespace = firstQueryResult.getNamespace(); - this.serverAddress = firstQueryResult.getAddress(); - this.limit = limit; - this.comment = comment; - this.batchSize = batchSize; - this.decoder = notNull("decoder", decoder); - if (result != null) { - this.operationTime = result.getTimestamp(OPERATION_TIME, null); - this.postBatchResumeToken = getPostBatchResumeTokenFromResponse(result); - } - ServerCursor serverCursor = initFromQueryResult(firstQueryResult); - if (serverCursor != null) { - notNull("connectionSource", connectionSource); - } - firstBatchEmpty = firstQueryResult.getResults().isEmpty(); - Connection connectionToPin = null; - boolean releaseServerAndResources = false; - if (connection != null) { - this.maxWireVersion = connection.getDescription().getMaxWireVersion(); - if (limitReached()) { - releaseServerAndResources = true; - } else { - assertNotNull(connectionSource); - if (connectionSource.getServerDescription().getType() == ServerType.LOAD_BALANCER) { - connectionToPin = connection; - } - } - } - resourceManager = new ResourceManager(connectionSource, connectionToPin, serverCursor); - if (releaseServerAndResources) { - resourceManager.releaseServerAndClientResources(assertNotNull(connection)); - } - } - - @Override - public boolean hasNext() { - return assertNotNull(resourceManager.execute(MESSAGE_IF_CLOSED_AS_CURSOR, this::doHasNext)); - } - - private boolean doHasNext() { - if (nextBatch != null) { - return true; - } - - if (limitReached()) { - return false; - } - - while (resourceManager.serverCursor() != null) { - getMore(); - if (!resourceManager.operable()) { - throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_CURSOR); - } - if (nextBatch != null) { - return true; - } - } - - return false; - } - - @Override - public List next() { - return assertNotNull(resourceManager.execute(MESSAGE_IF_CLOSED_AS_ITERATOR, this::doNext)); - } - - @Override - public int available() { - return !resourceManager.operable() || nextBatch == null ? 0 : nextBatch.size(); - } - - private List doNext() { - if (!doHasNext()) { - throw new NoSuchElementException(); - } - - List retVal = nextBatch; - nextBatch = null; - return retVal; - } - - @Override - public void setBatchSize(final int batchSize) { - this.batchSize = batchSize; - } - - @Override - public int getBatchSize() { - return batchSize; - } - - @Override - public void remove() { - throw new UnsupportedOperationException("Not implemented yet!"); - } - - @Override - public void close() { - resourceManager.close(); - } - - @Nullable - @Override - public List tryNext() { - return resourceManager.execute(MESSAGE_IF_CLOSED_AS_CURSOR, () -> { - if (!tryHasNext()) { - return null; - } - return doNext(); - }); - } - - private boolean tryHasNext() { - if (nextBatch != null) { - return true; - } - - if (limitReached()) { - return false; - } - - if (resourceManager.serverCursor() != null) { - getMore(); - } - - return nextBatch != null; - } - - @Override - @Nullable - public ServerCursor getServerCursor() { - if (!resourceManager.operable()) { - throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_ITERATOR); - } - - return resourceManager.serverCursor(); - } - - @Override - public ServerAddress getServerAddress() { - if (!resourceManager.operable()) { - throw new IllegalStateException(MESSAGE_IF_CLOSED_AS_ITERATOR); - } - - return serverAddress; - } - - @Override - public BsonDocument getPostBatchResumeToken() { - return postBatchResumeToken; - } - - @Override - public BsonTimestamp getOperationTime() { - return operationTime; - } - - @Override - public boolean isFirstBatchEmpty() { - return firstBatchEmpty; - } - - @Override - public int getMaxWireVersion() { - return maxWireVersion; - } - - private void getMore() { - ServerCursor serverCursor = assertNotNull(resourceManager.serverCursor()); - resourceManager.executeWithConnection(connection -> { - ServerCursor nextServerCursor; - try { - nextServerCursor = initFromCommandResult(connection.command(namespace.getDatabaseName(), - asGetMoreCommandDocument(serverCursor.getId(), connection.getDescription()), - NO_OP_FIELD_NAME_VALIDATOR, - ReadPreference.primary(), - CommandResultDocumentCodec.create(decoder, "nextBatch"), - assertNotNull(resourceManager.connectionSource))); - } catch (MongoCommandException e) { - throw translateCommandException(e, serverCursor); - } - resourceManager.setServerCursor(nextServerCursor); - if (limitReached()) { - resourceManager.releaseServerAndClientResources(connection); - } - }); - } - - private BsonDocument asGetMoreCommandDocument(final long cursorId, final ConnectionDescription connectionDescription) { - BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursorId)) - .append("collection", new BsonString(namespace.getCollectionName())); - - int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count)); - if (batchSizeForGetMoreCommand != 0) { - document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand)); - } - if (maxTimeMS != 0) { - document.append("maxTimeMS", new BsonInt64(maxTimeMS)); - } - if (serverIsAtLeastVersionFourDotFour(connectionDescription)) { - putIfNotNull(document, "comment", comment); - } - return document; - } - - @Nullable - private ServerCursor initFromQueryResult(final QueryResult queryResult) { - nextBatch = queryResult.getResults().isEmpty() ? null : queryResult.getResults(); - count += queryResult.getResults().size(); - LOGGER.debug(format("Received batch of %d documents with cursorId %d from server %s", queryResult.getResults().size(), - queryResult.getCursorId(), queryResult.getAddress())); - return queryResult.getCursor(); - } - - @Nullable - private ServerCursor initFromCommandResult(final BsonDocument getMoreCommandResultDocument) { - QueryResult queryResult = getMoreCursorDocumentToQueryResult(getMoreCommandResultDocument.getDocument(CURSOR), serverAddress); - postBatchResumeToken = getPostBatchResumeTokenFromResponse(getMoreCommandResultDocument); - operationTime = getMoreCommandResultDocument.getTimestamp(OPERATION_TIME, null); - return initFromQueryResult(queryResult); - } - - private boolean limitReached() { - return Math.abs(limit) != 0 && count >= Math.abs(limit); - } - - @Nullable - private BsonDocument getPostBatchResumeTokenFromResponse(final BsonDocument result) { - BsonDocument cursor = result.getDocument(CURSOR, null); - if (cursor != null) { - return cursor.getDocument(POST_BATCH_RESUME_TOKEN, null); - } - return null; - } - - /** - * This class maintains all resources that must be released in {@link QueryBatchCursor#close()}. - * It also implements a {@linkplain #doClose() deferred close action} such that it is totally ordered with other operations of - * {@link QueryBatchCursor} (methods {@link #tryStartOperation()}/{@link #endOperation()} must be used properly to enforce the order) - * despite the method {@link QueryBatchCursor#close()} being called concurrently with those operations. - * This total order induces the happens-before order. - *

    - * The deferred close action does not violate externally observable idempotence of {@link QueryBatchCursor#close()}, - * because {@link QueryBatchCursor#close()} is allowed to release resources "eventually". - *

    - * Only methods explicitly documented as thread-safe are thread-safe, - * others are not and rely on the total order mentioned above. - */ - @ThreadSafe - private final class ResourceManager { - private final Lock lock; - private volatile State state; - @Nullable - private volatile ConnectionSource connectionSource; - @Nullable - private volatile Connection pinnedConnection; - @Nullable - private volatile ServerCursor serverCursor; - private volatile boolean skipReleasingServerResourcesOnClose; - - ResourceManager(@Nullable final ConnectionSource connectionSource, - @Nullable final Connection connectionToPin, @Nullable final ServerCursor serverCursor) { - lock = new StampedLock().asWriteLock(); - state = State.IDLE; - if (serverCursor != null) { - this.connectionSource = (assertNotNull(connectionSource)).retain(); - if (connectionToPin != null) { - this.pinnedConnection = connectionToPin.retain(); - connectionToPin.markAsPinned(Connection.PinningMode.CURSOR); - } - } - skipReleasingServerResourcesOnClose = false; - this.serverCursor = serverCursor; - } - - /** - * Thread-safe. - */ - boolean operable() { - return state.operable(); - } - - /** - * Thread-safe. - * Executes {@code operation} within the {@link #tryStartOperation()}/{@link #endOperation()} bounds. - * - * @throws IllegalStateException If {@linkplain QueryBatchCursor#close() closed}. - */ - @Nullable - R execute(final String exceptionMessageIfClosed, final Supplier operation) throws IllegalStateException { - if (!tryStartOperation()) { - throw new IllegalStateException(exceptionMessageIfClosed); - } - try { - return operation.get(); - } finally { - endOperation(); - } - } - - /** - * Thread-safe. - * Returns {@code true} iff started an operation. - * If {@linkplain #operable() closed}, then returns false, otherwise completes abruptly. - * @throws IllegalStateException Iff another operation is in progress. - */ - private boolean tryStartOperation() throws IllegalStateException { - lock.lock(); - try { - State localState = state; - if (!localState.operable()) { - return false; - } else if (localState == State.IDLE) { - state = State.OPERATION_IN_PROGRESS; - return true; - } else if (localState == State.OPERATION_IN_PROGRESS) { - throw new IllegalStateException("Another operation is currently in progress, concurrent operations are not supported"); - } else { - throw fail(state.toString()); - } - } finally { - lock.unlock(); - } - } - - /** - * Thread-safe. - */ - private void endOperation() { - boolean doClose = false; - lock.lock(); - try { - State localState = state; - if (localState == State.OPERATION_IN_PROGRESS) { - state = State.IDLE; - } else if (localState == State.CLOSE_PENDING) { - state = State.CLOSED; - doClose = true; - } else { - fail(localState.toString()); - } - } finally { - lock.unlock(); - } - if (doClose) { - doClose(); - } - } - - /** - * Thread-safe. - */ - void close() { - boolean doClose = false; - lock.lock(); - try { - State localState = state; - if (localState == State.OPERATION_IN_PROGRESS) { - state = State.CLOSE_PENDING; - } else if (localState != State.CLOSED) { - state = State.CLOSED; - doClose = true; - } - } finally { - lock.unlock(); - } - if (doClose) { - doClose(); - } - } - - /** - * This method is never executed concurrently with either itself or other operations - * demarcated by {@link #tryStartOperation()}/{@link #endOperation()}. - */ - private void doClose() { - try { - if (skipReleasingServerResourcesOnClose) { - serverCursor = null; - } else if (serverCursor != null) { - Connection connection = connection(); - try { - releaseServerResources(connection); - } finally { - connection.release(); - } - } - } catch (MongoException e) { - // ignore exceptions when releasing server resources - } finally { - // guarantee that regardless of exceptions, `serverCursor` is null and client resources are released - serverCursor = null; - releaseClientResources(); - } - } - - void onCorruptedConnection(final Connection corruptedConnection) { - assertTrue(state.inProgress()); - // if `pinnedConnection` is corrupted, then we cannot kill `serverCursor` via such a connection - Connection localPinnedConnection = pinnedConnection; - if (localPinnedConnection != null) { - assertTrue(corruptedConnection == localPinnedConnection); - skipReleasingServerResourcesOnClose = true; - } - } - - void executeWithConnection(final Consumer action) { - Connection connection = connection(); - try { - action.accept(connection); - } catch (MongoSocketException e) { - try { - onCorruptedConnection(connection); - } catch (Exception suppressed) { - e.addSuppressed(suppressed); - } - throw e; - } finally { - connection.release(); - } - } - - private Connection connection() { - assertTrue(state != State.IDLE); - if (pinnedConnection == null) { - return assertNotNull(connectionSource).getConnection(); - } else { - return assertNotNull(pinnedConnection).retain(); - } - } - - /** - * Thread-safe. - */ - @Nullable - ServerCursor serverCursor() { - return serverCursor; - } - - void setServerCursor(@Nullable final ServerCursor serverCursor) { - assertTrue(state.inProgress()); - assertNotNull(this.serverCursor); - // without `connectionSource` we will not be able to kill `serverCursor` later - assertNotNull(connectionSource); - this.serverCursor = serverCursor; - if (serverCursor == null) { - releaseClientResources(); - } - } - - - void releaseServerAndClientResources(final Connection connection) { - try { - releaseServerResources(assertNotNull(connection)); - } finally { - releaseClientResources(); - } - } - - private void releaseServerResources(final Connection connection) { - try { - ServerCursor localServerCursor = serverCursor; - if (localServerCursor != null) { - killServerCursor(namespace, localServerCursor, assertNotNull(connection)); - } - } finally { - serverCursor = null; - } - } - - private void killServerCursor(final MongoNamespace namespace, final ServerCursor serverCursor, final Connection connection) { - connection.command(namespace.getDatabaseName(), asKillCursorsCommandDocument(namespace, serverCursor), - NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), assertNotNull(connectionSource)); - } - - private BsonDocument asKillCursorsCommandDocument(final MongoNamespace namespace, final ServerCursor serverCursor) { - return new BsonDocument("killCursors", new BsonString(namespace.getCollectionName())) - .append("cursors", new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))); - } - - private void releaseClientResources() { - assertNull(serverCursor); - ConnectionSource localConnectionSource = connectionSource; - if (localConnectionSource != null) { - localConnectionSource.release(); - connectionSource = null; - } - Connection localPinnedConnection = pinnedConnection; - if (localPinnedConnection != null) { - localPinnedConnection.release(); - pinnedConnection = null; - } - } - } - - private enum State { - IDLE(true, false), - OPERATION_IN_PROGRESS(true, true), - /** - * Implies {@link #OPERATION_IN_PROGRESS}. - */ - CLOSE_PENDING(false, true), - CLOSED(false, false); - - private final boolean operable; - private final boolean inProgress; - - State(final boolean operable, final boolean inProgress) { - this.operable = operable; - this.inProgress = inProgress; - } - - boolean operable() { - return operable; - } - - boolean inProgress() { - return inProgress; - } - } -} diff --git a/driver-core/src/main/com/mongodb/internal/operation/QueryHelper.java b/driver-core/src/main/com/mongodb/internal/operation/QueryHelper.java deleted file mode 100644 index 053e4fc8817..00000000000 --- a/driver-core/src/main/com/mongodb/internal/operation/QueryHelper.java +++ /dev/null @@ -1,35 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation; - -import com.mongodb.MongoCommandException; -import com.mongodb.MongoCursorNotFoundException; -import com.mongodb.MongoQueryException; -import com.mongodb.ServerCursor; - -final class QueryHelper { - static MongoQueryException translateCommandException(final MongoCommandException commandException, final ServerCursor cursor) { - if (commandException.getErrorCode() == 43) { - return new MongoCursorNotFoundException(cursor.getId(), commandException.getResponse(), cursor.getAddress()); - } else { - return new MongoQueryException(commandException.getResponse(), commandException.getServerAddress()); - } - } - - private QueryHelper() { - } -} diff --git a/driver-core/src/main/com/mongodb/internal/operation/SingleBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/SingleBatchCursor.java new file mode 100644 index 00000000000..8a673ee93d9 --- /dev/null +++ b/driver-core/src/main/com/mongodb/internal/operation/SingleBatchCursor.java @@ -0,0 +1,91 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.mongodb.internal.operation; + +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; + +import java.util.List; +import java.util.NoSuchElementException; + +import static java.util.Collections.emptyList; + +class SingleBatchCursor implements BatchCursor { + + static SingleBatchCursor createEmptySingleBatchCursor(final ServerAddress serverAddress, final int batchSize) { + return new SingleBatchCursor<>(emptyList(), batchSize, serverAddress); + } + + private final List batch; + private final ServerAddress serverAddress; + private final int batchSize; + private boolean hasNext; + + SingleBatchCursor(final List batch, final int batchSize, final ServerAddress serverAddress) { + this.batch = batch; + this.serverAddress = serverAddress; + this.batchSize = batchSize; + this.hasNext = !batch.isEmpty(); + } + + @Override + public boolean hasNext() { + return hasNext; + } + + @Override + public List next() { + if (hasNext) { + hasNext = false; + return batch; + } + throw new NoSuchElementException(); + } + + @Override + public int available() { + return hasNext ? 1 : 0; + } + + @Override + public void setBatchSize(final int batchSize) { + // NOOP + } + + @Override + public int getBatchSize() { + return batchSize; + } + + @Override + public List tryNext() { + return hasNext ? next() : null; + } + + @Override + public ServerCursor getServerCursor() { + return null; + } + + @Override + public ServerAddress getServerAddress() { + return serverAddress; + } + + @Override + public void close() { + } +} diff --git a/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java b/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java index 67d5acf9c37..476a2676b8e 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java +++ b/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java @@ -18,7 +18,6 @@ import com.mongodb.MongoException; import com.mongodb.ReadPreference; -import com.mongodb.ServerAddress; import com.mongodb.internal.VisibleForTesting; import com.mongodb.internal.async.SingleResultCallback; import com.mongodb.internal.async.function.AsyncCallbackBiFunction; @@ -32,7 +31,6 @@ import com.mongodb.internal.binding.WriteBinding; import com.mongodb.internal.connection.Connection; import com.mongodb.internal.connection.OperationContext; -import com.mongodb.internal.connection.QueryResult; import com.mongodb.internal.operation.retry.AttachmentKeys; import com.mongodb.internal.validator.NoOpFieldNameValidator; import com.mongodb.lang.Nullable; @@ -56,7 +54,6 @@ import static com.mongodb.internal.operation.OperationHelper.ResourceSupplierInternalException; import static com.mongodb.internal.operation.OperationHelper.canRetryRead; import static com.mongodb.internal.operation.OperationHelper.canRetryWrite; -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult; import static com.mongodb.internal.operation.WriteConcernHelper.throwOnWriteConcernError; final class SyncOperationHelper { @@ -303,14 +300,16 @@ static CommandWriteTransformer writeConcernErrorTransformer( }; } - static BatchCursor cursorDocumentToBatchCursor(final BsonDocument cursorDocument, final Decoder decoder, - final BsonValue comment, final ConnectionSource source, final Connection connection, final int batchSize) { - return new QueryBatchCursor<>(cursorDocumentToQueryResult(cursorDocument, source.getServerDescription().getAddress()), - 0, batchSize, 0, decoder, comment, source, connection); + static CommandReadTransformer> singleBatchCursorTransformer(final String resultsFieldName) { + return (result, source, connection) -> + new SingleBatchCursor<>(BsonDocumentWrapperHelper.toList(result, resultsFieldName), 0, + connection.getDescription().getServerAddress()); } - static QueryResult getMoreCursorDocumentToQueryResult(final BsonDocument cursorDocument, final ServerAddress serverAddress) { - return cursorDocumentToQueryResult(cursorDocument, serverAddress, "nextBatch"); + static BatchCursor cursorDocumentToBatchCursor(final BsonDocument cursorDocument, final Decoder decoder, + final BsonValue comment, final ConnectionSource source, final Connection connection, final int batchSize) { + return new CommandBatchCursor<>(connection.getDescription().getServerAddress(), cursorDocument, 0, batchSize, + 0, decoder, comment, source, connection); } private SyncOperationHelper() { diff --git a/driver-core/src/test/functional/com/mongodb/OperationFunctionalSpecification.groovy b/driver-core/src/test/functional/com/mongodb/OperationFunctionalSpecification.groovy index ddbb9f29a0d..372fdd4b82d 100644 --- a/driver-core/src/test/functional/com/mongodb/OperationFunctionalSpecification.groovy +++ b/driver-core/src/test/functional/com/mongodb/OperationFunctionalSpecification.groovy @@ -202,13 +202,6 @@ class OperationFunctionalSpecification extends Specification { } } - def consumeAsyncResults(cursor) { - def batch = next(cursor, true) - while (batch != null) { - batch = next(cursor, true) - } - } - void testOperation(Map params) { params.async = params.async != null ? params.async : false params.result = params.result != null ? params.result : null diff --git a/driver-core/src/test/functional/com/mongodb/client/model/OperationTest.java b/driver-core/src/test/functional/com/mongodb/client/model/OperationTest.java index 9a215c7260c..75ec2c60e2b 100644 --- a/driver-core/src/test/functional/com/mongodb/client/model/OperationTest.java +++ b/driver-core/src/test/functional/com/mongodb/client/model/OperationTest.java @@ -18,14 +18,17 @@ import com.mongodb.ClusterFixture; import com.mongodb.MongoNamespace; +import com.mongodb.async.FutureResultCallback; import com.mongodb.client.test.CollectionHelper; import com.mongodb.internal.connection.ServerHelper; +import com.mongodb.internal.validator.NoOpFieldNameValidator; import com.mongodb.lang.Nullable; import org.bson.BsonArray; import org.bson.BsonDocument; import org.bson.BsonDouble; import org.bson.BsonValue; import org.bson.Document; +import org.bson.FieldNameValidator; import org.bson.codecs.BsonDocumentCodec; import org.bson.codecs.DecoderContext; import org.bson.codecs.DocumentCodec; @@ -39,8 +42,11 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import java.util.stream.Collectors; +import static com.mongodb.ClusterFixture.TIMEOUT; import static com.mongodb.ClusterFixture.checkReferenceCountReachesTarget; import static com.mongodb.ClusterFixture.getAsyncBinding; import static com.mongodb.ClusterFixture.getBinding; @@ -54,6 +60,7 @@ public abstract class OperationTest { protected static final DocumentCodec DOCUMENT_DECODER = new DocumentCodec(); + protected static final FieldNameValidator NO_OP_FIELD_NAME_VALIDATOR = new NoOpFieldNameValidator(); @BeforeEach public void beforeEach() { @@ -77,15 +84,15 @@ private CollectionHelper getCollectionHelper(final MongoNamespace return new CollectionHelper<>(new BsonDocumentCodec(), namespace); } - private String getDatabaseName() { + protected String getDatabaseName() { return ClusterFixture.getDefaultDatabaseName(); } - private String getCollectionName() { + protected String getCollectionName() { return "test"; } - MongoNamespace getNamespace() { + protected MongoNamespace getNamespace() { return new MongoNamespace(getDatabaseName(), getCollectionName()); } @@ -97,7 +104,6 @@ public static BsonDocument toBsonDocument(final BsonDocument bsonDocument) { return getDefaultCodecRegistry().get(BsonDocument.class).decode(bsonDocument.asBsonReader(), DecoderContext.builder().build()); } - protected List assertPipeline(final String stageAsString, final Bson stage) { List pipeline = Collections.singletonList(stage); return assertPipeline(stageAsString, pipeline); @@ -159,4 +165,25 @@ protected List aggregateWithWindowFields(@Nullable final Object partitio .map(doc -> doc.get("result")) .collect(toList()); } + + protected void ifNotNull(@Nullable final T maybeNull, final Consumer consumer) { + if (maybeNull != null) { + consumer.accept(maybeNull); + } + } + + protected void sleep(final long ms) { + try { + Thread.sleep(ms); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + protected T block(final Consumer> consumer) { + FutureResultCallback cb = new FutureResultCallback<>(); + consumer.accept(cb); + return cb.get(TIMEOUT, TimeUnit.SECONDS); + } } diff --git a/driver-core/src/test/functional/com/mongodb/internal/connection/ServerHelper.java b/driver-core/src/test/functional/com/mongodb/internal/connection/ServerHelper.java index ecbf4befb73..ef0f99966a0 100644 --- a/driver-core/src/test/functional/com/mongodb/internal/connection/ServerHelper.java +++ b/driver-core/src/test/functional/com/mongodb/internal/connection/ServerHelper.java @@ -64,9 +64,20 @@ public static void waitForLastRelease(final ServerAddress address, final Cluster private static void checkPool(final ServerAddress address, final Cluster cluster) { ConcurrentPool pool = connectionPool( cluster.selectServer(new ServerAddressSelector(address), new OperationContext()).getServer()); - if (pool.getInUseCount() > 0) { - throw new IllegalStateException("Connection pool in use count is " + pool.getInUseCount()); + + int counter = 0; + while (counter < 5) { + if (pool.getInUseCount() == 0) { + return; + } + try { + Thread.sleep(500); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + counter++; } + throw new IllegalStateException("Connection pool in use count is " + pool.getInUseCount()); } private static ConcurrentPool connectionPool(final Server server) { diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncCommandBatchCursorFunctionalTest.java b/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncCommandBatchCursorFunctionalTest.java new file mode 100644 index 00000000000..09c3a51be25 --- /dev/null +++ b/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncCommandBatchCursorFunctionalTest.java @@ -0,0 +1,440 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + + +import com.mongodb.MongoCursorNotFoundException; +import com.mongodb.MongoException; +import com.mongodb.MongoQueryException; +import com.mongodb.ReadPreference; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; +import com.mongodb.client.model.CreateCollectionOptions; +import com.mongodb.client.model.OperationTest; +import com.mongodb.internal.binding.AsyncConnectionSource; +import com.mongodb.internal.connection.AsyncConnection; +import org.bson.BsonArray; +import org.bson.BsonBoolean; +import org.bson.BsonDocument; +import org.bson.BsonInt32; +import org.bson.BsonInt64; +import org.bson.BsonString; +import org.bson.BsonTimestamp; +import org.bson.Document; +import org.bson.codecs.BsonDocumentCodec; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static com.mongodb.ClusterFixture.checkReferenceCountReachesTarget; +import static com.mongodb.ClusterFixture.getAsyncBinding; +import static com.mongodb.ClusterFixture.getConnection; +import static com.mongodb.ClusterFixture.getReferenceCountAfterTimeout; +import static com.mongodb.ClusterFixture.getWriteConnectionSource; +import static com.mongodb.ClusterFixture.isSharded; +import static com.mongodb.internal.operation.QueryOperationHelper.makeAdditionalGetMoreCall; +import static java.util.Collections.singletonList; +import static java.util.stream.Stream.generate; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeFalse; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +public class AsyncCommandBatchCursorFunctionalTest extends OperationTest { + + private AsyncConnectionSource connectionSource; + private AsyncConnection connection; + private AsyncCommandBatchCursor cursor; + + @BeforeEach + void setup() throws Throwable { + List documents = IntStream.rangeClosed(1, 10) + .mapToObj(i -> new BsonDocument("i", new BsonInt32(i))) + .collect(Collectors.toList()); + getCollectionHelper().insertDocuments(documents); + + connectionSource = getWriteConnectionSource(getAsyncBinding()); + connection = getConnection(connectionSource); + } + + @AfterEach + void cleanup() { + ifNotNull(cursor, AsyncCommandBatchCursor::close); + ifNotNull(connectionSource, cs -> { + getReferenceCountAfterTimeout(cs, 1); + cs.release(); + }); + ifNotNull(connection, c -> { + getReferenceCountAfterTimeout(c, 1); + c.release(); + }); + } + + @Test + @DisplayName("server cursor should not be null") + void theServerCursorShouldNotBeNull() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.getServerCursor()); + } + + + @Test + @DisplayName("should get Exceptions for operations on the cursor after closing") + void shouldGetExceptionsForOperationsOnTheCursorAfterClosing() { + BsonDocument commandResult = executeFindCommand(5); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + cursor.close(); + assertDoesNotThrow(() -> cursor.close()); + + checkReferenceCountReachesTarget(connectionSource, 1); + assertThrows(MongoException.class, this::cursorNext); + assertNull(cursor.getServerCursor()); + } + + @Test + @DisplayName("should throw an Exception when going off the end") + void shouldThrowAnExceptionWhenGoingOffTheEnd() { + BsonDocument commandResult = executeFindCommand(1); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 2, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + cursorNext(); + cursorNext(); + + assertThrows(MongoException.class, this::cursorNext); + } + + + @Test + @DisplayName("test normal exhaustion") + void testNormalExhaustion() { + BsonDocument commandResult = executeFindCommand(); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(10, cursorFlatten().size()); + } + + @ParameterizedTest(name = "{index} => limit={0}, batchSize={1}, expectedTotal={2}") + @MethodSource + @DisplayName("test limit exhaustion") + void testLimitExhaustion(final int limit, final int batchSize, final int expectedTotal) { + BsonDocument commandResult = executeFindCommand(limit, batchSize); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, limit, batchSize, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + + assertEquals(expectedTotal, cursorFlatten().size()); + + checkReferenceCountReachesTarget(connectionSource, 1); + checkReferenceCountReachesTarget(connection, 1); + } + + @ParameterizedTest(name = "{index} => awaitData={0}, maxTimeMS={1}") + @MethodSource + @DisplayName("should block waiting for next batch on a tailable cursor") + void shouldBlockWaitingForNextBatchOnATailableCursor(final boolean awaitData, final int maxTimeMS) { + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, awaitData); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, maxTimeMS, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertFalse(cursor.isClosed()); + assertEquals(1, cursorNext().get(0).get("_id")); + + new Thread(() -> { + sleep(100); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 2).append("ts", new BsonTimestamp(6, 0))); + }).start(); + + assertFalse(cursor.isClosed()); + assertEquals(2, cursorNext().get(0).get("_id")); + } + + @Test + @DisplayName("test tailable interrupt") + void testTailableInterrupt() throws InterruptedException { + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, true); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + CountDownLatch latch = new CountDownLatch(1); + AtomicInteger seen = new AtomicInteger(); + Thread thread = new Thread(() -> { + try { + cursorNext(); + seen.incrementAndGet(); + cursorNext(); + seen.incrementAndGet(); + } catch (Exception e) { + // pass + } finally { + latch.countDown(); + } + }); + + thread.start(); + sleep(1000); + thread.interrupt(); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 2)); + latch.await(); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + assertEquals(1, seen.intValue()); + } + + @Test + @DisplayName("should kill cursor if limit is reached on initial query") + void shouldKillCursorIfLimitIsReachedOnInitialQuery() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(5); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 5, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursorNext()); + assertTrue(cursor.isClosed()); + assertNull(cursor.getServerCursor()); + } + + @Test + @DisplayName("should kill cursor if limit is reached on getMore") + void shouldKillCursorIfLimitIsReachedOnGetMore() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 5, 3, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + ServerCursor serverCursor = cursor.getServerCursor(); + assertNotNull(serverCursor); + assertNotNull(cursorNext()); + assertNotNull(cursorNext()); + + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + assertThrows(MongoQueryException.class, () -> + makeAdditionalGetMoreCall(getNamespace(), serverCursor, connection) + ); + } + + @Test + @DisplayName("should release connection source if limit is reached on initial query") + void shouldReleaseConnectionSourceIfLimitIsReachedOnInitialQuery() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(5); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 5, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNull(cursor.getServerCursor()); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connectionSource, 1)); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + } + + @Test + @DisplayName("should release connection source if limit is reached on getMore") + void shouldReleaseConnectionSourceIfLimitIsReachedOnGetMore() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(3); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 5, 3, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursorNext()); + assertNotNull(cursorNext()); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connectionSource, 1)); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + } + + @Test + @DisplayName("test limit with get more") + void testLimitWithGetMore() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 5, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursorNext()); + assertNotNull(cursorNext()); + assertNotNull(cursorNext()); + + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connectionSource, 1)); + assertTrue(cursor.isClosed()); + } + + @Test + @DisplayName("test limit with large documents") + void testLimitWithLargeDocuments() { + String bigString = generate(() -> "x") + .limit(16000) + .collect(Collectors.joining()); + + IntStream.range(11, 1000).forEach(i -> + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", i).append("s", bigString)) + ); + + BsonDocument commandResult = executeFindCommand(300, 0); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 300, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(300, cursorFlatten().size()); + } + + @Test + @DisplayName("should respect batch size") + void shouldRespectBatchSize() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(2, cursor.getBatchSize()); + assertEquals(2, cursorNext().size()); + assertEquals(2, cursorNext().size()); + + cursor.setBatchSize(3); + assertEquals(3, cursor.getBatchSize()); + assertEquals(3, cursorNext().size()); + assertEquals(3, cursorNext().size()); + } + + @Test + @DisplayName("should throw cursor not found exception") + void shouldThrowCursorNotFoundException() throws Throwable { + BsonDocument commandResult = executeFindCommand(2); + cursor = new AsyncCommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + ServerCursor serverCursor = cursor.getServerCursor(); + assertNotNull(serverCursor); + AsyncConnection localConnection = getConnection(connectionSource); + this.block(cb -> localConnection.commandAsync(getNamespace().getDatabaseName(), + new BsonDocument("killCursors", new BsonString(getNamespace().getCollectionName())) + .append("cursors", new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))), + NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), connectionSource, cb)); + localConnection.release(); + + cursorNext(); + + MongoCursorNotFoundException exception = assertThrows(MongoCursorNotFoundException.class, this::cursorNext); + assertEquals(serverCursor.getId(), exception.getCursorId()); + assertEquals(serverCursor.getAddress(), exception.getServerAddress()); + } + + + private static Stream shouldBlockWaitingForNextBatchOnATailableCursor() { + return Stream.of( + arguments(true, 0), + arguments(true, 100), + arguments(false, 0)); + } + + private static Stream testLimitExhaustion() { + return Stream.of( + arguments(5, 2, 5), + arguments(5, -2, 2), + arguments(-5, -2, 5), + arguments(-5, 2, 5), + arguments(2, 5, 2), + arguments(2, -5, 2), + arguments(-2, 5, 2), + arguments(-2, -5, 2) + ); + } + + private ServerAddress getServerAddress() { + return connection.getDescription().getServerAddress(); + } + + private BsonDocument executeFindCommand() { + return executeFindCommand(0); + } + + private BsonDocument executeFindCommand(final int batchSize) { + return executeFindCommand(new BsonDocument(), 0, batchSize, false, false); + } + + private BsonDocument executeFindCommand(final int limit, final int batchSize) { + return executeFindCommand(new BsonDocument(), limit, batchSize, false, false); + } + + private BsonDocument executeFindCommand(final BsonDocument filter, final int limit, final int batchSize, final boolean tailable, + final boolean awaitData) { + return executeFindCommand(filter, limit, batchSize, tailable, awaitData, ReadPreference.primary()); + } + + private BsonDocument executeFindCommand(final BsonDocument filter, final int limit, final int batchSize, + final boolean tailable, final boolean awaitData, final ReadPreference readPreference) { + BsonDocument findCommand = new BsonDocument("find", new BsonString(getCollectionName())) + .append("filter", filter) + .append("tailable", BsonBoolean.valueOf(tailable)) + .append("awaitData", BsonBoolean.valueOf(awaitData)); + + findCommand.append("limit", new BsonInt32(Math.abs(limit))); + if (limit >= 0) { + if (batchSize < 0 && Math.abs(batchSize) < limit) { + findCommand.append("limit", new BsonInt32(Math.abs(batchSize))); + } else { + findCommand.append("batchSize", new BsonInt32(Math.abs(batchSize))); + } + } + + BsonDocument results = block(cb -> connection.commandAsync(getDatabaseName(), findCommand, + NO_OP_FIELD_NAME_VALIDATOR, readPreference, + CommandResultDocumentCodec.create(DOCUMENT_DECODER, "firstBatch"), + connectionSource, cb)); + + assertNotNull(results); + return results; + } + + private List cursorNext() { + return block(cb -> cursor.next(cb)); + } + + private List cursorFlatten() { + List results = new ArrayList<>(); + while (!cursor.isClosed()) { + results.addAll(cursorNext()); + } + return results; + } +} diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncQueryBatchCursorFunctionalSpecification.groovy b/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncQueryBatchCursorFunctionalSpecification.groovy deleted file mode 100644 index 3d6f0c8b7a7..00000000000 --- a/driver-core/src/test/functional/com/mongodb/internal/operation/AsyncQueryBatchCursorFunctionalSpecification.groovy +++ /dev/null @@ -1,445 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation - -import com.mongodb.MongoCursorNotFoundException -import com.mongodb.MongoException -import com.mongodb.MongoTimeoutException -import com.mongodb.OperationFunctionalSpecification -import com.mongodb.ReadPreference -import com.mongodb.ServerCursor -import com.mongodb.WriteConcern -import com.mongodb.async.FutureResultCallback -import com.mongodb.client.model.CreateCollectionOptions -import com.mongodb.client.syncadapter.SyncConnection -import com.mongodb.internal.binding.AsyncConnectionSource -import com.mongodb.internal.binding.AsyncReadBinding -import com.mongodb.internal.connection.AsyncConnection -import com.mongodb.internal.connection.Connection -import com.mongodb.internal.connection.QueryResult -import com.mongodb.internal.validator.NoOpFieldNameValidator -import org.bson.BsonArray -import org.bson.BsonBoolean -import org.bson.BsonDocument -import org.bson.BsonInt32 -import org.bson.BsonInt64 -import org.bson.BsonNull -import org.bson.BsonString -import org.bson.BsonTimestamp -import org.bson.Document -import org.bson.codecs.BsonDocumentCodec -import org.bson.codecs.DocumentCodec -import spock.lang.IgnoreIf -import util.spock.annotations.Slow - -import java.util.concurrent.CountDownLatch - -import static com.mongodb.ClusterFixture.getAsyncBinding -import static com.mongodb.ClusterFixture.getAsyncCluster -import static com.mongodb.ClusterFixture.getBinding -import static com.mongodb.ClusterFixture.getConnection -import static com.mongodb.ClusterFixture.getReadConnectionSource -import static com.mongodb.ClusterFixture.getReferenceCountAfterTimeout -import static com.mongodb.ClusterFixture.isDiscoverableReplicaSet -import static com.mongodb.ClusterFixture.isSharded -import static com.mongodb.ClusterFixture.serverVersionLessThan -import static com.mongodb.internal.connection.ServerHelper.waitForLastRelease -import static com.mongodb.internal.connection.ServerHelper.waitForRelease -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult -import static com.mongodb.internal.operation.QueryOperationHelper.makeAdditionalGetMoreCall -import static java.util.Collections.singletonList -import static java.util.concurrent.TimeUnit.SECONDS -import static org.junit.Assert.assertEquals -import static org.junit.Assert.fail - -@IgnoreIf({ isSharded() && serverVersionLessThan(3, 2) }) -class AsyncQueryBatchCursorFunctionalSpecification extends OperationFunctionalSpecification { - AsyncConnectionSource connectionSource - AsyncQueryBatchCursor cursor - AsyncConnection connection - - def setup() { - def documents = [] - for (int i = 0; i < 10; i++) { - documents.add(new BsonDocument('_id', new BsonInt32(i))) - } - collectionHelper.insertDocuments(documents, - isDiscoverableReplicaSet() ? WriteConcern.MAJORITY : WriteConcern.ACKNOWLEDGED, - getBinding()) - setUpConnectionAndSource(getAsyncBinding()) - } - - private void setUpConnectionAndSource(final AsyncReadBinding binding) { - connectionSource = getReadConnectionSource(binding) - connection = getConnection(connectionSource) - } - - def cleanup() { - cursor?.close() - cleanupConnectionAndSource() - } - - private void cleanupConnectionAndSource() { - connection?.release() - connectionSource?.release() - waitForLastRelease(connectionSource.getServerDescription().getAddress(), getAsyncCluster()) - waitForRelease(connectionSource, 0) - } - - def 'should exhaust single batch'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(), 0, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - expect: - nextBatch().size() == 10 - } - - def 'should not retain connection and source after cursor is exhausted on first batch'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(), 0, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - nextBatch() - - then: - connection.count == 1 - connectionSource.count == 1 - } - - def 'should not retain connection and source after cursor is exhausted on getMore'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(1, 0), 1, 1, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - nextBatch() - - then: - getReferenceCountAfterTimeout(connection, 1) == 1 - getReferenceCountAfterTimeout(connectionSource, 1) == 1 - } - - def 'should not retain connection and source after cursor is exhausted after first batch'() { - when: - cursor = new AsyncQueryBatchCursor(executeQuery(10, 10), 10, 10, 0, new DocumentCodec(), null, connectionSource, - connection) - - then: - getReferenceCountAfterTimeout(connection, 1) == 1 - getReferenceCountAfterTimeout(connectionSource, 1) == 1 - } - - def 'should exhaust single batch with limit'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(1, 0), 1, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - expect: - nextBatch().size() == 1 - cursor.isClosed() || !nextBatch() && cursor.isClosed() - } - - def 'should exhaust multiple batches with limit'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(limit, batchSize), limit, batchSize, 0, new DocumentCodec(), null, - connectionSource, connection) - - when: - def next = nextBatch() - def total = 0 - while (next) { - total += next.size() - if (cursor.isClosed()) { - break - } - next = nextBatch() - } - - then: - total == expectedTotal - - where: - limit | batchSize | expectedTotal - 5 | 2 | 5 - 5 | -2 | 2 - -5 | 2 | 5 - -5 | -2 | 5 - 2 | 5 | 2 - 2 | -5 | 2 - -2 | 5 | 2 - -2 | -5 | 2 - } - - def 'should exhaust multiple batches'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(3), 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - expect: - nextBatch().size() == 3 - nextBatch().size() == 2 - nextBatch().size() == 2 - nextBatch().size() == 2 - nextBatch().size() == 1 - !nextBatch() - } - - def 'should respect batch size'() { - when: - cursor = new AsyncQueryBatchCursor(executeQuery(3), 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - then: - cursor.batchSize == 2 - - when: - nextBatch() - cursor.batchSize = 4 - - then: - nextBatch().size() == 4 - } - - def 'should close when exhausted'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(), 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - cursor.close() - waitForRelease(connectionSource, 1) - - then: - connectionSource.count == 1 - - when: - nextBatch() - - then: - thrown(MongoException) - } - - def 'should close when not exhausted'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(3), 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - cursor.close() - - then: - waitForRelease(connectionSource, 1) - } - - @Slow - def 'should block waiting for first batch on a tailable cursor'() { - given: - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(4, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, false) - - when: - cursor = new AsyncQueryBatchCursor(firstBatch, 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - def latch = new CountDownLatch(1) - Thread.start { - sleep(500) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 2).append('ts', new BsonTimestamp(5, 0))) - latch.countDown() - } - - def batch = nextBatch() - - then: - batch.size() == 1 - batch[0].get('_id') == 2 - - cleanup: - def cleanedUp = latch.await(10, SECONDS) // Workaround for codenarc bug - if (!cleanedUp) { - throw new MongoTimeoutException('Timed out waiting for documents to be inserted') - } - } - - @Slow - def 'should block waiting for next batch on a tailable cursor'() { - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(5, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, awaitData) - - - when: - cursor = new AsyncQueryBatchCursor(firstBatch, 0, 2, maxTimeMS, new DocumentCodec(), null, connectionSource, connection) - def batch = nextBatch() - - then: - batch.size() == 1 - batch[0].get('_id') == 1 - - when: - def latch = new CountDownLatch(1) - Thread.start { - sleep(500) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 2).append('ts', new BsonTimestamp(6, 0))) - latch.countDown() - } - - batch = nextBatch() - - then: - batch.size() == 1 - batch[0].get('_id') == 2 - - cleanup: - def cleanedUp = latch.await(10, SECONDS) - if (!cleanedUp) { - throw new MongoTimeoutException('Timed out waiting for documents to be inserted') - } - - where: - awaitData | maxTimeMS - true | 0 - true | 100 - false | 0 - } - - @Slow - def 'should unblock if closed while waiting for more data from tailable cursor'() { - given: - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), Document.parse('{}')) - def firstBatch = executeQuery(new BsonDocument('_id', BsonNull.VALUE), 0, 1, true, true) - - when: - cursor = new AsyncQueryBatchCursor(firstBatch, 0, 1, 500, new DocumentCodec(), null, connectionSource, connection) - Thread.start { - Thread.sleep(SECONDS.toMillis(2)) - cursor.close() - } - def batch = nextBatch() - - then: - cursor.isClosed() - batch == null - //both connection and connectionSource have reference count 1 when we pass them to the AsyncQueryBatchCursor constructor - connection.getCount() == 1 - waitForRelease(connectionSource, 1) - } - - def 'should respect limit'() { - given: - cursor = new AsyncQueryBatchCursor(executeQuery(6, 3), 6, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - expect: - nextBatch().size() == 3 - nextBatch().size() == 2 - nextBatch().size() == 1 - !nextBatch() - } - - @IgnoreIf({ isSharded() }) - def 'should kill cursor if limit is reached on initial query'() throws InterruptedException { - given: - def firstBatch = executeQuery(5) - - cursor = new AsyncQueryBatchCursor(firstBatch, 5, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - while (connection.getCount() > 1) { - Thread.sleep(5) - } - makeAdditionalGetMoreCall(getNamespace(), firstBatch.cursor, new SyncConnection(connection)) - - then: - thrown(MongoCursorNotFoundException) - } - - @SuppressWarnings('BracesForTryCatchFinally') - @IgnoreIf({ isSharded() }) - def 'should throw cursor not found exception'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new AsyncQueryBatchCursor(firstBatch, 0, 2, 0, new DocumentCodec(), null, connectionSource, connection) - - def connection = new SyncConnection(getConnection(connectionSource)) - def serverCursor = cursor.cursor.get() - connection.command(getNamespace().databaseName, - new BsonDocument('killCursors', new BsonString(namespace.getCollectionName())) - .append('cursors', new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))), - new NoOpFieldNameValidator(), ReadPreference.primary(), - new BsonDocumentCodec() - , connectionSource) - connection.release() - nextBatch() - - then: - try { - nextBatch() - fail('expected MongoCursorNotFoundException but no exception was thrown') - } catch (MongoCursorNotFoundException e) { - assertEquals(serverCursor.getId(), e.getCursorId()) - assertEquals(serverCursor.getAddress(), e.getServerAddress()) - } catch (ignored) { - fail('Expected MongoCursorNotFoundException to be thrown but got ' + ignored.getClass()) - } - } - - List nextBatch() { - def futureResultCallback = new FutureResultCallback() - cursor.next(futureResultCallback) - futureResultCallback.get() - } - - private QueryResult executeQuery() { - executeQuery(0) - } - - private QueryResult executeQuery(int batchSize) { - executeQuery(0, batchSize) - } - - private QueryResult executeQuery(int limit, int batchSize) { - executeQuery(new BsonDocument(), limit, batchSize, false, false) - } - - private QueryResult executeQuery(BsonDocument filter, int limit, int batchSize, boolean tailable, boolean awaitData) { - def findCommand = new BsonDocument('find', new BsonString(getCollectionName())) - .append('filter', filter) - .append('tailable', BsonBoolean.valueOf(tailable)) - .append('awaitData', BsonBoolean.valueOf(awaitData)) - - findCommand.append('limit', new BsonInt32(Math.abs(limit))) - - if (limit >= 0) { - if (batchSize < 0 && Math.abs(batchSize) < limit) { - findCommand.append('limit', new BsonInt32(Math.abs(batchSize))) - } else { - findCommand.append('batchSize', new BsonInt32(Math.abs(batchSize))) - } - } - - def futureResultCallback = new FutureResultCallback() - connection.commandAsync(getDatabaseName(), findCommand, NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), - CommandResultDocumentCodec.create(new DocumentCodec(), 'firstBatch'), connectionSource, - futureResultCallback) - def response = futureResultCallback.get() - cursorDocumentToQueryResult(response.getDocument('cursor'), connection.getDescription().getServerAddress()) - } - - private void makeAdditionalGetMoreCall(ServerCursor serverCursor, Connection connection) { - connection.command(getNamespace().databaseName, - new BsonDocument('getMore', new BsonInt64(serverCursor.getId())) - .append('collection', new BsonString(namespace.getCollectionName())), - NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), connectionSource.getSessionContext(), - connectionSource.getServerApi()) - } -} diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/CommandBatchCursorFunctionalTest.java b/driver-core/src/test/functional/com/mongodb/internal/operation/CommandBatchCursorFunctionalTest.java new file mode 100644 index 00000000000..10d0997f5d1 --- /dev/null +++ b/driver-core/src/test/functional/com/mongodb/internal/operation/CommandBatchCursorFunctionalTest.java @@ -0,0 +1,554 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + + +import com.mongodb.MongoCursorNotFoundException; +import com.mongodb.MongoQueryException; +import com.mongodb.ReadPreference; +import com.mongodb.ServerAddress; +import com.mongodb.ServerCursor; +import com.mongodb.client.model.CreateCollectionOptions; +import com.mongodb.client.model.OperationTest; +import com.mongodb.internal.binding.ConnectionSource; +import com.mongodb.internal.connection.Connection; +import org.bson.BsonArray; +import org.bson.BsonBoolean; +import org.bson.BsonDocument; +import org.bson.BsonInt32; +import org.bson.BsonInt64; +import org.bson.BsonString; +import org.bson.BsonTimestamp; +import org.bson.Document; +import org.bson.codecs.BsonDocumentCodec; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.ArrayList; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static com.mongodb.ClusterFixture.checkReferenceCountReachesTarget; +import static com.mongodb.ClusterFixture.getBinding; +import static com.mongodb.ClusterFixture.getReferenceCountAfterTimeout; +import static com.mongodb.ClusterFixture.isSharded; +import static com.mongodb.internal.operation.QueryOperationHelper.makeAdditionalGetMoreCall; +import static java.util.Collections.singletonList; +import static java.util.stream.Stream.generate; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeFalse; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +public class CommandBatchCursorFunctionalTest extends OperationTest { + + private ConnectionSource connectionSource; + private Connection connection; + private CommandBatchCursor cursor; + + @BeforeEach + void setup() { + List documents = IntStream.rangeClosed(1, 10) + .mapToObj(i -> new BsonDocument("i", new BsonInt32(i))) + .collect(Collectors.toList()); + getCollectionHelper().insertDocuments(documents); + + connectionSource = getBinding().getWriteConnectionSource(); + connection = connectionSource.getConnection(); + } + + @AfterEach + void cleanup() { + ifNotNull(cursor, CommandBatchCursor::close); + ifNotNull(connectionSource, cs -> { + getReferenceCountAfterTimeout(cs, 1); + cs.release(); + }); + ifNotNull(connection, c -> { + getReferenceCountAfterTimeout(c, 1); + c.release(); + }); + } + + @Test + @DisplayName("server cursor should not be null") + void theServerCursorShouldNotBeNull() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.getServerCursor()); + } + + @Test + @DisplayName("test server address should not be null") + void theServerAddressShouldNotNull() { + BsonDocument commandResult = executeFindCommand(); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.getServerAddress()); + } + + @Test + @DisplayName("should get Exceptions for operations on the cursor after closing") + void shouldGetExceptionsForOperationsOnTheCursorAfterClosing() { + BsonDocument commandResult = executeFindCommand(); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + cursor.close(); + + assertDoesNotThrow(() -> cursor.close()); + assertThrows(IllegalStateException.class, () -> cursor.hasNext()); + assertThrows(IllegalStateException.class, () -> cursor.next()); + assertThrows(IllegalStateException.class, () -> cursor.getServerCursor()); + } + + @Test + @DisplayName("should throw an Exception when going off the end") + void shouldThrowAnExceptionWhenGoingOffTheEnd() { + BsonDocument commandResult = executeFindCommand(1); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 2, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + cursor.next(); + cursor.next(); + assertThrows(NoSuchElementException.class, () -> cursor.next()); + } + + @Test + @DisplayName("test cursor remove") + void testCursorRemove() { + BsonDocument commandResult = executeFindCommand(); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertThrows(UnsupportedOperationException.class, () -> cursor.remove()); + } + + @Test + @DisplayName("test normal exhaustion") + void testNormalExhaustion() { + BsonDocument commandResult = executeFindCommand(); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(10, cursorFlatten().size()); + } + + @ParameterizedTest(name = "{index} => limit={0}, batchSize={1}, expectedTotal={2}") + @MethodSource + @DisplayName("test limit exhaustion") + void testLimitExhaustion(final int limit, final int batchSize, final int expectedTotal) { + BsonDocument commandResult = executeFindCommand(limit, batchSize); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, limit, batchSize, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(expectedTotal, cursorFlatten().size()); + + checkReferenceCountReachesTarget(connectionSource, 1); + checkReferenceCountReachesTarget(connection, 1); + } + + @ParameterizedTest(name = "{index} => awaitData={0}, maxTimeMS={1}") + @MethodSource + @DisplayName("should block waiting for next batch on a tailable cursor") + void shouldBlockWaitingForNextBatchOnATailableCursor(final boolean awaitData, final int maxTimeMS) { + + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, awaitData); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, maxTimeMS, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertTrue(cursor.hasNext()); + assertEquals(1, cursor.next().get(0).get("_id")); + + new Thread(() -> { + sleep(100); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 2).append("ts", new BsonTimestamp(6, 0))); + }).start(); + + assertTrue(cursor.hasNext()); + assertEquals(2, cursor.next().get(0).get("_id")); + } + + @Test + @DisplayName("test tryNext with tailable") + void testTryNextWithTailable() { + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, true); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + + List nextBatch = cursor.tryNext(); + assertNotNull(nextBatch); + assertEquals(1, nextBatch.get(0).get("_id")); + + nextBatch = cursor.tryNext(); + assertNull(nextBatch); + + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 2).append("ts", new BsonTimestamp(6, 0))); + + nextBatch = cursor.tryNext(); + assertNotNull(nextBatch); + assertEquals(2, nextBatch.get(0).get("_id")); + } + + @Test + @DisplayName("hasNext should throw when cursor is closed in another thread") + void hasNextShouldThrowWhenCursorIsClosedInAnotherThread() throws InterruptedException { + + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, true); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertTrue(cursor.hasNext()); + assertEquals(1, cursor.next().get(0).get("_id")); + + CountDownLatch latch = new CountDownLatch(1); + new Thread(() -> { + sleep(100); + cursor.close(); + latch.countDown(); + }).start(); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + assertThrows(IllegalStateException.class, () -> cursor.hasNext()); + } + + @Test + @DisplayName("test maxTimeMS") + void testMaxTimeMS() { + assumeFalse(isSharded()); + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + long maxTimeMS = 500; + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, true); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, maxTimeMS, DOCUMENT_DECODER, + null, connectionSource, connection); + + + List nextBatch = cursor.tryNext(); + assertNotNull(nextBatch); + + long startTime = System.currentTimeMillis(); + nextBatch = cursor.tryNext(); + long endTime = System.currentTimeMillis(); + + assertNull(nextBatch); + + // RACY TEST: no guarantee assertion will fire within the given timeframe + assertTrue(endTime - startTime < (maxTimeMS + 200)); + } + + @Test + @DisplayName("test tailable interrupt") + void testTailableInterrupt() throws InterruptedException { + getCollectionHelper().create(getCollectionName(), new CreateCollectionOptions().capped(true).sizeInBytes(1000)); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 1).append("ts", new BsonTimestamp(5, 0))); + + BsonDocument commandResult = executeFindCommand(new BsonDocument("ts", + new BsonDocument("$gte", new BsonTimestamp(5, 0))), 0, 2, true, true); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + CountDownLatch latch = new CountDownLatch(1); + AtomicInteger seen = new AtomicInteger(); + Thread thread = new Thread(() -> { + try { + cursor.next(); + seen.incrementAndGet(); + cursor.next(); + seen.incrementAndGet(); + } catch (Exception e) { + // pass + } finally { + latch.countDown(); + } + }); + + thread.start(); + sleep(1000); + thread.interrupt(); + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", 2)); + latch.await(); + + assertTrue(latch.await(5, TimeUnit.SECONDS)); + assertEquals(1, seen.intValue()); + } + + @Test + @DisplayName("should kill cursor if limit is reached on initial query") + void shouldKillCursorIfLimitIsReachedOnInitialQuery() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(5); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 5, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.next()); + assertFalse(cursor.hasNext()); + assertNull(cursor.getServerCursor()); + } + + @Test + @DisplayName("should kill cursor if limit is reached on getMore") + void shouldKillCursorIfLimitIsReachedOnGetMore() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 5, 3, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + ServerCursor serverCursor = cursor.getServerCursor(); + assertNotNull(serverCursor); + assertNotNull(cursor.next()); + assertNotNull(cursor.next()); + + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + assertThrows(MongoQueryException.class, () -> + makeAdditionalGetMoreCall(getNamespace(), serverCursor, connection) + ); + } + + @Test + @DisplayName("should release connection source if limit is reached on initial query") + void shouldReleaseConnectionSourceIfLimitIsReachedOnInitialQuery() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(5); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 5, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNull(cursor.getServerCursor()); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connectionSource, 1)); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + } + + @Test + @DisplayName("should release connection source if limit is reached on getMore") + void shouldReleaseConnectionSourceIfLimitIsReachedOnGetMore() { + assumeFalse(isSharded()); + BsonDocument commandResult = executeFindCommand(3); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 5, 3, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.next()); + assertNotNull(cursor.next()); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connectionSource, 1)); + assertDoesNotThrow(() -> checkReferenceCountReachesTarget(connection, 1)); + } + + @Test + @DisplayName("test limit with get more") + void testLimitWithGetMore() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 5, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertNotNull(cursor.next()); + assertNotNull(cursor.next()); + assertNotNull(cursor.next()); + assertFalse(cursor.hasNext()); + } + + @Test + @DisplayName("test limit with large documents") + void testLimitWithLargeDocuments() { + String bigString = generate(() -> "x") + .limit(16000) + .collect(Collectors.joining()); + + IntStream.range(11, 1000).forEach(i -> + getCollectionHelper().insertDocuments(DOCUMENT_DECODER, new Document("_id", i).append("s", bigString)) + ); + + BsonDocument commandResult = executeFindCommand(300, 0); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 300, 0, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(300, cursorFlatten().size()); + } + + @Test + @DisplayName("should respect batch size") + void shouldRespectBatchSize() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(2, cursor.getBatchSize()); + assertEquals(2, cursor.next().size()); + assertEquals(2, cursor.next().size()); + + cursor.setBatchSize(3); + assertEquals(3, cursor.getBatchSize()); + assertEquals(3, cursor.next().size()); + assertEquals(3, cursor.next().size()); + } + + @Test + @DisplayName("should throw cursor not found exception") + void shouldThrowCursorNotFoundException() { + BsonDocument commandResult = executeFindCommand(2); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + ServerCursor serverCursor = cursor.getServerCursor(); + assertNotNull(serverCursor); + Connection localConnection = connectionSource.getConnection(); + localConnection.command(getNamespace().getDatabaseName(), + new BsonDocument("killCursors", new BsonString(getNamespace().getCollectionName())) + .append("cursors", new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))), + NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), connectionSource); + localConnection.release(); + + cursor.next(); + + MongoCursorNotFoundException exception = assertThrows(MongoCursorNotFoundException.class, () -> cursor.next()); + assertEquals(serverCursor.getId(), exception.getCursorId()); + assertEquals(serverCursor.getAddress(), exception.getServerAddress()); + } + + @Test + @DisplayName("should report available documents") + void shouldReportAvailableDocuments() { + BsonDocument commandResult = executeFindCommand(3); + cursor = new CommandBatchCursor<>(getServerAddress(), commandResult, 0, 2, 0, DOCUMENT_DECODER, + null, connectionSource, connection); + + assertEquals(3, cursor.available()); + + cursor.next(); + assertEquals(0, cursor.available()); + + assertTrue(cursor.hasNext()); + assertEquals(2, cursor.available()); + + cursor.next(); + assertEquals(0, cursor.available()); + + assertTrue(cursor.hasNext()); + assertEquals(2, cursor.available()); + + cursor.close(); + assertEquals(0, cursor.available()); + } + + + private static Stream shouldBlockWaitingForNextBatchOnATailableCursor() { + return Stream.of( + arguments(true, 0), + arguments(true, 100), + arguments(false, 0)); + } + + private static Stream testLimitExhaustion() { + return Stream.of( + arguments(5, 2, 5), + arguments(5, -2, 2), + arguments(-5, -2, 5), + arguments(-5, 2, 5), + arguments(2, 5, 2), + arguments(2, -5, 2), + arguments(-2, 5, 2), + arguments(-2, -5, 2) + ); + } + + private ServerAddress getServerAddress() { + return connection.getDescription().getServerAddress(); + } + + private BsonDocument executeFindCommand() { + return executeFindCommand(0); + } + + private BsonDocument executeFindCommand(final int batchSize) { + return executeFindCommand(new BsonDocument(), 0, batchSize, false, false); + } + + private BsonDocument executeFindCommand(final int limit, final int batchSize) { + return executeFindCommand(new BsonDocument(), limit, batchSize, false, false); + } + + private BsonDocument executeFindCommand(final BsonDocument filter, final int limit, final int batchSize, final boolean tailable, + final boolean awaitData) { + return executeFindCommand(filter, limit, batchSize, tailable, awaitData, ReadPreference.primary()); + } + + private BsonDocument executeFindCommand(final BsonDocument filter, final int limit, final int batchSize, + final boolean tailable, final boolean awaitData, final ReadPreference readPreference) { + BsonDocument findCommand = new BsonDocument("find", new BsonString(getCollectionName())) + .append("filter", filter) + .append("tailable", BsonBoolean.valueOf(tailable)) + .append("awaitData", BsonBoolean.valueOf(awaitData)); + + findCommand.append("limit", new BsonInt32(Math.abs(limit))); + if (limit >= 0) { + if (batchSize < 0 && Math.abs(batchSize) < limit) { + findCommand.append("limit", new BsonInt32(Math.abs(batchSize))); + } else { + findCommand.append("batchSize", new BsonInt32(Math.abs(batchSize))); + } + } + + BsonDocument results = connection.command(getDatabaseName(), findCommand, + NO_OP_FIELD_NAME_VALIDATOR, readPreference, + CommandResultDocumentCodec.create(DOCUMENT_DECODER, "firstBatch"), + connectionSource); + + assertNotNull(results); + return results; + } + + private List cursorFlatten() { + List results = new ArrayList<>(); + while (cursor.hasNext()) { + results.addAll(cursor.next()); + } + return results; + } + +} diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/ListCollectionsOperationSpecification.groovy b/driver-core/src/test/functional/com/mongodb/internal/operation/ListCollectionsOperationSpecification.groovy index 9e2d8937818..18b1724fdff 100644 --- a/driver-core/src/test/functional/com/mongodb/internal/operation/ListCollectionsOperationSpecification.groovy +++ b/driver-core/src/test/functional/com/mongodb/internal/operation/ListCollectionsOperationSpecification.groovy @@ -16,13 +16,10 @@ package com.mongodb.internal.operation - import com.mongodb.MongoExecutionTimeoutException import com.mongodb.MongoNamespace import com.mongodb.OperationFunctionalSpecification import com.mongodb.ReadPreference -import com.mongodb.ServerAddress -import com.mongodb.ServerCursor import com.mongodb.async.FutureResultCallback import com.mongodb.client.model.CreateCollectionOptions import com.mongodb.connection.ConnectionDescription @@ -34,7 +31,6 @@ import com.mongodb.internal.binding.ConnectionSource import com.mongodb.internal.binding.ReadBinding import com.mongodb.internal.connection.AsyncConnection import com.mongodb.internal.connection.Connection -import com.mongodb.internal.connection.QueryResult import org.bson.BsonBoolean import org.bson.BsonDocument import org.bson.BsonDouble @@ -84,7 +80,7 @@ class ListCollectionsOperationSpecification extends OperationFunctionalSpecifica cursor.next(callback) then: - callback.get() == null + callback.get() == [] cleanup: collectionHelper.dropDatabase(madeUpDatabase) @@ -380,7 +376,7 @@ class ListCollectionsOperationSpecification extends OperationFunctionalSpecifica cursor.getBatchSize() == 2 cleanup: - consumeAsyncResults(cursor) + cursor?.close() } @IgnoreIf({ isSharded() }) @@ -479,11 +475,6 @@ class ListCollectionsOperationSpecification extends OperationFunctionalSpecifica threeSixConnectionDescription : Stub(ConnectionDescription) { getMaxWireVersion() >> 3 }, - queryResult: Stub(QueryResult) { - getNamespace() >> new MongoNamespace('db', 'coll') - getResults() >> [] - getCursor() >> new ServerCursor(1, Stub(ServerAddress)) - }, commandResult: new BsonDocument('ok', new BsonDouble(1.0)) .append('cursor', new BsonDocument('id', new BsonInt64(1)).append('ns', new BsonString('db.coll')) .append('firstBatch', new BsonArrayWrapper([]))) diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/ListIndexesOperationSpecification.groovy b/driver-core/src/test/functional/com/mongodb/internal/operation/ListIndexesOperationSpecification.groovy index 4ca91524e9f..20f842c3fff 100644 --- a/driver-core/src/test/functional/com/mongodb/internal/operation/ListIndexesOperationSpecification.groovy +++ b/driver-core/src/test/functional/com/mongodb/internal/operation/ListIndexesOperationSpecification.groovy @@ -16,13 +16,10 @@ package com.mongodb.internal.operation - import com.mongodb.MongoExecutionTimeoutException import com.mongodb.MongoNamespace import com.mongodb.OperationFunctionalSpecification import com.mongodb.ReadPreference -import com.mongodb.ServerAddress -import com.mongodb.ServerCursor import com.mongodb.async.FutureResultCallback import com.mongodb.connection.ConnectionDescription import com.mongodb.internal.async.AsyncBatchCursor @@ -34,7 +31,6 @@ import com.mongodb.internal.binding.ReadBinding import com.mongodb.internal.bulk.IndexRequest import com.mongodb.internal.connection.AsyncConnection import com.mongodb.internal.connection.Connection -import com.mongodb.internal.connection.QueryResult import org.bson.BsonDocument import org.bson.BsonDouble import org.bson.BsonInt32 @@ -76,7 +72,7 @@ class ListIndexesOperationSpecification extends OperationFunctionalSpecification cursor.next(callback) then: - callback.get() == null + callback.get() == [] } @@ -210,7 +206,7 @@ class ListIndexesOperationSpecification extends OperationFunctionalSpecification cursor.getBatchSize() == 2 cleanup: - consumeAsyncResults(cursor) + cursor?.close() } @IgnoreIf({ isSharded() }) @@ -310,11 +306,6 @@ class ListIndexesOperationSpecification extends OperationFunctionalSpecification threeSixConnectionDescription : Stub(ConnectionDescription) { getMaxWireVersion() >> 3 }, - queryResult: Stub(QueryResult) { - getNamespace() >> new MongoNamespace('db', 'coll') - getResults() >> [] - getCursor() >> new ServerCursor(1, Stub(ServerAddress)) - }, commandResult: new BsonDocument('ok', new BsonDouble(1.0)) .append('cursor', new BsonDocument('id', new BsonInt64(1)).append('ns', new BsonString('db.coll')) .append('firstBatch', new BsonArrayWrapper([]))) diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/QueryBatchCursorFunctionalSpecification.groovy b/driver-core/src/test/functional/com/mongodb/internal/operation/QueryBatchCursorFunctionalSpecification.groovy deleted file mode 100644 index 9c77bb41b89..00000000000 --- a/driver-core/src/test/functional/com/mongodb/internal/operation/QueryBatchCursorFunctionalSpecification.groovy +++ /dev/null @@ -1,642 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation - -import com.mongodb.MongoCursorNotFoundException -import com.mongodb.MongoTimeoutException -import com.mongodb.OperationFunctionalSpecification -import com.mongodb.ReadPreference -import com.mongodb.ServerCursor -import com.mongodb.WriteConcern -import com.mongodb.client.model.CreateCollectionOptions -import com.mongodb.internal.binding.ConnectionSource -import com.mongodb.internal.connection.Connection -import com.mongodb.internal.connection.QueryResult -import com.mongodb.internal.validator.NoOpFieldNameValidator -import org.bson.BsonArray -import org.bson.BsonBoolean -import org.bson.BsonDocument -import org.bson.BsonInt32 -import org.bson.BsonInt64 -import org.bson.BsonString -import org.bson.BsonTimestamp -import org.bson.Document -import org.bson.codecs.BsonDocumentCodec -import org.bson.codecs.DocumentCodec -import spock.lang.IgnoreIf -import util.spock.annotations.Slow - -import java.util.concurrent.CountDownLatch -import java.util.concurrent.TimeUnit - -import static com.mongodb.ClusterFixture.checkReferenceCountReachesTarget -import static com.mongodb.ClusterFixture.getBinding -import static com.mongodb.ClusterFixture.isDiscoverableReplicaSet -import static com.mongodb.ClusterFixture.isSharded -import static com.mongodb.ClusterFixture.serverVersionLessThan -import static com.mongodb.internal.operation.OperationHelper.cursorDocumentToQueryResult -import static com.mongodb.internal.operation.QueryOperationHelper.makeAdditionalGetMoreCall -import static java.util.Collections.singletonList -import static org.junit.Assert.assertEquals -import static org.junit.Assert.fail - -class QueryBatchCursorFunctionalSpecification extends OperationFunctionalSpecification { - ConnectionSource connectionSource - QueryBatchCursor cursor - - def setup() { - def documents = [] - for (int i = 0; i < 10; i++) { - documents.add(new BsonDocument('_id', new BsonInt32(i))) - } - collectionHelper.insertDocuments(documents, - isDiscoverableReplicaSet() ? WriteConcern.MAJORITY : WriteConcern.ACKNOWLEDGED, - getBinding()) - connectionSource = getBinding().getReadConnectionSource() - } - - def cleanup() { - cursor?.close() - connectionSource?.release() - } - - def 'server cursor should not be null'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 0, new DocumentCodec(), null, connectionSource) - - then: - cursor.getServerCursor() != null - } - - def 'test server address'() { - given: - def firstBatch = executeQuery() - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 0, new DocumentCodec(), null, connectionSource) - then: - cursor.getServerAddress() != null - } - - def 'should get Exceptions for operations on the cursor after closing'() { - given: - def firstBatch = executeQuery() - - cursor = new QueryBatchCursor(firstBatch, 0, 0, new DocumentCodec(), null, connectionSource) - - when: - cursor.close() - cursor.close() - - and: - cursor.next() - - then: - thrown(IllegalStateException) - - when: - cursor.hasNext() - - then: - thrown(IllegalStateException) - - when: - cursor.getServerCursor() - - then: - thrown(IllegalStateException) - } - - def 'should throw an Exception when going off the end'() { - given: - def firstBatch = executeQuery(1) - - cursor = new QueryBatchCursor(firstBatch, 2, 0, new DocumentCodec(), null, connectionSource) - when: - cursor.next() - cursor.next() - cursor.next() - - then: - thrown(NoSuchElementException) - } - - def 'test normal exhaustion'() { - given: - def firstBatch = executeQuery() - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 0, new DocumentCodec(), null, connectionSource) - - then: - cursor.iterator().sum { it.size() } == 10 - } - - def 'test limit exhaustion'() { - given: - def firstBatch = executeQuery(limit, batchSize) - def connection = connectionSource.getConnection() - - when: - cursor = new QueryBatchCursor(firstBatch, limit, batchSize, 0, new DocumentCodec(), null, connectionSource, connection) - - then: - cursor.iterator().sum { it.size() } == expectedTotal - - cleanup: - connection?.release() - - where: - limit | batchSize | expectedTotal - 5 | 2 | 5 - 5 | -2 | 2 - -5 | 2 | 5 - -5 | -2 | 5 - 2 | 5 | 2 - 2 | -5 | 2 - -2 | 5 | 2 - -2 | -5 | 2 - } - - def 'test remove'() { - given: - def firstBatch = executeQuery() - - cursor = new QueryBatchCursor(firstBatch, 0, 0, new DocumentCodec(), null, connectionSource) - - when: - cursor.remove() - - then: - thrown(UnsupportedOperationException) - } - - @SuppressWarnings('EmptyCatchBlock') - @Slow - def 'should block waiting for next batch on a tailable cursor'() { - given: - def connection = connectionSource.getConnection() - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(5, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, awaitData) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, maxTimeMS, new DocumentCodec(), null, connectionSource, connection) - - then: - cursor.hasNext() - cursor.next().iterator().next().get('_id') == 1 - - when: - def latch = new CountDownLatch(1) - Thread.start { - try { - sleep(500) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 2).append('ts', new BsonTimestamp(6, 0))) - } catch (ignored) { - //pass - } finally { - latch.countDown() - } - } - - // Note: this test is racy. - // The sleep above does not guarantee that we're testing what we're trying to, which is the loop in the hasNext() method. - then: - cursor.hasNext() - cursor.next().iterator().next().get('_id') == 2 - - cleanup: - def cleanedUp = latch.await(10, TimeUnit.SECONDS) - if (!cleanedUp) { - throw new MongoTimeoutException('Timed out waiting for documents to be inserted') - } - connection?.release() - - where: - awaitData | maxTimeMS - true | 0 - true | 100 - false | 0 - } - - @Slow - def 'test try next with tailable'() { - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(5, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, true) - - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - - then: - cursor.tryNext().iterator().next().get('_id') == 1 - - then: - !cursor.tryNext() - - when: - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 2).append('ts', new BsonTimestamp(6, 0))) - def nextBatch = cursor.tryNext() - - then: - nextBatch - nextBatch.iterator().next().get('_id') == 2 - } - - @Slow - def 'hasNext should throw when cursor is closed in another thread'() { - Connection conn = connectionSource.getConnection() - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(5, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, true) - cursor = new QueryBatchCursor(firstBatch, 0, 2, 0, new DocumentCodec(), null, connectionSource, conn) - cursor.next() - def closeCompleted = new CountDownLatch(1) - - // wait a second then close the cursor - new Thread({ - sleep(1000) - cursor.close() - closeCompleted.countDown() - } as Runnable).start() - - when: - cursor.hasNext() - - then: - thrown(Exception) - closeCompleted.await(5, TimeUnit.SECONDS) - conn.getCount() == 1 - - cleanup: - conn.release() - } - - @IgnoreIf({ serverVersionLessThan(3, 2) || isSharded() }) - @Slow - def 'test maxTimeMS'() { - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1).append('ts', new BsonTimestamp(5, 0))) - def firstBatch = executeQuery(new BsonDocument('ts', new BsonDocument('$gte', new BsonTimestamp(5, 0))), 0, 2, true, true) - - def connection = connectionSource.getConnection() - def maxTimeMS = 10 - cursor = new QueryBatchCursor(firstBatch, 0, 2, maxTimeMS, new DocumentCodec(), null, connectionSource, connection) - cursor.tryNext() - long startTime = System.currentTimeMillis() - - when: - def result = cursor.tryNext() - - then: - result == null - // RACY TEST: no guarantee assertion will fire within the given timeframe - System.currentTimeMillis() - startTime < (maxTimeMS + 200) - - cleanup: - connection?.release() - } - - @SuppressWarnings('EmptyCatchBlock') - @Slow - def 'test tailable interrupt'() throws InterruptedException { - collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000)) - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 1)) - - def firstBatch = executeQuery(new BsonDocument(), 0, 2, true, true) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - - CountDownLatch latch = new CountDownLatch(1) - def seen = 0 - def thread = Thread.start { - try { - cursor.next() - seen = 1 - cursor.next() - seen = 2 - } catch (ignored) { - // pass - } finally { - latch.countDown() - } - } - sleep(1000) - thread.interrupt() - collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', 2)) - latch.await() - - then: - seen == 1 - } - - @IgnoreIf({ isSharded() }) - def 'should kill cursor if limit is reached on initial query'() throws InterruptedException { - given: - def firstBatch = executeQuery(5) - def connection = connectionSource.getConnection() - - cursor = new QueryBatchCursor(firstBatch, 5, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - when: - makeAdditionalGetMoreCall(getNamespace(), firstBatch.cursor, connection) - - then: - thrown(MongoCursorNotFoundException) - - cleanup: - connection?.release() - } - - @IgnoreIf({ isSharded() }) - @Slow - def 'should kill cursor if limit is reached on get more'() throws InterruptedException { - given: - def firstBatch = executeQuery(3) - - cursor = new QueryBatchCursor(firstBatch, 5, 3, new DocumentCodec(), null, connectionSource) - ServerCursor serverCursor = cursor.getServerCursor() - - cursor.next() - cursor.next() - - Thread.sleep(1000) //Note: waiting for some time for killCursor operation to be performed on a server. - when: - makeAdditionalGetMoreCall(getNamespace(), serverCursor, connectionSource) - - then: - thrown(MongoCursorNotFoundException) - } - - def 'should release connection source if limit is reached on initial query'() throws InterruptedException { - given: - def firstBatch = executeQuery(5) - def connection = connectionSource.getConnection() - - when: - cursor = new QueryBatchCursor(firstBatch, 5, 0, 0, new DocumentCodec(), null, connectionSource, connection) - - then: - checkReferenceCountReachesTarget(connectionSource, 1) - - cleanup: - connection?.release() - } - - def 'should release connection source if limit is reached on get more'() throws InterruptedException { - given: - def firstBatch = executeQuery(3) - - cursor = new QueryBatchCursor(firstBatch, 5, 3, new DocumentCodec(), null, connectionSource) - - when: - cursor.next() - cursor.next() - - then: - checkReferenceCountReachesTarget(connectionSource, 1) - } - - def 'test limit with get more'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 5, 2, new DocumentCodec(), null, connectionSource) - - then: - cursor.next() != null - cursor.next() != null - cursor.next() != null - !cursor.hasNext() - } - - @Slow - def 'test limit with large documents'() { - given: - char[] array = 'x' * 16000 - String bigString = new String(array) - - (11..1000).each { collectionHelper.insertDocuments(new DocumentCodec(), new Document('_id', it).append('s', bigString)) } - def firstBatch = executeQuery(300, 0) - - when: - cursor = new QueryBatchCursor(firstBatch, 300, 0, new DocumentCodec(), null, connectionSource) - - then: - cursor.iterator().sum { it.size() } == 300 - } - - def 'should respect batch size'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - - then: - cursor.batchSize == 2 - - when: - def nextBatch = cursor.next() - - then: - nextBatch.size() == 2 - - when: - nextBatch = cursor.next() - - then: - nextBatch.size() == 2 - - when: - cursor.batchSize = 3 - nextBatch = cursor.next() - - then: - cursor.batchSize == 3 - nextBatch.size() == 3 - - when: - nextBatch = cursor.next() - - then: - nextBatch.size() == 3 - } - - def 'test normal loop with get more'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - def results = cursor.iterator().collectMany { it*.get('_id') } - - then: - results == (0..9).toList() - !cursor.hasNext() - } - - def 'test next without has next with get more'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - - then: - (0..4).each { cursor.next() } - !cursor.hasNext() - !cursor.hasNext() - - when: - cursor.next() - - then: - thrown(NoSuchElementException) - } - - @SuppressWarnings('BracesForTryCatchFinally') - @IgnoreIf({ isSharded() }) - def 'should throw cursor not found exception'() { - given: - def firstBatch = executeQuery(2) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - def serverCursor = cursor.getServerCursor() - def connection = connectionSource.getConnection() - connection.command(getNamespace().databaseName, - new BsonDocument('killCursors', new BsonString(namespace.getCollectionName())) - .append('cursors', new BsonArray(singletonList(new BsonInt64(serverCursor.getId())))), - new NoOpFieldNameValidator(), ReadPreference.primary(), new BsonDocumentCodec(), connectionSource) - connection.release() - cursor.next() - - then: - try { - cursor.next() - } catch (MongoCursorNotFoundException e) { - assertEquals(serverCursor.getId(), e.getCursorId()) - assertEquals(serverCursor.getAddress(), e.getServerAddress()) - } catch (ignored) { - fail('Expected MongoCursorNotFoundException to be thrown but got ' + ignored.getClass()) - } - } - - def 'should report available documents'() { - given: - def firstBatch = executeQuery(3) - - when: - cursor = new QueryBatchCursor(firstBatch, 0, 2, new DocumentCodec(), null, connectionSource) - - then: - cursor.available() == 3 - - when: - cursor.hasNext() - - then: - cursor.available() == 3 - - when: - cursor.next() - - then: - cursor.available() == 0 - - when: - cursor.hasNext() - - then: - cursor.available() == 2 - - when: - cursor.next() - - then: - cursor.available() == 0 - - when: - cursor.hasNext() - - then: - cursor.available() == 2 - - when: - cursor.close() - - then: - cursor.available() == 0 - } - - private QueryResult executeQuery() { - executeQuery(0) - } - - private QueryResult executeQuery(int batchSize) { - executeQuery(new BsonDocument(), 0, batchSize, false, false, ReadPreference.primary()) - } - - private QueryResult executeQuery(int batchSize, ReadPreference readPreference) { - executeQuery(new BsonDocument(), 0, batchSize, false, false, readPreference) - } - - private QueryResult executeQuery(int limit, int batchSize) { - executeQuery(new BsonDocument(), limit, batchSize, false, false, ReadPreference.primary()) - } - - - private QueryResult executeQuery(BsonDocument filter, int limit, int batchSize, boolean tailable, boolean awaitData) { - executeQuery(filter, limit, batchSize, tailable, awaitData, ReadPreference.primary()) - } - - private QueryResult executeQuery(BsonDocument filter, int limit, int batchSize, boolean tailable, boolean awaitData, - ReadPreference readPreference) { - def connection = connectionSource.getConnection() - try { - def findCommand = new BsonDocument('find', new BsonString(getCollectionName())) - .append('filter', filter) - .append('tailable', BsonBoolean.valueOf(tailable)) - .append('awaitData', BsonBoolean.valueOf(awaitData)) - - findCommand.append('limit', new BsonInt32(Math.abs(limit))) - - if (limit >= 0) { - if (batchSize < 0 && Math.abs(batchSize) < limit) { - findCommand.append('limit', new BsonInt32(Math.abs(batchSize))) - } else { - findCommand.append('batchSize', new BsonInt32(Math.abs(batchSize))) - } - } - - def response = connection.command(getDatabaseName(), findCommand, - NO_OP_FIELD_NAME_VALIDATOR, readPreference, - CommandResultDocumentCodec.create(new DocumentCodec(), 'firstBatch'), connectionSource) - cursorDocumentToQueryResult(response.getDocument('cursor'), connection.getDescription().getServerAddress()) - } finally { - connection.release() - } - } -} diff --git a/driver-core/src/test/functional/com/mongodb/internal/operation/QueryOperationHelper.groovy b/driver-core/src/test/functional/com/mongodb/internal/operation/QueryOperationHelper.groovy index c11d113f80c..8c609306f75 100644 --- a/driver-core/src/test/functional/com/mongodb/internal/operation/QueryOperationHelper.groovy +++ b/driver-core/src/test/functional/com/mongodb/internal/operation/QueryOperationHelper.groovy @@ -22,9 +22,10 @@ import com.mongodb.MongoNamespace import com.mongodb.MongoQueryException import com.mongodb.ReadPreference import com.mongodb.ServerCursor +import com.mongodb.async.FutureResultCallback import com.mongodb.internal.IgnorableRequestContext -import com.mongodb.internal.binding.ConnectionSource import com.mongodb.internal.binding.StaticBindingContext +import com.mongodb.internal.connection.AsyncConnection import com.mongodb.internal.connection.Connection import com.mongodb.internal.connection.NoOpSessionContext import com.mongodb.internal.connection.OperationContext @@ -53,18 +54,8 @@ class QueryOperationHelper { } } - static void makeAdditionalGetMoreCall(MongoNamespace namespace, ServerCursor serverCursor, - ConnectionSource connectionSource) { - def connection = connectionSource.getConnection() - try { - makeAdditionalGetMoreCall(namespace, serverCursor, connection) - } finally { - connection.release() - } - } - static void makeAdditionalGetMoreCall(MongoNamespace namespace, ServerCursor serverCursor, Connection connection) { - try { + makeAdditionalGetMoreCallHandleError(serverCursor) { connection.command(namespace.databaseName, new BsonDocument('getMore', new BsonInt64(serverCursor.getId())) .append('collection', new BsonString(namespace.getCollectionName())), @@ -72,6 +63,26 @@ class QueryOperationHelper { new BsonDocumentCodec(), new StaticBindingContext(new NoOpSessionContext(), getServerApi(), IgnorableRequestContext.INSTANCE, new OperationContext())) + } + } + + static void makeAdditionalGetMoreCall(MongoNamespace namespace, ServerCursor serverCursor, AsyncConnection connection) { + def callback = new FutureResultCallback<>() + connection.commandAsync(namespace.databaseName, + new BsonDocument('getMore', new BsonInt64(serverCursor.getId())) + .append('collection', new BsonString(namespace.getCollectionName())), + new NoOpFieldNameValidator(), ReadPreference.primary(), + new BsonDocumentCodec(), + new StaticBindingContext(new NoOpSessionContext(), getServerApi(), IgnorableRequestContext.INSTANCE, + new OperationContext()), + callback + ) + makeAdditionalGetMoreCallHandleError(serverCursor) { callback.get() } + } + + static void makeAdditionalGetMoreCallHandleError(ServerCursor serverCursor, Runnable runnable) { + try { + runnable.run() } catch (MongoCommandException e) { if (e.getErrorCode() == 43) { throw new MongoCursorNotFoundException(serverCursor.getId(), e.getResponse(), serverCursor.getAddress()) diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncChangeStreamBatchCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncChangeStreamBatchCursorSpecification.groovy index 1142ce5f91c..f58f58f3cb7 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncChangeStreamBatchCursorSpecification.groovy +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncChangeStreamBatchCursorSpecification.groovy @@ -27,11 +27,12 @@ import static java.util.concurrent.TimeUnit.SECONDS class AsyncChangeStreamBatchCursorSpecification extends Specification { - def 'should call the underlying AsyncQueryBatchCursor'() { + def 'should call the underlying AsyncCommandBatchCursor'() { given: def changeStreamOpertation = Stub(ChangeStreamOperation) def binding = Mock(AsyncReadBinding) - def wrapped = Mock(AsyncQueryBatchCursor) + 1 * binding.retain() >> binding + def wrapped = Mock(AsyncCommandBatchCursor) def callback = Stub(SingleResultCallback) def cursor = new AsyncChangeStreamBatchCursor(changeStreamOpertation, wrapped, binding, null, ServerVersionHelper.FOUR_DOT_FOUR_WIRE_VERSION) @@ -66,7 +67,8 @@ class AsyncChangeStreamBatchCursorSpecification extends Specification { def 'should not close the cursor in next if the cursor was closed before next completed'() { def changeStreamOpertation = Stub(ChangeStreamOperation) def binding = Mock(AsyncReadBinding) - def wrapped = Mock(AsyncQueryBatchCursor) + 1 * binding.retain() >> binding + def wrapped = Mock(AsyncCommandBatchCursor) def callback = Stub(SingleResultCallback) def cursor = new AsyncChangeStreamBatchCursor(changeStreamOpertation, wrapped, binding, null, ServerVersionHelper.FOUR_DOT_FOUR_WIRE_VERSION) @@ -91,7 +93,8 @@ class AsyncChangeStreamBatchCursorSpecification extends Specification { def 'should throw a MongoException when next/tryNext is called after the cursor is closed'() { def changeStreamOpertation = Stub(ChangeStreamOperation) def binding = Mock(AsyncReadBinding) - def wrapped = Mock(AsyncQueryBatchCursor) + 1 * binding.retain() >> binding + def wrapped = Mock(AsyncCommandBatchCursor) def cursor = new AsyncChangeStreamBatchCursor(changeStreamOpertation, wrapped, binding, null, ServerVersionHelper.FOUR_DOT_FOUR_WIRE_VERSION) diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncQueryBatchCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncCommandBatchCursorSpecification.groovy similarity index 68% rename from driver-core/src/test/unit/com/mongodb/internal/operation/AsyncQueryBatchCursorSpecification.groovy rename to driver-core/src/test/unit/com/mongodb/internal/operation/AsyncCommandBatchCursorSpecification.groovy index 5efcbc736ab..a1878a20bb5 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncQueryBatchCursorSpecification.groovy +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncCommandBatchCursorSpecification.groovy @@ -30,7 +30,6 @@ import com.mongodb.connection.ServerVersion import com.mongodb.internal.async.SingleResultCallback import com.mongodb.internal.binding.AsyncConnectionSource import com.mongodb.internal.connection.AsyncConnection -import com.mongodb.internal.connection.QueryResult import org.bson.BsonArray import org.bson.BsonDocument import org.bson.BsonInt32 @@ -42,17 +41,18 @@ import spock.lang.Specification import static OperationUnitSpecification.getMaxWireVersionForServerVersion import static com.mongodb.ReadPreference.primary +import static com.mongodb.internal.operation.CommandBatchCursorHelper.MESSAGE_IF_CLOSED_AS_CURSOR -class AsyncQueryBatchCursorSpecification extends Specification { +class AsyncCommandBatchCursorSpecification extends Specification { def 'should generate expected command with batchSize and maxTimeMS'() { given: def connection = referenceCountedAsyncConnection() def connectionSource = getAsyncConnectionSource(connection) - def firstBatch = new QueryResult(NAMESPACE, [], 42, SERVER_ADDRESS) - def cursor = new AsyncQueryBatchCursor(firstBatch, 0, batchSize, maxTimeMS, CODEC, null, connectionSource, - connection) + def firstBatch = createCommandResult([]) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, batchSize, maxTimeMS, CODEC, + null, connectionSource, connection) def expectedCommand = new BsonDocument('getMore': new BsonInt64(CURSOR_ID)) .append('collection', new BsonString(NAMESPACE.getCollectionName())) if (batchSize != 0) { @@ -62,7 +62,7 @@ class AsyncQueryBatchCursorSpecification extends Specification { expectedCommand.append('maxTimeMS', new BsonInt64(expectedMaxTimeFieldValue)) } - def reply = documentResponse([], 0) + def reply = getMoreResponse([], 0) when: def batch = nextBatch(cursor) @@ -71,7 +71,7 @@ class AsyncQueryBatchCursorSpecification extends Specification { 1 * connection.commandAsync(NAMESPACE.getDatabaseName(), expectedCommand, *_) >> { it.last().onResult(reply, null) } - batch == null + batch.isEmpty() then: !cursor.isClosed() @@ -95,14 +95,15 @@ class AsyncQueryBatchCursorSpecification extends Specification { def serverVersion = new ServerVersion([3, 6, 0]) def connection = referenceCountedAsyncConnection(serverVersion) def connectionSource = getAsyncConnectionSource(connection) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) when: - def cursor = new AsyncQueryBatchCursor(firstBatch, 0, 0, 0, CODEC, null, connectionSource, connection) cursor.close() then: - if (firstBatch.getCursor() != null) { - 1 * connection.commandAsync(NAMESPACE.databaseName, createKillCursorsDocument(firstBatch.cursor), _, primary(), *_) >> { + if (cursor.getServerCursor() != null) { + 1 * connection.commandAsync(NAMESPACE.databaseName, createKillCursorsDocument(cursor.getServerCursor()), _, primary(), *_) >> { it.last().onResult(null, null) } } @@ -112,15 +113,18 @@ class AsyncQueryBatchCursorSpecification extends Specification { connectionSource.getCount() == 0 where: - firstBatch << [queryResult(), queryResult(FIRST_BATCH, 0)] + firstBatch << [createCommandResult(), createCommandResult(FIRST_BATCH, 0)] } def 'should return the expected results from next'() { given: - def connectionSource = getAsyncConnectionSource(referenceCountedAsyncConnection()) + def connection = referenceCountedAsyncConnection() + def connectionSource = getAsyncConnectionSource(connection) when: - def cursor = new AsyncQueryBatchCursor(queryResult(FIRST_BATCH, 0), 0, 0, 0, CODEC, null, connectionSource, null) + def firstBatch = createCommandResult(FIRST_BATCH, 0) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) then: nextBatch(cursor) == FIRST_BATCH @@ -136,23 +140,23 @@ class AsyncQueryBatchCursorSpecification extends Specification { then: def exception = thrown(MongoException) - exception.getMessage() == 'next() called after the cursor was closed.' + exception.getMessage() == MESSAGE_IF_CLOSED_AS_CURSOR } def 'should respect the limit'() { given: def serverVersion = new ServerVersion([3, 6, 0]) - def connectionA = referenceCountedAsyncConnection(serverVersion) - def connectionB = referenceCountedAsyncConnection(serverVersion) + def connectionA = referenceCountedAsyncConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedAsyncConnection(serverVersion, 'connectionB') def connectionSource = getAsyncConnectionSource(connectionA, connectionB) def firstBatch = [new Document('_id', 1), new Document('_id', 2), new Document('_id', 3)] - def secondBatch = [new Document('_id', 4), new Document('_id', 5)] - def thirdBatch = [new Document('_id', 6)] + def secondBatch = [new Document('_id', 4), new Document('_id', 5), new Document('_id', 6)] + def thirdBatch = [new Document('_id', 7)] when: - def cursor = new AsyncQueryBatchCursor(queryResult(firstBatch), 6, 2, 0, CODEC, null, connectionSource, - connectionA) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, createCommandResult(firstBatch), 7, 3, 0, CODEC, + null, connectionSource, connectionA) def batch = nextBatch(cursor) then: @@ -162,7 +166,7 @@ class AsyncQueryBatchCursorSpecification extends Specification { batch = nextBatch(cursor) then: - 1 * connectionA.commandAsync(*_) >> { it.last().onResult(documentResponse(secondBatch), null) } + 1 * connectionA.commandAsync(*_) >> { it.last().onResult(getMoreResponse(secondBatch), null) } then: batch == secondBatch @@ -176,20 +180,14 @@ class AsyncQueryBatchCursorSpecification extends Specification { 1 * connectionB.commandAsync(*_) >> { connectionB.getCount() == 1 connectionSource.getCount() == 1 - it.last().onResult(documentResponse(thirdBatch, 0), null) + it.last().onResult(getMoreResponse(thirdBatch, 0), null) } then: batch == thirdBatch connectionB.getCount() == 0 connectionSource.getCount() == 0 - - when: - batch = nextBatch(cursor) - - then: - batch == null - connectionSource.getCount() == 0 + cursor.isClosed() } @@ -198,16 +196,15 @@ class AsyncQueryBatchCursorSpecification extends Specification { def serverVersion = new ServerVersion([3, 6, 0]) def connection = referenceCountedAsyncConnection(serverVersion) def connectionSource = getAsyncConnectionSource(connection) - def queryResult = queryResult() + def firstBatch = createCommandResult() when: - def cursor = new AsyncQueryBatchCursor(queryResult, 1, 0, 0, CODEC, null, connectionSource, connection) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 1, 0, 0, CODEC, + null, connectionSource, connection) then: - 1 * connection.commandAsync(NAMESPACE.databaseName, createKillCursorsDocument(queryResult.cursor), _, primary(), - *_) >> { - it.last().onResult(null, null) - } + 1 * connection.commandAsync(NAMESPACE.databaseName, + createKillCursorsDocument(new ServerCursor(42, SERVER_ADDRESS)), _, primary(), *_) >> { it.last().onResult(null, null) } when: cursor.close() @@ -222,53 +219,59 @@ class AsyncQueryBatchCursorSpecification extends Specification { def 'should handle getMore when there are empty results but there is a cursor'() { given: - def connection = referenceCountedAsyncConnection(serverVersion) - def connectionSource = getAsyncConnectionSource(connection) + def connectionA = referenceCountedAsyncConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedAsyncConnection(serverVersion, 'connectionB') + def connectionSource = getAsyncConnectionSource(connectionA, connectionB) when: - def cursor = new AsyncQueryBatchCursor(queryResult([], 42), 3, 0, 0, CODEC, null, connectionSource, connection) + def firstBatch = createCommandResult([], CURSOR_ID) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 3, 0, 0, CODEC, + null, connectionSource, connectionA) 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) + it.last().onResult(response, null) } - 1 * connection.commandAsync(*_) >> { - connection.getCount() == 1 + 1 * connectionB.commandAsync(*_) >> { + connectionB.getCount() == 1 connectionSource.getCount() == 1 - it.last().onResult(response2, null) + it.last().onResult(response2, null) } then: batch == SECOND_BATCH then: - connection.getCount() == 0 + connectionA.getCount() == 0 + connectionB.getCount() == 0 connectionSource.getCount() == 0 when: cursor.close() then: - 0 * connection._ + 0 * connectionA._ + 0 * connectionB._ connectionSource.getCount() == 0 where: - serverVersion | response | response2 - new ServerVersion([3, 6, 0]) | documentResponse([]) | documentResponse(SECOND_BATCH, 0) + serverVersion | response | response2 + new ServerVersion([3, 6, 0]) | getMoreResponse([]) | getMoreResponse(SECOND_BATCH, 0) } def 'should kill the cursor in the getMore if limit is reached'() { given: def connection = referenceCountedAsyncConnection(serverVersion) def connectionSource = getAsyncConnectionSource(connection) - def initialResult = queryResult() + def firstBatch = createCommandResult() when: - def cursor = new AsyncQueryBatchCursor(initialResult, 3, 0, 0, CODEC, null, connectionSource, connection) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 3, 0, 0, CODEC, + null, connectionSource, connection) def batch = nextBatch(cursor) then: @@ -281,7 +284,7 @@ class AsyncQueryBatchCursorSpecification extends Specification { 1 * connection.commandAsync(*_) >> { it.last().onResult(response, null) } - 1 * connection.commandAsync(NAMESPACE.databaseName, createKillCursorsDocument(initialResult.cursor), _, primary(), _, + 1 * connection.commandAsync(NAMESPACE.databaseName, createKillCursorsDocument(cursor.getServerCursor()), _, primary(), _, connectionSource, *_) >> { it.last().onResult(null, null) } @@ -299,19 +302,20 @@ class AsyncQueryBatchCursorSpecification extends Specification { where: serverVersion | response - new ServerVersion([3, 2, 0]) | documentResponse(SECOND_BATCH) + new ServerVersion([3, 2, 0]) | getMoreResponse(SECOND_BATCH) } def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore returns a response'() { given: def serverVersion = new ServerVersion([3, 6, 0]) - def connectionA = referenceCountedAsyncConnection(serverVersion) - def connectionB = referenceCountedAsyncConnection(serverVersion) + def connectionA = referenceCountedAsyncConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedAsyncConnection(serverVersion, 'connectionB') def connectionSource = getAsyncConnectionSource(serverType, connectionA, connectionB) - def initialResult = queryResult() + def firstBatch = createCommandResult() when: - def cursor = new AsyncQueryBatchCursor(initialResult, 0, 0, 0, CODEC, null, connectionSource, connectionA) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) def batch = nextBatch(cursor) then: @@ -324,10 +328,10 @@ class AsyncQueryBatchCursorSpecification extends Specification { numberOfInvocations * connectionA.commandAsync(*_) >> { // Simulate the user calling close while the getMore is in flight cursor.close() - ((SingleResultCallback) it.last()).onResult(response, null) + ((SingleResultCallback) it.last()).onResult(response, null) } >> { - // `killCursors` command - ((SingleResultCallback) it.last()).onResult(response2, null) + // `killCursors` command + ((SingleResultCallback) it.last()).onResult(null, null) } then: @@ -335,26 +339,29 @@ class AsyncQueryBatchCursorSpecification extends Specification { then: connectionA.getCount() == 0 + connectionB.getCount() == 0 + connectionSource.getCount() == 0 cursor.isClosed() where: - response | response2 | getMoreResponseHasCursor | serverType | numberOfInvocations - documentResponse([]) | documentResponse([], 0) | true | ServerType.LOAD_BALANCER | 2 - documentResponse([], 0) | null | false | ServerType.LOAD_BALANCER | 1 - documentResponse([]) | documentResponse([], 0) | true | ServerType.STANDALONE | 1 - documentResponse([], 0) | null | false | ServerType.STANDALONE | 1 + response | serverType | numberOfInvocations + getMoreResponse([]) | ServerType.LOAD_BALANCER | 2 + getMoreResponse([], 0) | ServerType.LOAD_BALANCER | 1 + getMoreResponse([]) | ServerType.STANDALONE | 2 + getMoreResponse([], 0) | ServerType.STANDALONE | 1 } def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore throws exception'() { given: - def serverVersion = new ServerVersion([3, 2, 0]) - def connectionA = referenceCountedAsyncConnection(serverVersion) - def connectionB = referenceCountedAsyncConnection(serverVersion) + def serverVersion = new ServerVersion([4, 4, 0]) + def connectionA = referenceCountedAsyncConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedAsyncConnection(serverVersion, 'connectionB') def connectionSource = getAsyncConnectionSource(serverType, connectionA, connectionB) - def initialResult = queryResult() + def firstBatch = createCommandResult() when: - def cursor = new AsyncQueryBatchCursor(initialResult, 0, 0, 0, CODEC, null, connectionSource, connectionA) + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) def batch = nextBatch(cursor) then: @@ -367,10 +374,10 @@ class AsyncQueryBatchCursorSpecification extends Specification { numberOfInvocations * connectionA.commandAsync(*_) >> { // Simulate the user calling close while the getMore is throwing a MongoException cursor.close() - ((SingleResultCallback) it.last()).onResult(null, MONGO_EXCEPTION) - } >> { - // `killCursors` command - ((SingleResultCallback) it.last()).onResult(null, null) + ((SingleResultCallback) it.last()).onResult(null, MONGO_EXCEPTION) + } >> { + // `killCursors` command + ((SingleResultCallback) it.last()).onResult(null, null) } then: @@ -390,15 +397,12 @@ class AsyncQueryBatchCursorSpecification extends Specification { given: def connection = referenceCountedAsyncConnection() def connectionSource = getAsyncConnectionSourceWithResult(ServerType.STANDALONE) { [null, MONGO_EXCEPTION] } - def cursor = new AsyncQueryBatchCursor(queryResult(), 0, 0, 0, CODEC, null, connectionSource, connection) + def firstBatch = createCommandResult() + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) when: cursor.close() - nextBatch(cursor) - - then: - def exception = thrown(MongoException) - exception.getMessage() == 'next() called after the cursor was closed.' then: cursor.isClosed() @@ -412,7 +416,9 @@ class AsyncQueryBatchCursorSpecification extends Specification { def connectionSource = getAsyncConnectionSourceWithResult(ServerType.STANDALONE) { [null, MONGO_EXCEPTION] } when: - def cursor = new AsyncQueryBatchCursor(queryResult(), 0, 0, 0, CODEC, null, connectionSource, connection) + def firstBatch = createCommandResult() + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) then: nextBatch(cursor) @@ -424,39 +430,35 @@ class AsyncQueryBatchCursorSpecification extends Specification { thrown(MongoException) then: + connection.getCount() == 0 connectionSource.getCount() == 1 - - when: - cursor.close() - - then: - connectionSource.getCount() == 0 } def 'should handle errors when calling getMore'() { given: def serverVersion = new ServerVersion([3, 6, 0]) - def connectionA = referenceCountedAsyncConnection(serverVersion) - def connectionB = referenceCountedAsyncConnection(serverVersion) + def connectionA = referenceCountedAsyncConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedAsyncConnection(serverVersion, 'connectionB') def connectionSource = getAsyncConnectionSource(connectionA, connectionB) when: - def cursor = new AsyncQueryBatchCursor(queryResult([]), 0, 0, 0, CODEC, null, connectionSource, - connectionA) + def firstBatch = createCommandResult() + def cursor = new AsyncCommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) then: connectionSource.getCount() == 1 when: nextBatch(cursor) + nextBatch(cursor) then: 1 * connectionA.commandAsync(*_) >> { connectionA.getCount() == 1 connectionSource.getCount() == 1 - it.last().onResult(null, exception) + it.last().onResult(null, exception) } - then: thrown(MongoException) @@ -468,13 +470,22 @@ class AsyncQueryBatchCursorSpecification extends Specification { cursor.close() then: - connectionSource.getCount() == 1 + 1 * connectionB.commandAsync(*_) >> { + connectionB.getCount() == 1 + connectionSource.getCount() == 1 + it.last().onResult(null, null) + } + + then: + connectionA.getCount() == 0 + connectionB.getCount() == 0 + connectionSource.getCount() == 0 where: exception << [COMMAND_EXCEPTION, MONGO_EXCEPTION] } - List nextBatch(AsyncQueryBatchCursor cursor) { + List nextBatch(AsyncCommandBatchCursor cursor) { def futureResultCallback = new FutureResultCallback() cursor.next(futureResultCallback) futureResultCallback.get() @@ -490,25 +501,32 @@ class AsyncQueryBatchCursorSpecification extends Specification { private static final COMMAND_EXCEPTION = new MongoCommandException(BsonDocument.parse('{"ok": false, "errmsg": "error"}'), SERVER_ADDRESS) - private static BsonDocument documentResponse(results, cursorId = 42) { - new BsonDocument('ok', new BsonInt32(1)).append('cursor', - new BsonDocument('id', new BsonInt64(cursorId)).append('ns', - new BsonString(NAMESPACE.getFullName())) - .append('nextBatch', new BsonArrayWrapper(results))) + private static BsonDocument getMoreResponse(results, cursorId = CURSOR_ID) { + createCommandResult(results, cursorId, "nextBatch") + } + + private static BsonDocument createCommandResult(List results = FIRST_BATCH, Long cursorId = CURSOR_ID, + String fieldNameContainingBatch = "firstBatch") { + new BsonDocument("ok", new BsonInt32(1)) + .append("cursor", + new BsonDocument("ns", new BsonString(NAMESPACE.fullName)) + .append("id", new BsonInt64(cursorId)) + .append(fieldNameContainingBatch, new BsonArrayWrapper(results))) } - private static QueryResult queryResult(results = FIRST_BATCH, cursorId = 42) { - new QueryResult(NAMESPACE, results, cursorId, SERVER_ADDRESS) + private static BsonDocument createKillCursorsDocument(ServerCursor serverCursor) { + new BsonDocument('killCursors', new BsonString(NAMESPACE.getCollectionName())) + .append('cursors', new BsonArray(Collections.singletonList(new BsonInt64(serverCursor.id)))) } - def referenceCountedAsyncConnection() { + AsyncConnection referenceCountedAsyncConnection() { referenceCountedAsyncConnection(new ServerVersion([3, 2, 0])) } - def referenceCountedAsyncConnection(ServerVersion serverVersion) { + AsyncConnection referenceCountedAsyncConnection(ServerVersion serverVersion, String name = 'connection') { def released = false def counter = 0 - def mock = Mock(AsyncConnection) { + def mock = Mock(AsyncConnection, name: name) { _ * getDescription() >> Stub(ConnectionDescription) { getMaxWireVersion() >> getMaxWireVersionForServerVersion(serverVersion.getVersionList()) } @@ -581,10 +599,4 @@ class AsyncQueryBatchCursorSpecification extends Specification { mock.getCount() >> { counter } mock } - - BsonDocument createKillCursorsDocument(ServerCursor serverCursor) { - new BsonDocument('killCursors', new BsonString(NAMESPACE.getCollectionName())) - .append('cursors', new BsonArray(Collections.singletonList(new BsonInt64(serverCursor.id)))) - } - } diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchCursorTest.java b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchCursorTest.java new file mode 100644 index 00000000000..561a4cf9f31 --- /dev/null +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchCursorTest.java @@ -0,0 +1,84 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.MongoException; +import com.mongodb.async.FutureResultCallback; +import org.bson.Document; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.concurrent.TimeUnit; + +import static com.mongodb.ClusterFixture.TIMEOUT; +import static com.mongodb.internal.operation.AsyncSingleBatchCursor.createEmptyAsyncSingleBatchCursor; +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertIterableEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + + +class AsyncSingleBatchCursorTest { + + private static final List SINGLE_BATCH = asList(new Document("a", 1), new Document("b", 2)); + + @Test + @DisplayName("should work as expected") + void shouldWorkAsExpected() { + try (AsyncSingleBatchCursor cursor = new AsyncSingleBatchCursor<>(SINGLE_BATCH, 0)) { + + assertIterableEquals(SINGLE_BATCH, nextBatch(cursor)); + assertIterableEquals(emptyList(), nextBatch(cursor)); + assertTrue(cursor.isClosed()); + + assertThrows(MongoException.class, () -> nextBatch(cursor)); + } + } + + @Test + @DisplayName("should work as expected emptyCursor") + void shouldWorkAsExpectedEmptyCursor() { + try (AsyncSingleBatchCursor cursor = createEmptyAsyncSingleBatchCursor(0)) { + assertIterableEquals(emptyList(), nextBatch(cursor)); + assertTrue(cursor.isClosed()); + + assertThrows(MongoException.class, () -> nextBatch(cursor)); + } + } + + @Test + @DisplayName("should not support setting batch size") + void shouldNotSupportSettingBatchSize() { + try (AsyncSingleBatchCursor cursor = new AsyncSingleBatchCursor<>(SINGLE_BATCH, 0)) { + + assertEquals(0, cursor.getBatchSize()); + + cursor.setBatchSize(1); + assertEquals(0, cursor.getBatchSize()); + } + } + + List nextBatch(final AsyncSingleBatchCursor cursor) { + FutureResultCallback> futureResultCallback = new FutureResultCallback<>(); + cursor.next(futureResultCallback); + return futureResultCallback.get(TIMEOUT, TimeUnit.MILLISECONDS); + } + +} diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchQueryCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchQueryCursorSpecification.groovy deleted file mode 100644 index 22f9035404f..00000000000 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/AsyncSingleBatchQueryCursorSpecification.groovy +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation - -import com.mongodb.MongoException -import com.mongodb.MongoNamespace -import com.mongodb.ServerAddress -import com.mongodb.async.FutureResultCallback -import com.mongodb.internal.connection.QueryResult -import org.bson.Document -import spock.lang.Specification - -class AsyncSingleBatchQueryCursorSpecification extends Specification { - - def 'should work as expected'() { - given: - def cursor = new AsyncSingleBatchQueryCursor(firstBatch) - - when: - def batch = nextBatch(cursor) - - then: - batch == firstBatch.getResults() - - then: - nextBatch(cursor) == null - - when: - nextBatch(cursor) - - then: - thrown(MongoException) - } - - def 'should not support setting batchsize'() { - given: - def cursor = new AsyncSingleBatchQueryCursor(firstBatch) - - when: - cursor.setBatchSize(1) - - then: - cursor.getBatchSize() == 0 - } - - - List nextBatch(AsyncSingleBatchQueryCursor cursor) { - def futureResultCallback = new FutureResultCallback() - cursor.next(futureResultCallback) - futureResultCallback.get() - } - - def firstBatch = new QueryResult(new MongoNamespace('db', 'coll'), [new Document('a', 1)], 0, new ServerAddress()) -} diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorSpecification.groovy index e654c2ef5ca..3d5219c15f5 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorSpecification.groovy +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorSpecification.groovy @@ -23,11 +23,11 @@ import spock.lang.Specification class ChangeStreamBatchCursorSpecification extends Specification { - def 'should call the underlying QueryBatchCursor'() { + def 'should call the underlying CommandBatchCursor'() { given: def changeStreamOperation = Stub(ChangeStreamOperation) def binding = Stub(ReadBinding) - def wrapped = Mock(QueryBatchCursor) + def wrapped = Mock(CommandBatchCursor) def resumeToken = new BsonDocument('_id': new BsonInt32(1)) def cursor = new ChangeStreamBatchCursor(changeStreamOperation, wrapped, binding, resumeToken, ServerVersionHelper.FOUR_DOT_FOUR_WIRE_VERSION) diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/CommandBatchCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/CommandBatchCursorSpecification.groovy new file mode 100644 index 00000000000..a5538500b75 --- /dev/null +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/CommandBatchCursorSpecification.groovy @@ -0,0 +1,645 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation + +import com.mongodb.MongoCommandException +import com.mongodb.MongoException +import com.mongodb.MongoNamespace +import com.mongodb.MongoSocketException +import com.mongodb.MongoSocketOpenException +import com.mongodb.ServerAddress +import com.mongodb.ServerCursor +import com.mongodb.connection.ConnectionDescription +import com.mongodb.connection.ServerConnectionState +import com.mongodb.connection.ServerDescription +import com.mongodb.connection.ServerType +import com.mongodb.connection.ServerVersion +import com.mongodb.internal.binding.ConnectionSource +import com.mongodb.internal.connection.Connection +import org.bson.BsonArray +import org.bson.BsonDocument +import org.bson.BsonInt32 +import org.bson.BsonInt64 +import org.bson.BsonString +import org.bson.Document +import org.bson.codecs.DocumentCodec +import spock.lang.Specification + +import static com.mongodb.ReadPreference.primary +import static com.mongodb.internal.operation.OperationUnitSpecification.getMaxWireVersionForServerVersion + +class CommandBatchCursorSpecification extends Specification { + + def 'should generate expected command with batchSize and maxTimeMS'() { + given: + def connection = referenceCountedConnection() + def connectionSource = getConnectionSource(connection) + + def firstBatch = createCommandResult([]) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, batchSize, maxTimeMS, CODEC, + null, connectionSource, connection) + def expectedCommand = new BsonDocument('getMore': new BsonInt64(CURSOR_ID)) + .append('collection', new BsonString(NAMESPACE.getCollectionName())) + if (batchSize != 0) { + expectedCommand.append('batchSize', new BsonInt32(batchSize)) + } + if (expectedMaxTimeFieldValue != null) { + expectedCommand.append('maxTimeMS', new BsonInt64(expectedMaxTimeFieldValue)) + } + + def reply = getMoreResponse([], 0) + + when: + cursor.hasNext() + + then: + 1 * connection.command(NAMESPACE.getDatabaseName(), expectedCommand, *_) >> reply + + then: + !cursor.isClosed() + + then: + cursor.close() + + then: + connection.getCount() == 0 + connectionSource.getCount() == 0 + + where: + batchSize | maxTimeMS | expectedMaxTimeFieldValue + 0 | 0 | null + 2 | 0 | null + 0 | 100 | 100 + } + + def 'should close the cursor'() { + given: + def serverVersion = new ServerVersion([3, 6, 0]) + def connection = referenceCountedConnection(serverVersion) + def connectionSource = getConnectionSource(connection) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) + + when: + cursor.close() + + then: + if (cursor.getServerCursor() != null) { + 1 * connection.command(NAMESPACE.databaseName, createKillCursorsDocument(cursor.getServerCursor()), _, primary(), *_) + } + + then: + connection.getCount() == 0 + connectionSource.getCount() == 0 + + where: + firstBatch << [createCommandResult(FIRST_BATCH, 42), createCommandResult(FIRST_BATCH, 0)] + } + + def 'should return the expected results from next'() { + given: + def connection = referenceCountedConnection() + def connectionSource = getConnectionSource(connection) + + when: + def firstBatch = createCommandResult(FIRST_BATCH, 0) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) + + then: + cursor.next() == FIRST_BATCH + + then: + connectionSource.getCount() == 0 + + then: + // Unlike the AsyncCommandBatchCursor - the cursor isn't automatically closed + !cursor.isClosed() + } + + def 'should respect the limit'() { + given: + def serverVersion = new ServerVersion([3, 6, 0]) + def connectionA = referenceCountedConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedConnection(serverVersion, 'connectionB') + def connectionSource = getConnectionSource(connectionA, connectionB) + + def firstBatch = [new Document('_id', 1), new Document('_id', 2), new Document('_id', 3)] + def secondBatch = [new Document('_id', 4), new Document('_id', 5), new Document('_id', 6)] + def thirdBatch = [new Document('_id', 7)] + + when: + def cursor = new CommandBatchCursor(SERVER_ADDRESS, createCommandResult(firstBatch), 7, 2, 0, CODEC, + null, connectionSource, connectionA) + def batch = cursor.next() + + then: + batch == firstBatch + + when: + batch = cursor.next() + + then: + 1 * connectionA.command(*_) >> getMoreResponse(secondBatch) + + then: + batch == secondBatch + connectionA.getCount() == 0 + connectionSource.getCount() == 1 + + when: + batch = cursor.next() + + then: + 1 * connectionB.command(*_) >> { + connectionB.getCount() == 1 + connectionSource.getCount() == 1 + getMoreResponse(thirdBatch, 0) + } + + then: + batch == thirdBatch + connectionB.getCount() == 0 + connectionSource.getCount() == 0 + !cursor.hasNext() + } + + + def 'should close the cursor immediately if the limit has been reached'() { + given: + def serverVersion = new ServerVersion([3, 6, 0]) + def connection = referenceCountedConnection(serverVersion) + def connectionSource = getConnectionSource(connection) + def firstBatch = createCommandResult() + + when: + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 1, 0, 0, CODEC, + null, connectionSource, connection) + + then: + 1 * connection.command(NAMESPACE.databaseName, + createKillCursorsDocument(new ServerCursor(42, SERVER_ADDRESS)), _, primary(), *_) >> null + + when: + cursor.close() + + then: + 0 * connection.command(_, _, _, _, _) + + then: + connection.getCount() == 0 + connectionSource.getCount() == 0 + } + + def 'should handle getMore when there are empty results but there is a cursor'() { + given: + def connectionA = referenceCountedConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedConnection(serverVersion, 'connectionB') + def connectionSource = getConnectionSource(connectionA, connectionB) + + when: + def firstBatch = createCommandResult([], CURSOR_ID) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 3, 0, 0, CODEC, + null, connectionSource, connectionA) + def batch = cursor.next() + + then: + 1 * connectionA.command(*_) >> { + connectionA.getCount() == 1 + connectionSource.getCount() == 1 + response + } + + 1 * connectionB.command(*_) >> { + connectionB.getCount() == 1 + connectionSource.getCount() == 1 + response2 + } + + then: + batch == SECOND_BATCH + + then: + connectionA.getCount() == 0 + connectionB.getCount() == 0 + connectionSource.getCount() == 0 + + when: + cursor.close() + + then: + 0 * connectionA._ + 0 * connectionB._ + connectionSource.getCount() == 0 + + where: + serverVersion | response | response2 + new ServerVersion([3, 6, 0]) | getMoreResponse([]) | getMoreResponse(SECOND_BATCH, 0) + } + + def 'should kill the cursor in the getMore if limit is reached'() { + given: + def connection = referenceCountedConnection(serverVersion) + def connectionSource = getConnectionSource(connection) + def firstBatch = createCommandResult() + + when: + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 3, 0, 0, CODEC, + null, connectionSource, connection) + def batch = cursor.next() + + then: + batch == FIRST_BATCH + + when: + cursor.next() + + then: + 1 * connection.command(*_) >> response + 1 * connection.command(NAMESPACE.databaseName, createKillCursorsDocument(cursor.getServerCursor()), _, primary(), _, + connectionSource, *_) >> null + + then: + connection.getCount() == 0 + connectionSource.getCount() == 0 + + when: + cursor.close() + + then: + 0 * connection.command(*_) + connectionSource.getCount() == 0 + + where: + serverVersion | response + new ServerVersion([3, 2, 0]) | getMoreResponse(SECOND_BATCH) + } + + def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore returns a response'() { + given: + def serverVersion = new ServerVersion([3, 6, 0]) + def connectionA = referenceCountedConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedConnection(serverVersion, 'connectionB') + def connectionSource = getConnectionSource(serverType, connectionA, connectionB) + def firstBatch = createCommandResult() + + when: + CommandBatchCursor cursor = new CommandBatchCursor<>(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) + List batch = cursor.next() + + then: + batch == FIRST_BATCH + + when: + cursor.next() + + then: + // simulate the user calling `close` while `getMore` is in flight + // in LB mode the same connection is used to execute both `getMore` and `killCursors` + numberOfInvocations * connectionA.command(*_) >> { + // `getMore` command + cursor.close() + response + } >> { + // `killCursors` command + null + } + + then: + IllegalStateException e = thrown() + e.getMessage() == 'Cursor has been closed' + + then: + connectionA.getCount() == 0 + connectionB.getCount() == 0 + connectionSource.getCount() == 0 + cursor.isClosed() + + where: + response | serverType | numberOfInvocations + getMoreResponse([]) | ServerType.LOAD_BALANCER | 2 + getMoreResponse([], 0) | ServerType.LOAD_BALANCER | 1 + getMoreResponse([]) | ServerType.STANDALONE | 2 + getMoreResponse([], 0) | ServerType.STANDALONE | 1 + } + + def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore throws exception'() { + given: + def serverVersion = new ServerVersion([4, 4, 0]) + def connectionA = referenceCountedConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedConnection(serverVersion, 'connectionB') + def connectionSource = getConnectionSource(serverType, connectionA, connectionB) + def firstBatch = createCommandResult() + + when: + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) + def batch = cursor.next() + + then: + batch == FIRST_BATCH + + when: + cursor.next() + + then: + numberOfInvocations * connectionA.command(*_) >> { + // Simulate the user calling close while the getMore is throwing a MongoException + cursor.close() + throw MONGO_EXCEPTION + } >> { + // `killCursors` command + null + } + + then: + thrown(MongoException) + + then: + connectionA.getCount() == 0 + cursor.isClosed() + + where: + serverType | numberOfInvocations + ServerType.LOAD_BALANCER | 2 + ServerType.STANDALONE | 1 + } + + def 'should handle errors when calling close'() { + given: + def connection = referenceCountedConnection() + def connectionSource = getConnectionSourceWithResult(ServerType.STANDALONE) { throw MONGO_EXCEPTION } + def firstBatch = createCommandResult() + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) + + when: + cursor.close() + + then: + cursor.isClosed() + connectionSource.getCount() == 0 + } + + + def 'should handle errors when getting a connection for getMore'() { + given: + def connection = referenceCountedConnection() + def connectionSource = getConnectionSourceWithResult(ServerType.STANDALONE) { throw MONGO_EXCEPTION } + + when: + def firstBatch = createCommandResult() + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connection) + + then: + cursor.next() + + when: + cursor.hasNext() + + then: + thrown(MongoException) + + then: + connection.getCount() == 0 + connectionSource.getCount() == 1 + } + + def 'should handle errors when calling getMore'() { + given: + def serverVersion = new ServerVersion([3, 6, 0]) + def connectionA = referenceCountedConnection(serverVersion, 'connectionA') + def connectionB = referenceCountedConnection(serverVersion, 'connectionB') + def connectionSource = getConnectionSource(connectionA, connectionB) + + when: + def firstBatch = createCommandResult() + def cursor = new CommandBatchCursor(SERVER_ADDRESS, firstBatch, 0, 0, 0, CODEC, + null, connectionSource, connectionA) + + then: + connectionSource.getCount() == 1 + + when: + cursor.next() + cursor.next() + + then: + 1 * connectionA.command(*_) >> { + connectionA.getCount() == 1 + connectionSource.getCount() == 1 + throw exception + } + + then: + thrown(MongoException) + + then: + connectionA.getCount() == 0 + connectionSource.getCount() == 1 + + when: + cursor.close() + + then: + 1 * connectionB.command(*_) >> { + connectionB.getCount() == 1 + connectionSource.getCount() == 1 + null + } + + then: + connectionA.getCount() == 0 + connectionB.getCount() == 0 + connectionSource.getCount() == 0 + + where: + exception << [COMMAND_EXCEPTION, MONGO_EXCEPTION] + } + + def 'should handle exceptions when closing'() { + given: + def connection = Mock(Connection) { + _ * getDescription() >> Stub(ConnectionDescription) { + getMaxWireVersion() >> 4 + } + _ * command(*_) >> { throw new MongoSocketException('No MongoD', SERVER_ADDRESS) } + } + def connectionSource = Stub(ConnectionSource) { + getServerApi() >> null + getConnection() >> { connection } + } + connectionSource.retain() >> connectionSource + + def initialResults = createCommandResult([]) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, initialResults, 0, 2, 100, new DocumentCodec(), + null, connectionSource, connection) + + when: + cursor.close() + + then: + notThrown(MongoSocketException) + + when: + cursor.close() + + then: + notThrown(Exception) + } + + def 'should handle exceptions when killing cursor and a connection can not be obtained'() { + given: + def connection = Mock(Connection) { + _ * getDescription() >> Stub(ConnectionDescription) { + getMaxWireVersion() >> 4 + } + } + def connectionSource = Stub(ConnectionSource) { + getConnection() >> { throw new MongoSocketOpenException("can't open socket", SERVER_ADDRESS, new IOException()) } + getServerApi() >> null + } + connectionSource.retain() >> connectionSource + + def initialResults = createCommandResult([]) + def cursor = new CommandBatchCursor(SERVER_ADDRESS, initialResults, 0, 2, 100, new DocumentCodec(), + null, connectionSource, connection) + + when: + cursor.close() + + then: + notThrown(MongoSocketException) + + when: + cursor.close() + + then: + notThrown(Exception) + } + + private static final MongoNamespace NAMESPACE = new MongoNamespace('db', 'coll') + private static final ServerAddress SERVER_ADDRESS = new ServerAddress() + private static final CURSOR_ID = 42 + private static final FIRST_BATCH = [new Document('_id', 1), new Document('_id', 2)] + private static final SECOND_BATCH = [new Document('_id', 3), new Document('_id', 4)] + private static final CODEC = new DocumentCodec() + private static final MONGO_EXCEPTION = new MongoException('error') + private static final COMMAND_EXCEPTION = new MongoCommandException(BsonDocument.parse('{"ok": false, "errmsg": "error"}'), + SERVER_ADDRESS) + + private static BsonDocument getMoreResponse(results, cursorId = CURSOR_ID) { + createCommandResult(results, cursorId, "nextBatch") + } + + private static BsonDocument createCommandResult(List results = FIRST_BATCH, Long cursorId = CURSOR_ID, + String fieldNameContainingBatch = "firstBatch") { + new BsonDocument("ok", new BsonInt32(1)) + .append("cursor", + new BsonDocument("ns", new BsonString(NAMESPACE.fullName)) + .append("id", new BsonInt64(cursorId)) + .append(fieldNameContainingBatch, new BsonArrayWrapper(results))) + } + + private static BsonDocument createKillCursorsDocument(ServerCursor serverCursor) { + new BsonDocument('killCursors', new BsonString(NAMESPACE.getCollectionName())) + .append('cursors', new BsonArray(Collections.singletonList(new BsonInt64(serverCursor.id)))) + } + + Connection referenceCountedConnection() { + referenceCountedConnection(new ServerVersion([3, 2, 0])) + } + + Connection referenceCountedConnection(ServerVersion serverVersion, String name = 'connection') { + def released = false + def counter = 0 + def mock = Mock(Connection, name: name) { + _ * getDescription() >> Stub(ConnectionDescription) { + getMaxWireVersion() >> getMaxWireVersionForServerVersion(serverVersion.getVersionList()) + } + } + mock.retain() >> { + if (released) { + throw new IllegalStateException('Tried to retain Connection when already released') + } else { + counter += 1 + } + mock + } + mock.release() >> { + counter -= 1 + if (counter == 0) { + released = true + } else if (counter < 0) { + throw new IllegalStateException('Tried to release Connection below 0') + } + counter + } + mock.getCount() >> { counter } + mock + } + + ConnectionSource getConnectionSource(Connection... connections) { + getConnectionSource(ServerType.STANDALONE, connections) + } + + ConnectionSource getConnectionSource(ServerType serverType, Connection... connections) { + def index = -1 + getConnectionSourceWithResult(serverType) { index += 1; connections.toList().get(index).retain() } + } + + def getConnectionSourceWithResult(ServerType serverType, Closure connectionCallbackResults) { + def released = false + int counter = 0 + def mock = Mock(ConnectionSource) + mock.getServerDescription() >> { + ServerDescription.builder() + .address(new ServerAddress()) + .type(serverType) + .state(ServerConnectionState.CONNECTED) + .build() + } + mock.getConnection() >> { + if (counter == 0) { + throw new IllegalStateException('Tried to use released ConnectionSource') + } + connectionCallbackResults() + } + mock.retain() >> { + if (released) { + throw new IllegalStateException('Tried to retain ConnectionSource when already released') + } else { + counter += 1 + } + mock + } + mock.release() >> { + counter -= 1 + if (counter == 0) { + released = true + } else if (counter < 0) { + throw new IllegalStateException('Tried to release ConnectionSource below 0') + } + counter + } + mock.getCount() >> { counter } + mock + } + +} diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorSpecification.groovy deleted file mode 100644 index db6831138e1..00000000000 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorSpecification.groovy +++ /dev/null @@ -1,357 +0,0 @@ -/* - * Copyright 2008-present MongoDB, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.mongodb.internal.operation - -import com.mongodb.MongoException -import com.mongodb.MongoNamespace -import com.mongodb.MongoSocketException -import com.mongodb.MongoSocketOpenException -import com.mongodb.ServerAddress -import com.mongodb.connection.ConnectionDescription -import com.mongodb.connection.ServerConnectionState -import com.mongodb.connection.ServerDescription -import com.mongodb.connection.ServerType -import com.mongodb.connection.ServerVersion -import com.mongodb.internal.binding.ConnectionSource -import com.mongodb.internal.connection.Connection -import com.mongodb.internal.connection.QueryResult -import org.bson.BsonDocument -import org.bson.BsonInt32 -import org.bson.BsonInt64 -import org.bson.BsonString -import org.bson.Document -import org.bson.codecs.BsonDocumentCodec -import org.bson.codecs.DocumentCodec -import spock.lang.Specification - -import static com.mongodb.internal.operation.OperationUnitSpecification.getMaxWireVersionForServerVersion - -class QueryBatchCursorSpecification extends Specification { - private static final MongoNamespace NAMESPACE = new MongoNamespace('db', 'coll') - private static final ServerAddress SERVER_ADDRESS = new ServerAddress() - - def 'should generate expected command with batchSize and maxTimeMS'() { - given: - def connection = Mock(Connection) { - _ * getDescription() >> Stub(ConnectionDescription) { - getMaxWireVersion() >> 4 - } - } - def connectionSource = Stub(ConnectionSource) { - getConnection() >> { connection } - getServerApi() >> null - } - connectionSource.retain() >> connectionSource - - def cursorId = 42 - - def firstBatch = new QueryResult(NAMESPACE, [], cursorId, SERVER_ADDRESS) - def cursor = new QueryBatchCursor(firstBatch, 0, batchSize, maxTimeMS, new BsonDocumentCodec(), null, connectionSource, - connection) - def expectedCommand = new BsonDocument('getMore': new BsonInt64(cursorId)) - .append('collection', new BsonString(NAMESPACE.getCollectionName())) - if (batchSize != 0) { - expectedCommand.append('batchSize', new BsonInt32(batchSize)) - } - if (expectedMaxTimeFieldValue != null) { - expectedCommand.append('maxTimeMS', new BsonInt64(expectedMaxTimeFieldValue)) - } - - def reply = new BsonDocument('ok', new BsonInt32(1)) - .append('cursor', - new BsonDocument('id', new BsonInt64(0)) - .append('ns', new BsonString(NAMESPACE.getFullName())) - .append('nextBatch', new BsonArrayWrapper([]))) - - when: - cursor.hasNext() - - then: - 1 * connection.command(NAMESPACE.getDatabaseName(), expectedCommand, _, _, _, connectionSource) >> { - reply - } - 1 * connection.release() - - where: - batchSize | maxTimeMS | expectedMaxTimeFieldValue - 0 | 0 | null - 2 | 0 | null - 0 | 100 | 100 - } - - def 'should handle exceptions when closing'() { - given: - def connection = Mock(Connection) { - _ * getDescription() >> Stub(ConnectionDescription) { - getMaxWireVersion() >> 4 - } - _ * command(*_) >> { throw new MongoSocketException('No MongoD', SERVER_ADDRESS) } - } - def connectionSource = Stub(ConnectionSource) { - getServerApi() >> null - getConnection() >> { connection } - } - connectionSource.retain() >> connectionSource - - def firstBatch = new QueryResult(NAMESPACE, [], 42, SERVER_ADDRESS) - def cursor = new QueryBatchCursor(firstBatch, 0, 2, 100, new DocumentCodec(), null, connectionSource, connection) - - when: - cursor.close() - - then: - notThrown(MongoSocketException) - - when: - cursor.close() - - then: - notThrown(Exception) - } - - def 'should handle exceptions when killing cursor and a connection can not be obtained'() { - given: - def connection = Mock(Connection) { - _ * getDescription() >> Stub(ConnectionDescription) { - getMaxWireVersion() >> 4 - } - } - def connectionSource = Stub(ConnectionSource) { - getConnection() >> { throw new MongoSocketOpenException("can't open socket", SERVER_ADDRESS, new IOException()) } - getServerApi() >> null - } - connectionSource.retain() >> connectionSource - - def firstBatch = new QueryResult(NAMESPACE, [], 42, SERVER_ADDRESS) - def cursor = new QueryBatchCursor(firstBatch, 0, 2, 100, new DocumentCodec(), null, connectionSource, connection) - - when: - cursor.close() - - then: - notThrown(MongoSocketException) - - when: - cursor.close() - - then: - notThrown(Exception) - } - - def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore returns a response'() { - given: - Connection conn = mockConnection(serverVersion) - ConnectionSource connSource - if (serverType == ServerType.LOAD_BALANCER) { - connSource = mockConnectionSource(SERVER_ADDRESS, serverType) - } else { - connSource = mockConnectionSource(SERVER_ADDRESS, serverType, conn, mockConnection(serverVersion)) - } - List firstBatch = [new Document()] - QueryResult initialResult = new QueryResult<>(NAMESPACE, firstBatch, 1, SERVER_ADDRESS) - Object getMoreResponse = useCommand - ? emptyGetMoreCommandResponse(NAMESPACE, getMoreResponseHasCursor ? 42 : 0) - : emptyGetMoreQueryResponse(NAMESPACE, SERVER_ADDRESS, getMoreResponseHasCursor ? 42 : 0) - - when: - QueryBatchCursor cursor = new QueryBatchCursor<>(initialResult, 0, 0, 0, new DocumentCodec(), null, connSource, conn) - List batch = cursor.next() - - then: - batch == firstBatch - - when: - cursor.next() - - then: - // simulate the user calling `close` while `getMore` is in flight - if (useCommand) { - // in LB mode the same connection is used to execute both `getMore` and `killCursors` - int numberOfInvocations = serverType == ServerType.LOAD_BALANCER - ? getMoreResponseHasCursor ? 2 : 1 - : 1 - numberOfInvocations * conn.command(*_) >> { - // `getMore` command - cursor.close() - getMoreResponse - } >> { - // `killCursors` command - null - } - } else { - 1 * conn.getMore(*_) >> { - cursor.close() - getMoreResponse - } - } - - then: - IllegalStateException e = thrown() - e.getMessage() == 'Cursor has been closed' - - then: - conn.getCount() == 1 - connSource.getCount() == 1 - - where: - serverVersion | useCommand | getMoreResponseHasCursor | serverType - new ServerVersion([5, 0, 0]) | true | true | ServerType.LOAD_BALANCER - new ServerVersion([5, 0, 0]) | true | false | ServerType.LOAD_BALANCER - new ServerVersion([3, 2, 0]) | true | true | ServerType.STANDALONE - new ServerVersion([3, 2, 0]) | true | false | ServerType.STANDALONE - } - - def 'should close cursor after getMore finishes if cursor was closed while getMore was in progress and getMore throws exception'() { - given: - Connection conn = mockConnection(serverVersion) - ConnectionSource connSource - if (serverType == ServerType.LOAD_BALANCER) { - connSource = mockConnectionSource(SERVER_ADDRESS, serverType) - } else { - connSource = mockConnectionSource(SERVER_ADDRESS, serverType, conn, mockConnection(serverVersion)) - } - List firstBatch = [new Document()] - QueryResult initialResult = new QueryResult<>(NAMESPACE, firstBatch, 1, SERVER_ADDRESS) - String exceptionMessage = 'test' - - when: - QueryBatchCursor cursor = new QueryBatchCursor<>(initialResult, 0, 0, 0, new DocumentCodec(), null, connSource, conn) - List batch = cursor.next() - - then: - batch == firstBatch - - when: - cursor.next() - - then: - // simulate the user calling `close` while `getMore` is in flight - if (useCommand) { - // in LB mode the same connection is used to execute both `getMore` and `killCursors` - int numberOfInvocations = serverType == ServerType.LOAD_BALANCER ? 2 : 1 - numberOfInvocations * conn.command(*_) >> { - // `getMore` command - cursor.close() - throw new MongoException(exceptionMessage) - } >> { - // `killCursors` command - null - } - } else { - 1 * conn.getMore(*_) >> { - cursor.close() - throw new MongoException(exceptionMessage) - } - } - - then: - MongoException e = thrown() - e.getMessage() == exceptionMessage - - then: - conn.getCount() == 1 - connSource.getCount() == 1 - - where: - serverVersion | useCommand | serverType - new ServerVersion([5, 0, 0]) | true | ServerType.LOAD_BALANCER - new ServerVersion([3, 2, 0]) | true | ServerType.STANDALONE - } - - /** - * Creates a {@link Connection} with {@link Connection#getCount()} returning 1. - */ - private Connection mockConnection(ServerVersion serverVersion) { - int refCounter = 1 - Connection mockConn = Mock(Connection) { - getDescription() >> Stub(ConnectionDescription) { - getMaxWireVersion() >> getMaxWireVersionForServerVersion(serverVersion.getVersionList()) - } - } - mockConn.retain() >> { - if (refCounter == 0) { - throw new IllegalStateException('Tried to retain Connection when already released') - } else { - refCounter += 1 - } - mockConn - } - mockConn.release() >> { - refCounter -= 1 - if (refCounter < 0) { - throw new IllegalStateException('Tried to release Connection below 0') - } - refCounter - } - mockConn.getCount() >> { refCounter } - mockConn - } - - private ConnectionSource mockConnectionSource(ServerAddress serverAddress, ServerType serverType, Connection... connections) { - int connIdx = 0 - int refCounter = 1 - ConnectionSource mockConnectionSource = Mock(ConnectionSource) - mockConnectionSource.getServerDescription() >> { - ServerDescription.builder() - .address(serverAddress) - .type(serverType) - .state(ServerConnectionState.CONNECTED) - .build() - } - mockConnectionSource.retain() >> { - if (refCounter == 0) { - throw new IllegalStateException('Tried to retain ConnectionSource when already released') - } else { - refCounter += 1 - } - mockConnectionSource - } - mockConnectionSource.release() >> { - refCounter -= 1 - if (refCounter < 0) { - throw new IllegalStateException('Tried to release ConnectionSource below 0') - } - refCounter - } - mockConnectionSource.getCount() >> { refCounter } - mockConnectionSource.getConnection() >> { - if (refCounter == 0) { - throw new IllegalStateException('Tried to use released ConnectionSource') - } - Connection conn - if (connIdx < connections.length) { - conn = connections[connIdx] - } else { - throw new IllegalStateException('Requested more than maxConnections=' + maxConnections) - } - connIdx++ - conn.retain() - } - mockConnectionSource - } - - private static BsonDocument emptyGetMoreCommandResponse(MongoNamespace namespace, long cursorId) { - new BsonDocument('ok', new BsonInt32(1)) - .append('cursor', new BsonDocument('id', new BsonInt64(cursorId)) - .append('ns', new BsonString(namespace.getFullName())) - .append('nextBatch', new BsonArrayWrapper([]))) - } - - private static QueryResult emptyGetMoreQueryResponse(MongoNamespace namespace, ServerAddress serverAddress, long cursorId) { - new QueryResult(namespace, [], cursorId, serverAddress) - } -} diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/SingleBatchCursorTest.java b/driver-core/src/test/unit/com/mongodb/internal/operation/SingleBatchCursorTest.java new file mode 100644 index 00000000000..a71f067f5d6 --- /dev/null +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/SingleBatchCursorTest.java @@ -0,0 +1,93 @@ +/* + * Copyright 2008-present MongoDB, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.mongodb.internal.operation; + +import com.mongodb.ServerAddress; +import org.bson.Document; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.NoSuchElementException; + +import static com.mongodb.assertions.Assertions.assertFalse; +import static com.mongodb.assertions.Assertions.assertNull; +import static com.mongodb.internal.connection.tlschannel.util.Util.assertTrue; +import static com.mongodb.internal.operation.SingleBatchCursor.createEmptySingleBatchCursor; +import static java.util.Arrays.asList; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertIterableEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + + +class SingleBatchCursorTest { + + private static final List SINGLE_BATCH = asList(new Document("a", 1), new Document("b", 2)); + private static final ServerAddress SERVER_ADDRESS = new ServerAddress(); + + @Test + @DisplayName("should work as expected") + void shouldWorkAsExpected() { + + try (SingleBatchCursor cursor = new SingleBatchCursor<>(SINGLE_BATCH, 0, SERVER_ADDRESS)) { + assertEquals(SERVER_ADDRESS, cursor.getServerAddress()); + assertEquals(1, cursor.available()); + assertNull(cursor.getServerCursor()); + + assertTrue(cursor.hasNext()); + assertIterableEquals(SINGLE_BATCH, cursor.next()); + assertEquals(0, cursor.available()); + + assertFalse(cursor.hasNext()); + assertThrows(NoSuchElementException.class, cursor::next); + } + } + + @Test + @DisplayName("should work as expected emptyCursor") + void shouldWorkAsExpectedEmptyCursor() { + try (SingleBatchCursor cursor = createEmptySingleBatchCursor(SERVER_ADDRESS, 0)) { + assertEquals(SERVER_ADDRESS, cursor.getServerAddress()); + assertEquals(0, cursor.available()); + assertNull(cursor.getServerCursor()); + + assertFalse(cursor.hasNext()); + assertThrows(NoSuchElementException.class, cursor::next); + } + } + + @Test + @DisplayName("should work as expected with try methods") + void shouldWorkAsExpectedWithTryMethods() { + try (SingleBatchCursor cursor = new SingleBatchCursor<>(SINGLE_BATCH, 0, SERVER_ADDRESS)) { + assertIterableEquals(SINGLE_BATCH, cursor.tryNext()); + assertNull(cursor.tryNext()); + } + } + + @Test + @DisplayName("should not support setting batch size") + void shouldNotSupportSettingBatchSize() { + try (SingleBatchCursor cursor = new SingleBatchCursor<>(SINGLE_BATCH, 0, SERVER_ADDRESS)) { + assertEquals(0, cursor.getBatchSize()); + + cursor.setBatchSize(1); + assertEquals(0, cursor.getBatchSize()); + } + } + +} diff --git a/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/internal/BatchCursorFluxTest.java b/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/internal/BatchCursorFluxTest.java index 91f44bfddb7..ac5730d0d00 100644 --- a/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/internal/BatchCursorFluxTest.java +++ b/driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/internal/BatchCursorFluxTest.java @@ -318,7 +318,6 @@ public void testBatchCursorDoesNotDropAnError() { Flux.fromIterable(asList(1, 2)) .flatMap(x -> Flux.from(collection.find())) .take(1) - ) .collectList() .block(TIMEOUT_DURATION);