Shared a CursorResourceManager between async and sync code

JAVA-5159

Author: rozza

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

@@ -19,6 +19,7 @@
 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;
@@ -32,10 +33,7 @@
 import com.mongodb.internal.connection.AsyncConnection;
 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 org.bson.BsonTimestamp;
 import org.bson.BsonValue;
 import org.bson.codecs.BsonDocumentCodec;
@@ -44,14 +42,10 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
 
 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.CommandBatchCursorHelper.FIRST_BATCH;
@@ -64,7 +58,6 @@
 import static com.mongodb.internal.operation.OperationHelper.LOGGER;
 import static java.lang.String.format;
 import static java.util.Collections.emptyList;
-import static java.util.Collections.singletonList;
 
 class AsyncCommandBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T> {
 
@@ -110,9 +103,9 @@ class AsyncCommandBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T>
             connectionToPin = connection;
         }
 
-        resourceManager = new ResourceManager(connectionSource, connectionToPin, commandCursor.getServerCursor());
+        resourceManager = new ResourceManager(namespace, connectionSource, connectionToPin, commandCursor.getServerCursor());
         if (releaseServerAndResources) {
-            resourceManager.releaseServerAndClientResources(connection.retain());
+            resourceManager.releaseServerAndClientResources(connection);
         }
     }
 
@@ -191,7 +184,7 @@ ServerCursor getServerCursor() {
     }
 
     private void getMore(final ServerCursor cursor, final SingleResultCallback<List<T>> callback) {
-        resourceManager.executeWithConnection(callback, connection -> getMore(connection, cursor, callback));
+        resourceManager.executeWithConnection(connection -> getMore(connection, cursor, callback), callback);
     }
 
     private void getMore(final AsyncConnection connection, final ServerCursor cursor, final SingleResultCallback<List<T>> callback) {
@@ -219,6 +212,7 @@ private void getMore(final AsyncConnection connection, final ServerCursor cursor
 
                     if (!resourceManager.operable()) {
                         resourceManager.releaseServerAndClientResources(connection);
+                        connection.release();
                         callback.onResult(emptyList(), null);
                         return;
                     }
@@ -230,9 +224,8 @@ private void getMore(final AsyncConnection connection, final ServerCursor cursor
                         resourceManager.endOperation();
                         if (limitReached()) {
                             resourceManager.releaseServerAndClientResources(connection);
-                        } else {
-                            connection.release();
                         }
+                        connection.release();
                         callback.onResult(commandCursor.getResults(), null);
                     }
                 });
@@ -256,268 +249,85 @@ private boolean limitReached() {
         return Math.abs(limit) != 0 && count.get() >= Math.abs(limit);
     }
 
