SERVER-33282 listen for interrupt while waiting on next event from the ARM

3 files changed, 36 insertions, 21 deletions

diff --git a/jstests/sharding/kill_pinned_cursor.js b/jstests/sharding/kill_pinned_cursor.js
index 9113db5c240..c2d90be1387 100644
--- a/jstests/sharding/kill_pinned_cursor.js
+++ b/jstests/sharding/kill_pinned_cursor.js
@@ -26,7 +26,7 @@
     st.shardColl(coll, {_id: 1}, {_id: 5}, {_id: 6}, kDBName, false);
 
     // Set up the first mongod to hang on a getMore request.
-    let cleanup = null;
+    let getMoreJoiner = null;
     let cursorId;
 
     try {
@@ -51,7 +51,7 @@
         let code = `let cursorId = ${cursorId.toString()};`;
         code += `let collName = "${coll.getName()}";`;
         code += `(${runGetMore.toString()})();`;
-        cleanup = startParallelShell(code, st.s.port);
+        getMoreJoiner = startParallelShell(code, st.s.port);
 
         // Sleep until we know the cursor is pinned on the mongod.
         assert.soon(() => shard0DB.serverStatus().metrics.cursor.open.pinned > 0);
@@ -65,11 +65,29 @@
         assert.eq(cmdRes.cursorsAlive, []);
         assert.eq(cmdRes.cursorsNotFound, []);
         assert.eq(cmdRes.cursorsUnknown, []);
+
+        // The getMore should finish now that we've killed the cursor (even though the failpoint is
+        // still enabled).
+        getMoreJoiner();
+        getMoreJoiner = null;
+
+        // Eventually the cursor should get reaped, at which point the next call to killCursors
+        // should report that nothing was killed.
+        let killRes = null;
+        assert.soon(function() {
+            killRes = mongosDB.runCommand({killCursors: coll.getName(), cursors: [cursorId]});
+            assert.commandWorked(killRes);
+            return killRes.cursorsKilled.length == 0;
+        });
+
+        assert.eq(killRes.cursorsAlive, []);
+        assert.eq(killRes.cursorsNotFound, [cursorId]);
+        assert.eq(killRes.cursorsUnknown, []);
     } finally {
         assert.commandWorked(
             shard0DB.adminCommand({configureFailPoint: kFailPointName, mode: "off"}));
-        if (cleanup) {
-            cleanup();
+        if (getMoreJoiner) {
+            getMoreJoiner();
         }
     }
 
@@ -77,18 +95,5 @@
     // failpoint.
     assert.soon(() => shard0DB.serverStatus().metrics.cursor.open.pinned == 0);
 
-    // Eventually the cursor should get reaped, at which point the next call to killCursors
-    // should report that nothing was killed.
-    let cmdRes = null;
-    assert.soon(function() {
-        cmdRes = mongosDB.runCommand({killCursors: coll.getName(), cursors: [cursorId]});
-        assert.commandWorked(cmdRes);
-        return cmdRes.cursorsKilled.length == 0;
-    });
-
-    assert.eq(cmdRes.cursorsAlive, []);
-    assert.eq(cmdRes.cursorsNotFound, [cursorId]);
-    assert.eq(cmdRes.cursorsUnknown, []);
-
     st.stop();
 })();
diff --git a/src/mongo/executor/task_executor.h b/src/mongo/executor/task_executor.h
index a36c9544ac3..8ef9e732233 100644
--- a/src/mongo/executor/task_executor.h
+++ b/src/mongo/executor/task_executor.h
@@ -186,12 +186,13 @@ public:
     virtual void waitForEvent(const EventHandle& event) = 0;
 
     /**
-     * Same as waitForEvent without an OperationContext, but returns Status::OK with
-     * cv_status::timeout if the event was not triggered within deadline.
+     * Same as waitForEvent without an OperationContext, but if the OperationContext gets
+     * interrupted, will return the kill code, or, if the the deadline passes, will return
+     * Status::OK with cv_status::timeout.
      */
     virtual StatusWith<stdx::cv_status> waitForEvent(OperationContext* opCtx,
                                                      const EventHandle& event,
-                                                     Date_t deadline) = 0;
+                                                     Date_t deadline = Date_t::max()) = 0;
 
     /**
      * Schedules "work" to be run by the executor ASAP.
diff --git a/src/mongo/s/query/router_stage_merge.cpp b/src/mongo/s/query/router_stage_merge.cpp
index 31c200bb004..8ff0dfc991a 100644
--- a/src/mongo/s/query/router_stage_merge.cpp
+++ b/src/mongo/s/query/router_stage_merge.cpp
@@ -52,6 +52,7 @@ StatusWith<ClusterQueryResult> RouterStageMerge::next(ExecContext execCtx) {
 
 StatusWith<ClusterQueryResult> RouterStageMerge::blockForNextNoTimeout(ExecContext execCtx) {
     invariant(_params->tailableMode != TailableMode::kTailableAndAwaitData);
+    invariant(getOpCtx());
     while (!_arm.ready()) {
         auto nextEventStatus = _arm.nextEvent();
         if (!nextEventStatus.isOK()) {
@@ -60,7 +61,15 @@ StatusWith<ClusterQueryResult> RouterStageMerge::blockForNextNoTimeout(ExecConte
         auto event = nextEventStatus.getValue();
 
         // Block until there are further results to return.
-        _executor->waitForEvent(event);
+        auto status = _executor->waitForEvent(getOpCtx(), event);
+
+        if (!status.isOK()) {
+            return status.getStatus();
+        }
+
+        // We have not provided a deadline, so if the wait returns without interruption, we do not
+        // expect to have timed out.
+        invariant(status.getValue() == stdx::cv_status::no_timeout);
     }
 
     return _arm.nextReady();