SERVER-32145 Avoid dropping lock before disposing of PlanExecutor

2 files changed, 28 insertions, 8 deletions

diff --git a/src/mongo/db/pipeline/document_source_cursor.cpp b/src/mongo/db/pipeline/document_source_cursor.cpp
index 90449e1b073..bba8a3dee72 100644
--- a/src/mongo/db/pipeline/document_source_cursor.cpp
+++ b/src/mongo/db/pipeline/document_source_cursor.cpp
@@ -31,7 +31,6 @@
 #include "mongo/db/pipeline/document_source_cursor.h"
 
 #include "mongo/db/catalog/collection.h"
-#include "mongo/db/db_raii.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/pipeline/document.h"
 #include "mongo/db/query/explain.h"
@@ -120,22 +119,28 @@ void DocumentSourceCursor::loadBatch() {
                 return;
             }
         }
-    }
 
-    // If we got here, there won't be any more documents, so destroy our PlanExecutor. Note we can't
-    // use dispose() since we want to keep the current batch.
-    cleanupExecutor();
+        // If we got here, there won't be any more documents, so destroy our PlanExecutor. Note we
+        // must hold a collection lock to destroy '_exec', but we can only assume that our locks are
+        // still held if '_exec' did not end in an error. If '_exec' encountered an error during a
+        // yield, the locks might be yielded.
+        if (state != PlanExecutor::DEAD && state != PlanExecutor::FAILURE) {
+            cleanupExecutor(autoColl);
+        }
+    }
 
     switch (state) {
         case PlanExecutor::ADVANCED:
         case PlanExecutor::IS_EOF:
             return;  // We've reached our limit or exhausted the cursor.
         case PlanExecutor::DEAD: {
+            cleanupExecutor();
             uasserted(ErrorCodes::QueryPlanKilled,
                       str::stream() << "collection or index disappeared when cursor yielded: "
                                     << WorkingSetCommon::toStatusString(resultObj));
         }
         case PlanExecutor::FAILURE: {
+            cleanupExecutor();
             uasserted(17285,
                       str::stream() << "cursor encountered an error: "
                                     << WorkingSetCommon::toStatusString(resultObj));
@@ -253,6 +258,14 @@ void DocumentSourceCursor::cleanupExecutor() {
     _exec.reset();
 }
 
+void DocumentSourceCursor::cleanupExecutor(const AutoGetCollectionForRead& readLock) {
+    invariant(_exec);
+    auto cursorManager =
+        readLock.getCollection() ? readLock.getCollection()->getCursorManager() : nullptr;
+    _exec->dispose(pExpCtx->opCtx, cursorManager);
+    _exec.reset();
+}
+
 DocumentSourceCursor::~DocumentSourceCursor() {
     invariant(!_exec);  // '_exec' should have been cleaned up via dispose() before destruction.
 }
diff --git a/src/mongo/db/pipeline/document_source_cursor.h b/src/mongo/db/pipeline/document_source_cursor.h
index 297b0c60925..cb1d0323c65 100644
--- a/src/mongo/db/pipeline/document_source_cursor.h
+++ b/src/mongo/db/pipeline/document_source_cursor.h
@@ -30,6 +30,7 @@
 
 #include <deque>
 
+#include "mongo/db/db_raii.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/pipeline/document_source.h"
 #include "mongo/db/pipeline/document_source_limit.h"
@@ -148,8 +149,8 @@ public:
 
 protected:
     /**
-     * Disposes of '_exec' and '_rangePreserver' if they haven't been disposed already. This
-     * involves taking a collection lock.
+     * Disposes of '_exec' if it hasn't been disposed already. This involves taking a collection
+     * lock.
      */
     void doDispose() final;
 
@@ -166,11 +167,17 @@ private:
     ~DocumentSourceCursor();
 
     /**
-     * Acquires locks to properly destroy and de-register '_exec'. '_exec' must be non-null.
+     * Acquires the appropriate locks, then destroys and de-registers '_exec'. '_exec' must be
+     * non-null.
      */
     void cleanupExecutor();
 
     /**
+     * Destroys and de-registers '_exec'. '_exec' must be non-null.
+     */
+    void cleanupExecutor(const AutoGetCollectionForRead& readLock);
+
+    /**
      * Reads a batch of data from '_exec'.
      */
     void loadBatch();