SERVER-46767 Provide a mapping from OperationKey to CursorID in the CursorManager

1 files changed, 111 insertions, 0 deletions

diff --git a/src/mongo/dbtests/cursor_manager_test.cpp b/src/mongo/dbtests/cursor_manager_test.cpp
index e6227d325e6..124de6ead10 100644
--- a/src/mongo/dbtests/cursor_manager_test.cpp
+++ b/src/mongo/dbtests/cursor_manager_test.cpp
@@ -626,5 +626,116 @@ TEST_F(CursorManagerTest, CanAccessFromOperationContext) {
     ASSERT(cursorManager);
 }
 
+TEST_F(CursorManagerTestCustomOpCtx, CursorsWithoutOperationKeys) {
+    auto opCtx = _queryServiceContext->makeOperationContext();
+    auto pinned = makeCursor(opCtx.get());
+    ASSERT_EQUALS(pinned.getCursor()->getOperationKey(), boost::none);
+}
+
+TEST_F(CursorManagerTestCustomOpCtx, OneCursorWithAnOperationKey) {
+    auto opKey = UUID::gen();
+    auto opCtx = _queryServiceContext->makeOperationContext();
+    opCtx->setOperationKey(opKey);
+    auto pinned = makeCursor(opCtx.get());
+
+    auto cursors = useCursorManager()->getCursorsForOpKeys({opKey});
+    ASSERT_EQ(cursors.size(), size_t(1));
+    auto cursorId = pinned.getCursor()->cursorid();
+    ASSERT(cursors.find(cursorId) != cursors.end());
+
+    // Remove the cursor from the manager and verify that we can't retrieve it.
+    pinned.release();
+    ASSERT_OK(useCursorManager()->killCursor(opCtx.get(), cursorId, false));
+    ASSERT(useCursorManager()->getCursorsForOpKeys({opKey}).empty());
+}
+
+TEST_F(CursorManagerTestCustomOpCtx, MultipleCursorsMultipleOperationKeys) {
+    auto opKey1 = UUID::gen();
+    auto opKey2 = UUID::gen();
+
+    CursorId cursor1;
+    CursorId cursor2;
+
+    // Cursor with operationKey 1.
+    {
+        auto opCtx1 = _queryServiceContext->makeOperationContext();
+        opCtx1->setOperationKey(opKey1);
+        cursor1 = makeCursor(opCtx1.get()).getCursor()->cursorid();
+    }
+
+    // Cursor with operationKey 2.
+    {
+        auto opCtx2 = _queryServiceContext->makeOperationContext();
+        opCtx2->setOperationKey(opKey2);
+        cursor2 = makeCursor(opCtx2.get()).getCursor()->cursorid();
+    }
+
+    // Cursor with no operation key.
+    {
+        auto opCtx3 = _queryServiceContext->makeOperationContext();
+        makeCursor(opCtx3.get()).getCursor();
+    }
+
+    // Retrieve cursors for each operation key - should be one for each.
+    auto cursors1 = useCursorManager()->getCursorsForOpKeys({opKey1});
+    ASSERT_EQ(cursors1.size(), size_t(1));
+    ASSERT(cursors1.find(cursor1) != cursors1.end());
+
+    auto cursors2 = useCursorManager()->getCursorsForOpKeys({opKey2});
+    ASSERT_EQ(cursors2.size(), size_t(1));
+    ASSERT(cursors2.find(cursor2) != cursors2.end());
+
+    // Retrieve cursors for both operation keys.
+    auto cursors = useCursorManager()->getCursorsForOpKeys({opKey1, opKey2});
+    ASSERT_EQ(cursors.size(), size_t(2));
+    ASSERT(cursors.find(cursor1) != cursors.end());
+    ASSERT(cursors.find(cursor2) != cursors.end());
+}
+
+TEST_F(CursorManagerTestCustomOpCtx, TimedOutCursorShouldNotBeReturnedForOpKeyLookup) {
+    auto opKey = UUID::gen();
+    auto opCtx = _queryServiceContext->makeOperationContext();
+    opCtx->setOperationKey(opKey);
+    auto clock = useClock();
+
+    auto cursor = makeCursor(opCtx.get());
+
+    ASSERT_EQ(1UL, useCursorManager()->numCursors());
+    ASSERT_EQ(0UL, useCursorManager()->timeoutCursors(opCtx.get(), Date_t()));
+
+    // Advance the clock and verify that the cursor times out.
+    cursor.release();
+    clock->advance(getDefaultCursorTimeoutMillis() + Milliseconds(1));
+    ASSERT_EQ(1UL, useCursorManager()->timeoutCursors(opCtx.get(), clock->now()));
+    ASSERT_EQ(0UL, useCursorManager()->numCursors());
+
+    // Verify that the timed out cursor is not returned when looking up by OperationKey.
+    auto cursors = useCursorManager()->getCursorsForOpKeys({opKey});
+    ASSERT_EQ(cursors.size(), size_t(0));
+}
+
+TEST_F(CursorManagerTestCustomOpCtx, CursorsMarkedAsKilledAreReturnedForOpKeyLookup) {
+    auto opKey = UUID::gen();
+    auto opCtx = _queryServiceContext->makeOperationContext();
+    opCtx->setOperationKey(opKey);
+
+    auto cursor = makeCursor(opCtx.get());
+
+    // Mark the OperationContext as killed.
+    {
+        stdx::lock_guard<Client> lkClient(*opCtx->getClient());
+        // A cursor will stay alive, but be marked as killed, if it is interrupted with a code other
+        // than ErrorCodes::Interrupted or ErrorCodes::CursorKilled and then unpinned.
+        opCtx->getServiceContext()->killOperation(lkClient, opCtx.get(), ErrorCodes::InternalError);
+    }
+    cursor.release();
+
+    // The cursor should still be present in the manager.
+    ASSERT_EQ(1UL, useCursorManager()->numCursors());
+
+    // Verify that the killed cursor is still returned when looking up by OperationKey.
+    auto cursors = useCursorManager()->getCursorsForOpKeys({opKey});
+    ASSERT_EQ(cursors.size(), size_t(1));
+}
 }  // namespace
 }  // namespace mongo