SERVER-58356 Cannot kill the dataSize operation

2 files changed, 50 insertions, 2 deletions

diff --git a/jstests/noPassthroughWithMongod/datasize_kill_op.js b/jstests/noPassthroughWithMongod/datasize_kill_op.js
new file mode 100644
index 00000000000..e034aa630ed
--- /dev/null
+++ b/jstests/noPassthroughWithMongod/datasize_kill_op.js
@@ -0,0 +1,41 @@
+// Test that the dataSize command can be interrupted. Failpoint is defined for mongod only,
+// therefore this test is running only in unsharded setup.
+
+(function() {
+"use strict";
+
+load("jstests/libs/fail_point_util.js");         // For configureFailPoint.
+load("jstests/libs/wait_for_command.js");        // For waitForCommand.
+load("jstests/libs/parallel_shell_helpers.js");  // For funWithArgs.
+
+const coll = db.foo;
+coll.drop();
+coll.insert({_id: 0, s: "asdasdasdasdasdasdasd"});
+
+const dataSizeCommand = {
+    "dataSize": "test.foo",
+    "keyPattern": {"_id": 1},
+    "min": {"_id": 0},
+    "max": {"_id": 1}
+};
+
+// Configure the failpoint.
+const failpoint = configureFailPoint(db, "hangBeforeDatasizeCount");
+
+// Launch a parallel shell that runs the dataSize command, that should fail due to interrupt.
+const awaitShell =
+    startParallelShell(funWithArgs(function(cmd) {
+                           assert.commandFailedWithCode(db.runCommand(cmd), ErrorCodes.Interrupted);
+                       }, dataSizeCommand), db.getMongo().port);
+failpoint.wait();
+
+// Find the command opid and kill it.
+const opId = waitForCommand("dataSizeCmd", op => (op["command"]["dataSize"] == "test.foo"), db);
+assert.commandWorked(db.killOp(opId));
+
+// The command is not killed just yet. It will be killed, after releasing the failpoint.
+assert.neq(waitForCommand("dataSizeCmd", op => (op["command"]["dataSize"] == "test.foo"), db), -1);
+
+failpoint.off();
+awaitShell();
+})();
diff --git a/src/mongo/db/commands/dbcommands.cpp b/src/mongo/db/commands/dbcommands.cpp
index debb9c109ed..7ce8384eeeb 100644
--- a/src/mongo/db/commands/dbcommands.cpp
+++ b/src/mongo/db/commands/dbcommands.cpp
@@ -59,6 +59,7 @@
 #include "mongo/db/commands/feature_compatibility_version.h"
 #include "mongo/db/commands/server_status.h"
 #include "mongo/db/concurrency/write_conflict_exception.h"
+#include "mongo/db/curop_failpoint_helpers.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/dbhelpers.h"
@@ -108,6 +109,9 @@
 namespace mongo {
 namespace {
 
+// Will cause 'CmdDatasize' to hang as it starts executing.
+MONGO_FAIL_POINT_DEFINE(hangBeforeDatasizeCount);
+
 /**
  * Returns a CollMod on the underlying buckets collection of the time-series collection.
  * Returns null if 'origCmd' is not for a time-series collection.
@@ -439,7 +443,7 @@ public:
                 return 1;
             }
             exec = InternalPlanner::collectionScan(
-                opCtx, &collection.getCollection(), PlanYieldPolicy::YieldPolicy::NO_YIELD);
+                opCtx, &collection.getCollection(), PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY);
         } else if (min.isEmpty() || max.isEmpty()) {
             errmsg = "only one of min or max specified";
             return false;
@@ -470,9 +474,12 @@ public:
                                               min,
                                               max,
                                               BoundInclusion::kIncludeStartKeyOnly,
-                                              PlanYieldPolicy::YieldPolicy::NO_YIELD);
+                                              PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY);
         }
 
+        CurOpFailpointHelpers::waitWhileFailPointEnabled(
+            &hangBeforeDatasizeCount, opCtx, "hangBeforeDatasizeCount", []() {});
+
         long long avgObjSize = collection->dataSize(opCtx) / numRecords;
 
         long long maxSize = jsobj["maxSize"].numberLong();