diff --git a/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java index 587237fcaf8..f07ee8b1878 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java +++ b/driver-core/src/main/com/mongodb/internal/operation/QueryBatchCursor.java @@ -26,6 +26,7 @@ import com.mongodb.annotations.ThreadSafe; import com.mongodb.connection.ConnectionDescription; import com.mongodb.connection.ServerType; +import com.mongodb.internal.VisibleForTesting; import com.mongodb.internal.binding.ConnectionSource; import com.mongodb.internal.connection.Connection; import com.mongodb.internal.connection.QueryResult; @@ -58,6 +59,7 @@ import static com.mongodb.assertions.Assertions.isTrueArgument; import static com.mongodb.assertions.Assertions.notNull; import static com.mongodb.internal.Locks.withLock; +import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE; import static com.mongodb.internal.operation.CursorHelper.getNumberToReturn; import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull; import static com.mongodb.internal.operation.SyncOperationHelper.getMoreCursorDocumentToQueryResult; @@ -359,7 +361,8 @@ private BsonDocument getPostBatchResumeTokenFromResponse(final BsonDocument resu * others are not and rely on the total order mentioned above. */ @ThreadSafe - private final class ResourceManager { + @VisibleForTesting(otherwise = PRIVATE) + final class ResourceManager { private final Lock lock; private volatile State state; @Nullable @@ -416,7 +419,8 @@ R execute(final String exceptionMessageIfClosed, final Supplier operation * If {@linkplain #operable() closed}, then returns false, otherwise completes abruptly. * @throws IllegalStateException Iff another operation is in progress. */ - private boolean tryStartOperation() throws IllegalStateException { + @VisibleForTesting(otherwise = PRIVATE) + boolean tryStartOperation() throws IllegalStateException { return withLock(lock, () -> { State localState = state; if (!localState.operable()) { @@ -435,7 +439,8 @@ private boolean tryStartOperation() throws IllegalStateException { /** * Thread-safe. */ - private void endOperation() { + @VisibleForTesting(otherwise = PRIVATE) + void endOperation() { boolean doClose = withLock(lock, () -> { State localState = state; if (localState == State.OPERATION_IN_PROGRESS) { @@ -459,7 +464,7 @@ private void endOperation() { void close() { boolean doClose = withLock(lock, () -> { State localState = state; - if (localState == State.OPERATION_IN_PROGRESS) { + if (localState.inProgress()) { state = State.CLOSE_PENDING; return false; } else if (localState != State.CLOSED) { diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorResourceManagerTest.java b/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorResourceManagerTest.java new file mode 100644 index 00000000000..28cd521f094 --- /dev/null +++ b/driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorResourceManagerTest.java @@ -0,0 +1,55 @@ +/* + * 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 com.mongodb.internal.binding.ConnectionSource; +import com.mongodb.internal.connection.QueryResult; +import org.bson.BsonDocument; +import org.bson.codecs.BsonDocumentCodec; +import org.junit.jupiter.api.Test; + +import static java.util.Collections.emptyList; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +final class QueryBatchCursorResourceManagerTest { + @Test + void doubleCloseExecutedConcurrentlyWithOperationBeingInProgressShouldNotFail() { + ConnectionSource connectionSourceMock = mock(ConnectionSource.class); + when(connectionSourceMock.retain()).thenReturn(connectionSourceMock); + when(connectionSourceMock.release()).thenReturn(1); + ServerAddress serverAddress = new ServerAddress(); + try (QueryBatchCursor cursor = new QueryBatchCursor<>( + new QueryResult<>(null, emptyList(), 0, serverAddress), + 1, 1, new BsonDocumentCodec())) { + QueryBatchCursor.ResourceManager cursorResourceManager = cursor.new ResourceManager( + connectionSourceMock, null, new ServerCursor(1, serverAddress)); + cursorResourceManager.tryStartOperation(); + try { + assertDoesNotThrow(() -> { + cursorResourceManager.close(); + cursorResourceManager.close(); + cursorResourceManager.setServerCursor(null); + }); + } finally { + cursorResourceManager.endOperation(); + } + } + } +}