SERVER-21600 add integration test for killCursors command

(cherry picked from commit 9f8ec07dd7ba300ea54899ff9c6c1d7ee638df3a)

9 files changed, 88 insertions, 1 deletions

diff --git a/src/mongo/db/commands/getmore_cmd.cpp b/src/mongo/db/commands/getmore_cmd.cpp
index 882215570f0..eae59f65b17 100644
--- a/src/mongo/db/commands/getmore_cmd.cpp
+++ b/src/mongo/db/commands/getmore_cmd.cpp
@@ -218,6 +218,10 @@ public:
                        str::stream() << "Cursor not found, cursor id: " << request.cursorid));
         }
 
+        // If the fail point is enabled, busy wait until it is disabled.
+        while (MONGO_FAIL_POINT(keepCursorPinnedDuringGetMore)) {
+        }
+
         if (request.nss.ns() != cursor->ns()) {
             return appendCommandStatus(
                 result,
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index cc4fcec81c6..ce2f45085a6 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -30,7 +30,6 @@ env.Library(
         "$BUILD_DIR/mongo/db/server_parameters",
         "command_request_response",
         "index_bounds",
-        "query_common",
     ],
 )
 
@@ -48,6 +47,7 @@ env.Library(
     ],
     LIBDEPS=[
         "internal_plans",
+        "query_common",
         "query_planner",
         "query_planner_test_lib",
         "$BUILD_DIR/mongo/db/curop",
diff --git a/src/mongo/db/query/find_common.cpp b/src/mongo/db/query/find_common.cpp
index 451fad77873..54e652be8a6 100644
--- a/src/mongo/db/query/find_common.cpp
+++ b/src/mongo/db/query/find_common.cpp
@@ -36,6 +36,8 @@
 
 namespace mongo {
 
+MONGO_FP_DECLARE(keepCursorPinnedDuringGetMore);
+
 bool FindCommon::enoughForFirstBatch(const LiteParsedQuery& pq,
                                      long long numDocs,
                                      int bytesBuffered) {
diff --git a/src/mongo/db/query/find_common.h b/src/mongo/db/query/find_common.h
index b999f2e95e1..6029d511157 100644
--- a/src/mongo/db/query/find_common.h
+++ b/src/mongo/db/query/find_common.h
@@ -26,11 +26,17 @@
  *    it in the license file.
  */
 
+#include "mongo/util/fail_point_service.h"
+
 namespace mongo {
 
 class BSONObj;
 class LiteParsedQuery;
 
+// Enabling this fail point will cause the getMore command to busy wait after pinning the cursor,
+// until the fail point is disabled.
+MONGO_FP_FORWARD_DECLARE(keepCursorPinnedDuringGetMore);
+
 /**
  * Suite of find/getMore related functions used in both the mongod and mongos query paths.
  */
diff --git a/src/mongo/s/cluster_cursor_stats.cpp b/src/mongo/s/cluster_cursor_stats.cpp
index cdf3f46ab4d..bbc460f7178 100644
--- a/src/mongo/s/cluster_cursor_stats.cpp
+++ b/src/mongo/s/cluster_cursor_stats.cpp
@@ -48,6 +48,7 @@ public:
 
         openBob.append("multiTarget", static_cast<long long>(stats.cursorsSharded));
         openBob.append("singleTarget", static_cast<long long>(stats.cursorsNotSharded));
+        openBob.append("pinned", static_cast<long long>(stats.cursorsPinned));
         openBob.append("total",
                        static_cast<long long>(stats.cursorsSharded + stats.cursorsNotSharded));
         openBob.done();
diff --git a/src/mongo/s/query/cluster_cursor_manager.cpp b/src/mongo/s/query/cluster_cursor_manager.cpp
index 529825c24df..941a770ea22 100644
--- a/src/mongo/s/query/cluster_cursor_manager.cpp
+++ b/src/mongo/s/query/cluster_cursor_manager.cpp
@@ -376,8 +376,15 @@ ClusterCursorManager::Stats ClusterCursorManager::stats() const {
             const CursorEntry& entry = cursorIdEntryPair.second;
 
             if (entry.getKillPending()) {
+                // Killed cursors do not count towards the number of pinned cursors or the number of
+                // open cursors.
                 continue;
             }
+
+            if (!entry.isCursorOwned()) {
+                ++stats.cursorsPinned;
+            }
+
             switch (entry.getCursorType()) {
                 case CursorType::NamespaceNotSharded:
                     ++stats.cursorsNotSharded;
diff --git a/src/mongo/s/query/cluster_cursor_manager.h b/src/mongo/s/query/cluster_cursor_manager.h
index 7c04414f15c..6ba90e67828 100644
--- a/src/mongo/s/query/cluster_cursor_manager.h
+++ b/src/mongo/s/query/cluster_cursor_manager.h
@@ -104,6 +104,9 @@ public:
 
         // Count of open cursors registered with CursorType::NamespaceNotSharded.
         size_t cursorsNotSharded = 0;
+
+        // Count of pinned cursors.
+        size_t cursorsPinned = 0;
     };
 
     /**
@@ -428,6 +431,10 @@ private:
             return _lastActive;
         }
 
+        bool isCursorOwned() const {
+            return static_cast<bool>(_cursor);
+        }
+
         /**
          * Releases the cursor from this entry.  If the cursor has already been released, returns
          * null.
diff --git a/src/mongo/s/query/cluster_cursor_manager_test.cpp b/src/mongo/s/query/cluster_cursor_manager_test.cpp
index 398a894d8df..8f5f4b54355 100644
--- a/src/mongo/s/query/cluster_cursor_manager_test.cpp
+++ b/src/mongo/s/query/cluster_cursor_manager_test.cpp
@@ -455,6 +455,7 @@ TEST_F(ClusterCursorManagerTest, ReapZombieCursorsSkipNonZombies) {
 TEST_F(ClusterCursorManagerTest, StatsInitAsZero) {
     ASSERT_EQ(0U, getManager()->stats().cursorsSharded);
     ASSERT_EQ(0U, getManager()->stats().cursorsNotSharded);
+    ASSERT_EQ(0U, getManager()->stats().cursorsPinned);
 }
 
 // Test that registering a sharded cursor updates the corresponding counter in stats().
@@ -475,6 +476,16 @@ TEST_F(ClusterCursorManagerTest, StatsRegisterNotShardedCursor) {
     ASSERT_EQ(1U, getManager()->stats().cursorsNotSharded);
 }
 
+// Test that checking out a cursor updates the pinned counter in stats().
+TEST_F(ClusterCursorManagerTest, StatsPinCursor) {
+    auto cursorId = getManager()->registerCursor(allocateMockCursor(),
+                                                 nss,
+                                                 ClusterCursorManager::CursorType::NamespaceSharded,
+                                                 ClusterCursorManager::CursorLifetime::Mortal);
+    auto pinnedCursor = getManager()->checkOutCursor(nss, cursorId);
+    ASSERT_EQ(1U, getManager()->stats().cursorsPinned);
+}
+
 // Test that registering multiple sharded and not-sharded cursors updates the corresponding
 // counters in stats().
 TEST_F(ClusterCursorManagerTest, StatsRegisterMultipleCursors) {
@@ -521,6 +532,18 @@ TEST_F(ClusterCursorManagerTest, StatsKillNotShardedCursor) {
     ASSERT_EQ(0U, getManager()->stats().cursorsNotSharded);
 }
 
+// Test that killing a pinned cursor decrements the corresponding counter in stats().
+TEST_F(ClusterCursorManagerTest, StatsKillPinnedCursor) {
+    auto cursorId = getManager()->registerCursor(allocateMockCursor(),
+                                                 nss,
+                                                 ClusterCursorManager::CursorType::NamespaceSharded,
+                                                 ClusterCursorManager::CursorLifetime::Mortal);
+    auto pinnedCursor = getManager()->checkOutCursor(nss, cursorId);
+    ASSERT_EQ(1U, getManager()->stats().cursorsPinned);
+    ASSERT_OK(getManager()->killCursor(nss, cursorId));
+    ASSERT_EQ(0U, getManager()->stats().cursorsPinned);
+}
+
 // Test that exhausting a sharded cursor decrements the corresponding counter in stats().
 TEST_F(ClusterCursorManagerTest, StatsExhaustShardedCursor) {
     auto cursorId = getManager()->registerCursor(allocateMockCursor(),
@@ -550,6 +573,38 @@ TEST_F(ClusterCursorManagerTest, StatsExhaustNotShardedCursor) {
     ASSERT_EQ(0U, getManager()->stats().cursorsNotSharded);
 }
 
+// Test that checking a pinned cursor in as exhausted decrements the corresponding counter in
+// stats().
+TEST_F(ClusterCursorManagerTest, StatsExhaustPinnedCursor) {
+    auto cursorId =
+        getManager()->registerCursor(allocateMockCursor(),
+                                     nss,
+                                     ClusterCursorManager::CursorType::NamespaceNotSharded,
+                                     ClusterCursorManager::CursorLifetime::Mortal);
+    auto pinnedCursor = getManager()->checkOutCursor(nss, cursorId);
+    ASSERT_OK(pinnedCursor.getStatus());
+    ASSERT_OK(pinnedCursor.getValue().next().getStatus());
+    ASSERT_EQ(1U, getManager()->stats().cursorsPinned);
+    pinnedCursor.getValue().returnCursor(ClusterCursorManager::CursorState::Exhausted);
+    ASSERT_EQ(0U, getManager()->stats().cursorsPinned);
+}
+
+// Test that checking a pinned cursor in as *not* exhausted decrements the corresponding counter in
+// stats().
+TEST_F(ClusterCursorManagerTest, StatsCheckInWithoutExhaustingPinnedCursor) {
+    auto cursorId =
+        getManager()->registerCursor(allocateMockCursor(),
+                                     nss,
+                                     ClusterCursorManager::CursorType::NamespaceNotSharded,
+                                     ClusterCursorManager::CursorLifetime::Mortal);
+    auto pinnedCursor = getManager()->checkOutCursor(nss, cursorId);
+    ASSERT_OK(pinnedCursor.getStatus());
+    ASSERT_OK(pinnedCursor.getValue().next().getStatus());
+    ASSERT_EQ(1U, getManager()->stats().cursorsPinned);
+    pinnedCursor.getValue().returnCursor(ClusterCursorManager::CursorState::NotExhausted);
+    ASSERT_EQ(0U, getManager()->stats().cursorsPinned);
+}
+
 // Test that getting the namespace for a cursor returns the correct namespace.
 TEST_F(ClusterCursorManagerTest, GetNamespaceForCursorIdBasic) {
     auto cursorId =
diff --git a/src/mongo/s/query/cluster_find.cpp b/src/mongo/s/query/cluster_find.cpp
index 7a9f7fc0f03..7cab6adcb28 100644
--- a/src/mongo/s/query/cluster_find.cpp
+++ b/src/mongo/s/query/cluster_find.cpp
@@ -56,6 +56,7 @@
 #include "mongo/s/query/store_possible_cursor.h"
 #include "mongo/s/stale_exception.h"
 #include "mongo/stdx/memory.h"
+#include "mongo/util/fail_point_service.h"
 #include "mongo/util/log.h"
 
 namespace mongo {
@@ -408,6 +409,10 @@ StatusWith<CursorResponse> ClusterFind::runGetMore(OperationContext* txn,
     }
     invariant(request.cursorid == pinnedCursor.getValue().getCursorId());
 
+    // If the fail point is enabled, busy wait until it is disabled.
+    while (MONGO_FAIL_POINT(keepCursorPinnedDuringGetMore)) {
+    }
+
     if (request.awaitDataTimeout) {
         auto status = pinnedCursor.getValue().setAwaitDataTimeout(*request.awaitDataTimeout);
         if (!status.isOK()) {