1 files changed, 24 insertions, 2 deletions

diff --git a/src/mongo/db/cursor_manager.cpp b/src/mongo/db/cursor_manager.cpp
index d80a33e62ae..f26f77ab50d 100644
--- a/src/mongo/db/cursor_manager.cpp
+++ b/src/mongo/db/cursor_manager.cpp
@@ -206,6 +206,24 @@ std::size_t CursorManager::timeoutCursors(OperationContext* opCtx, Date_t now) {
     return toDisposeWithoutMutex.size();
 }
 
+std::vector<CursorId> CursorManager::getCursorIdsForNamespace(const NamespaceString& nss) {
+    std::vector<CursorId> cursorIds;
+
+    // Lock and inspect one partition at a time in order to avoid contention. It is acceptable for
+    // the output not to include info about cursors opened/closed while iterating.
+    for (size_t partitionId = 0; partitionId < kNumPartitions; ++partitionId) {
+        auto lockedPartition = _cursorMap->lockOnePartitionById(partitionId);
+        for (auto it = lockedPartition->begin(); it != lockedPartition->end(); ++it) {
+            auto* cursor = it->second;
+            if (cursor->nss() == nss) {
+                cursorIds.push_back(cursor->cursorid());
+            }
+        }
+    }
+
+    return cursorIds;
+}
+
 StatusWith<ClientCursorPin> CursorManager::pinCursor(
     OperationContext* opCtx,
     CursorId id,
@@ -281,7 +299,7 @@ void CursorManager::unpin(OperationContext* opCtx,
     // interesting in proactively cleaning up that cursor's resources. In these cases, we
     // proactively delete the cursor. In other cases we preserve the error code so that the client
     // will see the reason the cursor was killed when asking for the next batch.
-    if (interruptStatus == ErrorCodes::Interrupted || interruptStatus == ErrorCodes::CursorKilled) {
+    if (interruptStatus == ErrorCodes::Interrupted || cursor->isKillPending()) {
         LOGV2(20530,
               "removing cursor {cursor_cursorid} after completing batch: {error}",
               "Removing cursor after completing batch",
@@ -289,7 +307,7 @@ void CursorManager::unpin(OperationContext* opCtx,
               "error"_attr = interruptStatus);
         return deregisterAndDestroyCursor(std::move(partition), opCtx, std::move(cursor));
     } else if (!interruptStatus.isOK()) {
-        cursor->markAsKilled(interruptStatus);
+        cursor->getExecutor()->markAsKilled(interruptStatus);
     }
 
     // The cursor will stay around in '_cursorMap', so release the unique pointer to avoid deleting
@@ -456,6 +474,10 @@ Status CursorManager::killCursor(OperationContext* opCtx, CursorId id) {
             cursor->_operationUsingCursor->getServiceContext()->killOperation(
                 lk, cursor->_operationUsingCursor, ErrorCodes::CursorKilled);
         }
+
+        // Mark that the cursor has been killed on the cursor object itself as well, as other errors
+        // e.g. MaxTimeMSExpired may override the CursorKilled status.
+        cursor->setKillPending(true);
         return Status::OK();
     }
     std::unique_ptr<ClientCursor, ClientCursor::Deleter> ownedCursor(cursor);