-    /**
-     * This class maintains all resources that must be released in {@link AsyncCommandBatchCursor#close()}.
-     * It also implements a {@linkplain #doClose() deferred close action} such that it is totally ordered with other operations of
-     * {@link AsyncCommandBatchCursor} (methods {@link #tryStartOperation()}/{@link #endOperation()} must be used properly to enforce the order)
-     * despite the method {@link AsyncCommandBatchCursor#close()} being called concurrently with those operations.
-     * This total order induces the happens-before order.
-     * <p>
-     * The deferred close action does not violate externally observable idempotence of {@link AsyncCommandBatchCursor#close()},
-     * because {@link AsyncCommandBatchCursor#close()} is allowed to release resources "eventually".
-     * <p>
-     * 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 = State.IDLE;
-        @Nullable
-        private volatile AsyncConnectionSource connectionSource;
-        @Nullable
-        private volatile AsyncConnection pinnedConnection;
-        @Nullable
-        private volatile ServerCursor serverCursor;
-        private volatile boolean skipReleasingServerResourcesOnClose;
-
-        ResourceManager(final AsyncConnectionSource connectionSource,
-                @Nullable final AsyncConnection connectionToPin,
-                @Nullable final ServerCursor serverCursor) {
-            lock = new ReentrantLock();
-            if (serverCursor != null) {
-                this.connectionSource = notNull("connectionSource", 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.
-         * Returns {@code true} iff started an operation.
-         * If {@linkplain State#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) {
-                    return false;
-                } 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 if (localState != State.CLOSED) {
-                    fail(localState.toString());
-                }
-            } finally {
-                lock.unlock();
-            }
-            if (doClose) {
-                doClose();
-            }
-        }
+    private static final class ResourceManager extends CursorResourceManager<AsyncConnectionSource, AsyncConnection> {
 
-        /**
-         * 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();
-            }
+        ResourceManager(
+                final MongoNamespace namespace,
+                final AsyncConnectionSource connectionSource,
+                @Nullable final AsyncConnection connectionToPin,
+                @Nullable  final ServerCursor serverCursor) {
+            super(namespace, connectionSource, connectionToPin, serverCursor);
         }
 
-        /**
-         * This method is never executed concurrently with either itself or other operations
-         * demarcated by {@link #tryStartOperation()}/{@link #endOperation()}.
-         */
-        private void doClose() {
-            if (skipReleasingServerResourcesOnClose) {
-                serverCursor = null;
-                releaseClientResources();
-            } else if (serverCursor != null) {
-                executeWithConnection((connection, t) -> {
-                    // guarantee that regardless of exceptions, `serverCursor` is null and client resources are released
-                    serverCursor = null;
-                    releaseClientResources();
-                }, resourceManager::releaseServerAndClientResources);
-            }
+        @Override
+        void markAsPinned(final AsyncConnection connectionToPin, final Connection.PinningMode pinningMode) {
+             connectionToPin.markAsPinned(pinningMode);
         }
 
-        void onCorruptedConnection(@Nullable final AsyncConnection corruptedConnection) {
-            // if `pinnedConnection` is corrupted, then we cannot kill `serverCursor` via such a connection
-            AsyncConnection localPinnedConnection = pinnedConnection;
-            if (localPinnedConnection != null) {
-                assertTrue(corruptedConnection == localPinnedConnection);
-                skipReleasingServerResourcesOnClose = true;
-            }
+        @Override
+        void executeWithConnection(final Consumer<AsyncConnection> action) {
+            throw new UnsupportedOperationException();
         }
 
-        <R> void executeWithConnection(
-                final SingleResultCallback<R> callback,
-                final Consumer<AsyncConnection> function) {
+        @Override
+        <R> void executeWithConnection(final Consumer<AsyncConnection> action, final SingleResultCallback<R> callback) {
             assertTrue(state != State.IDLE);
             if (pinnedConnection != null) {
-                executeWithConnection(assertNotNull(pinnedConnection).retain(), null, callback, function);
+                executeWithConnection(assertNotNull(pinnedConnection).retain(), null, action, callback);
             } else {
-                assertNotNull(connectionSource).getConnection((conn, t) -> executeWithConnection(conn, t, callback, function));
+                assertNotNull(connectionSource).getConnection((conn, t) -> executeWithConnection(conn, t, action, callback));
             }
         }
 
         <R> void executeWithConnection(
                 @Nullable final AsyncConnection connection,
                 @Nullable final Throwable t,
-                final SingleResultCallback<R> callback,
-                final Consumer<AsyncConnection> function) {
+                final Consumer<AsyncConnection> function,
+                final SingleResultCallback<R> callback) {
             assertTrue(connection != null || t != null);
             if (t != null) {
-                try {
-                    onCorruptedConnection(connection);
-                } catch (Exception suppressed) {
-                    t.addSuppressed(suppressed);
-                }
                 callback.onResult(null, t);
             } else {
                 AsyncCallbackSupplier<R> curriedFunction = c -> function.accept(connection);
-                curriedFunction.whenComplete(connection::release).get(callback);
+                curriedFunction.whenComplete(() -> {
+                    System.out.println("COMPLETED!!!!");
+                    connection.release();
+                }).get((result, error) -> {
+                    if (error instanceof MongoSocketException) {
+                        onCorruptedConnection(connection);
+                    }
+                    callback.onResult(result, error);
+                });
             }
         }
 
-        /**
-         * Thread-safe.
-         */
-        @Nullable
-        ServerCursor serverCursor() {
-            return serverCursor;
-        }
-
-        void setServerCursor(@Nullable final ServerCursor serverCursor) {
-            assertTrue(state.inProgress());
-            assertNotNull(connectionSource);
-            this.serverCursor = serverCursor;
-            if (serverCursor == null) {
-                releaseClientResources();
+        @Override
+        void doClose() {
+            if (skipReleasingServerResourcesOnClose) {
+                serverCursor = null;
             }
-        }
 
-        private void releaseServerAndClientResources(final AsyncConnection connection) {
-            lock.lock();
-            ServerCursor localServerCursor = serverCursor;
-            serverCursor = null;
-            lock.unlock();
-            if (localServerCursor != null) {
-                killServerCursor(namespace, localServerCursor, connection);
+            if (serverCursor != null) {
+                executeWithConnection(conn -> {
+                    releaseServerResources(conn);
+                    conn.release();
+                }, (r, t) -> {
+                    // guarantee that regardless of exceptions, `serverCursor` is null and client resources are released
+                    serverCursor = null;
+                    releaseClientResources();
+                });
             } else {
-                connection.release();
                 releaseClientResources();
             }
         }
 
-        private void killServerCursor(final MongoNamespace namespace, final ServerCursor serverCursor, final AsyncConnection connection) {
+        @Override
+        void killServerCursor(final MongoNamespace namespace, final ServerCursor serverCursor, final AsyncConnection connection) {
             connection
-                    .commandAsync(namespace.getDatabaseName(), asKillCursorsCommandDocument(namespace, serverCursor),
+                    .commandAsync(namespace.getDatabaseName(), getKillCursorsCommand(namespace, serverCursor),
                             NO_OP_FIELD_NAME_VALIDATOR, ReadPreference.primary(), new BsonDocumentCodec(), assertNotNull(connectionSource),
-                            (r, t) -> {
-                                connection.release();
-                                releaseClientResources();
-                            });
-        }
-
-        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);
-            lock.lock();
-            AsyncConnectionSource localConnectionSource = connectionSource;
-            connectionSource = null;
-
-            AsyncConnection localPinnedConnection = pinnedConnection;
-            pinnedConnection = null;
-            lock.unlock();
-
-            if (localConnectionSource != null) {
-                localConnectionSource.release();
-            }
-            if (localPinnedConnection != null) {
-                localPinnedConnection.release();
-            }
-        }
-    }
-
-    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;
+                            (r, t) -> releaseClientResources());
         }
     }
 }