SERVER-71764 Fix cancellation of hedged operations

5 files changed, 234 insertions, 66 deletions

diff --git a/src/mongo/executor/network_interface_integration_fixture.cpp b/src/mongo/executor/network_interface_integration_fixture.cpp
index 5b83f351f2c..2e50c4d1e39 100644
--- a/src/mongo/executor/network_interface_integration_fixture.cpp
+++ b/src/mongo/executor/network_interface_integration_fixture.cpp
@@ -116,8 +116,7 @@ void NetworkInterfaceIntegrationFixture::startCommand(const TaskExecutor::Callba
 }
 
 Future<RemoteCommandResponse> NetworkInterfaceIntegrationFixture::runCommand(
-    const TaskExecutor::CallbackHandle& cbHandle, RemoteCommandRequest request) {
-    RemoteCommandRequestOnAny rcroa{request};
+    const TaskExecutor::CallbackHandle& cbHandle, RemoteCommandRequestOnAny rcroa) {
 
     _onSchedulingCommand();
 
diff --git a/src/mongo/executor/network_interface_integration_fixture.h b/src/mongo/executor/network_interface_integration_fixture.h
index 30c4fd4132a..3e5eafbabda 100644
--- a/src/mongo/executor/network_interface_integration_fixture.h
+++ b/src/mongo/executor/network_interface_integration_fixture.h
@@ -85,7 +85,7 @@ public:
                       StartCommandCB onFinish);
 
     Future<RemoteCommandResponse> runCommand(const TaskExecutor::CallbackHandle& cbHandle,
-                                             RemoteCommandRequest request);
+                                             RemoteCommandRequestOnAny rcroa);
 
     Future<RemoteCommandOnAnyResponse> runCommandOnAny(const TaskExecutor::CallbackHandle& cbHandle,
                                                        RemoteCommandRequestOnAny request);
diff --git a/src/mongo/executor/network_interface_integration_test.cpp b/src/mongo/executor/network_interface_integration_test.cpp
index 0041b24db14..c2b7c7a6c0a 100644
--- a/src/mongo/executor/network_interface_integration_test.cpp
+++ b/src/mongo/executor/network_interface_integration_test.cpp
@@ -39,6 +39,7 @@
 #include "mongo/executor/network_connection_hook.h"
 #include "mongo/executor/network_interface_integration_fixture.h"
 #include "mongo/executor/test_network_connection_hook.h"
+#include "mongo/logv2/log.h"
 #include "mongo/rpc/factory.h"
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/rpc/message.h"
@@ -196,23 +197,28 @@ public:
                             << "secs" << 1000000000);
     }
 
-    /**
-     * Returns true if the given command is still running.
-     */
-    bool isCommandRunning(const std::string command) {
+    RemoteCommandResponse runCurrentOpForCommand(HostAndPort target, const std::string command) {
         const auto cmdObj =
             BSON("aggregate" << 1 << "pipeline"
                              << BSON_ARRAY(BSON("$currentOp" << BSON("localOps" << true))
                                            << BSON("$match" << BSON(("command." + command)
                                                                     << BSON("$exists" << true))))
-                             << "cursor" << BSONObj());
-        auto cs = fixture();
-        RemoteCommandRequest request{
-            cs.getServers().front(), "admin", cmdObj, BSONObj(), nullptr, kNoTimeout};
+                             << "cursor" << BSONObj() << "$readPreference"
+                             << BSON("mode"
+                                     << "nearest"));
+        RemoteCommandRequest request{target, "admin", cmdObj, BSONObj(), nullptr, kNoTimeout};
         auto res = runCommandSync(request);
-
         ASSERT_OK(res.status);
         ASSERT_OK(getStatusFromCommandResult(res.data));
+        return res;
+    }
+
+    /**
+     * Returns true if the given command is still running.
+     */
+    bool isCommandRunning(const std::string command,
+                          boost::optional<HostAndPort> target = boost::none) {
+        auto res = runCurrentOpForCommand(target.value_or(fixture().getServers().front()), command);
         return !res.data["cursor"]["firstBatch"].Array().empty();
     }
 
@@ -750,6 +756,160 @@ TEST_F(NetworkInterfaceTest, SetAlarm) {
     ASSERT_FALSE(swResult.isOK());
 }
 
+class HedgeCancellationTest : public NetworkInterfaceTest {
+public:
+    enum class CancellationMode { kAfterCompletion, kAfterScheduling };
+
+    void runTest(CancellationMode mode) {
+        if (fixture().type() != ConnectionString::ConnectionType::kReplicaSet) {
+            LOGV2(7176401, "Skipped: this test may only run against a replica-set");
+            return;
+        }
+
+        auto cbh = makeCallbackHandle();
+        auto future = [&] {
+            _blockCommandsOnAllServers(BSON_ARRAY("echo"
+                                                  << "_killOperations"));
+            ON_BLOCK_EXIT([&] { _unblockCommandsOnAllServers({"echo", "_killOperations"}); });
+
+            auto future = _scheduleHedgedEcho(cbh);
+
+            _waitForServersToStartRunningEcho();
+
+            if (mode == CancellationMode::kAfterScheduling) {
+                net().cancelCommand(cbh);
+            } else {
+                // Let the first node in the list of servers proceed with running the command by
+                // killing the blocked `echo` operation. This results in the completion of the
+                // operation and cancels all pending hedged operations.
+                _killRemoteOps(fixture().getServers().front(), "echo");
+            }
+
+            _waitForServersToStartRunningKillOperations(mode);
+
+            return future;
+        }();
+
+        LOGV2(7176402, "Wait for the remote command to finish");
+        std::move(future).ignoreValue().get();
+    }
+
+private:
+    void _runCommand(const HostAndPort& server, std::string db, BSONObj cmd) {
+        RemoteCommandRequest request{server, db, cmd, BSONObj(), nullptr, kNoTimeout};
+        request.sslMode = transport::kGlobalSSLMode;
+        auto res = runCommandSync(request);
+        ASSERT_OK(res.status);
+        ASSERT_OK(getStatusFromCommandResult(res.data));
+        ASSERT(!res.data["writeErrors"]);
+    }
+
+    void _configureFailPoint(const HostAndPort& server,
+                             std::string fpName,
+                             bool enable,
+                             BSONObj data = BSONObj()) {
+        BSONObjBuilder bob;
+        bob.append("configureFailPoint", fpName);
+        bob.append("mode", enable ? "alwaysOn" : "off");
+        if (!data.isEmpty())
+            bob.append("data", std::move(data));
+        _runCommand(server, "admin", bob.obj());
+    }
+
+    void _blockCommandsOnAllServers(BSONArray cmds) {
+        auto servers = fixture().getServers();
+        for (const auto& server : servers) {
+            _configureFailPoint(server,
+                                "failCommand",
+                                true,
+                                BSON("blockConnection" << true << "blockTimeMS" << 1000000
+                                                       << "failCommands" << cmds));
+        }
+    }
+
+    void _killRemoteOps(HostAndPort server, std::string cmd) {
+        auto res = runCurrentOpForCommand(server, cmd);
+        for (auto& op : res.data["cursor"]["firstBatch"].Array()) {
+            auto opid = op.Obj()["opid"];
+            _runCommand(server, "admin", BSON("killOp" << 1 << "op" << opid));
+        }
+    }
+
+    void _unblockCommandsOnAllServers(std::vector<std::string> cmds) {
+        LOGV2(7176403, "Disabling fail-points to unblock commands");
+        auto servers = fixture().getServers();
+        for (auto& server : servers) {
+            // Must kill (and unblock) the remote operations blocked behind the `failCommand`
+            // fail-point (if still running) before disabling it, otherwise it will hang forever.
+            for (auto& cmd : cmds) {
+                _killRemoteOps(server, cmd);
+            }
+            _configureFailPoint(server, "failCommand", false);
+        }
+    }
+
+    Future<RemoteCommandResponse> _scheduleHedgedEcho(const TaskExecutor::CallbackHandle& cbh) {
+        RemoteCommandRequest::Options rcrOptions;
+        rcrOptions.hedgeOptions.isHedgeEnabled = true;
+        rcrOptions.hedgeOptions.hedgeCount = fixture().getServers().size();
+        RemoteCommandRequestOnAny rcr(fixture().getServers(),
+                                      "admin",
+                                      makeEchoCmdObj(),
+                                      BSONObj(),
+                                      nullptr,
+                                      kNoTimeout,
+                                      std::move(rcrOptions));
+        // Only internal clients can run hedged operations.
+        resetIsInternalClient(true);
+        ON_BLOCK_EXIT([&] { resetIsInternalClient(false); });
+        LOGV2(7176404, "Scheduling the remote command");
+        return runCommand(cbh, std::move(rcr));
+    }
+
+    void _waitForServerToRunCommand(const HostAndPort& server, std::string command) {
+        ClockSource::StopWatch stopwatch;
+        while (!isCommandRunning(command, server) && stopwatch.elapsed() < kMaxWait) {
+            sleepmillis(100);
+        }
+    }
+
+    void _waitForServersToStartRunningEcho() {
+        LOGV2(7176405, "Waiting for all servers to start running the command");
+        const auto cmd = "echo";
+        auto servers = fixture().getServers();
+        for (auto& server : servers) {
+            _waitForServerToRunCommand(server, cmd);
+            ASSERT_TRUE(isCommandRunning(cmd, server));
+        }
+    }
+
+    void _waitForServersToStartRunningKillOperations(CancellationMode mode) {
+        LOGV2(7176406, "Wait for servers to receive $_killOperations");
+        const auto cmd = "_killOperations";
+        auto servers = fixture().getServers();
+        size_t idx = (mode == CancellationMode::kAfterCompletion) ? 1 : 0;
+        for (; idx < servers.size(); idx++) {
+            _waitForServerToRunCommand(servers[idx], cmd);
+            ASSERT_TRUE(isCommandRunning(cmd, servers[idx]));
+        }
+    }
+};
+
+TEST_F(HedgeCancellationTest, CancelAfterScheduling) {
+    // Cancel the hedged operation after it is scheduled on all targets and before completion.
+    // We should send `_killOperations` to all targets that have already acquired a connection
+    // and might have started/completed running the operation.
+    runTest(CancellationMode::kAfterScheduling);
+}
+
+TEST_F(HedgeCancellationTest, CancelAfterCompletion) {
+    // Waits until the hedged operation is scheduled on all targets, then cancels all pending
+    // operations after the first scheduled operation completes. We should send
+    // `_killOperations` to all targets except for the one used to fulfill the final promise
+    // (i.e. complete the operation).
+    runTest(CancellationMode::kAfterCompletion);
+}
+
 TEST_F(NetworkInterfaceInternalClientTest,
        IsMasterRequestContainsOutgoingWireVersionInternalClientInfo) {
     auto deferred = runCommand(makeCallbackHandle(), makeTestCommand(kNoTimeout, makeEchoCmdObj()));
diff --git a/src/mongo/executor/network_interface_tl.cpp b/src/mongo/executor/network_interface_tl.cpp
index f8fd8c1b7ab..cfc7d4b1452 100644
--- a/src/mongo/executor/network_interface_tl.cpp
+++ b/src/mongo/executor/network_interface_tl.cpp
@@ -739,69 +739,57 @@ void NetworkInterfaceTL::CommandState::fulfillFinalPromise(
 }
 
 NetworkInterfaceTL::RequestManager::RequestManager(CommandStateBase* cmdState_)
-    : cmdState{cmdState_},
-      requests(cmdState->maxConcurrentRequests(), std::weak_ptr<RequestState>()) {}
+    : cmdState{cmdState_}, requests(cmdState->maxConcurrentRequests()) {}
 
 void NetworkInterfaceTL::RequestManager::cancelRequests() {
+    std::vector<std::shared_ptr<RequestState>> requestsToCancel;
     {
         stdx::lock_guard<Latch> lk(mutex);
         isLocked = true;
 
-        if (sentIdx == 0) {
-            // We've canceled before any connections were acquired.
-            return;
+        for (size_t i = 0; i < sentIdx; i++) {
+            requestsToCancel.push_back(requests[i].request.lock());
         }
     }
 
-    for (size_t i = 0; i < requests.size(); i++) {
-        // This may cause the connection to be discarded before it receives the response to an
-        // earlier `_killOperations` command.
-        if (auto requestState = requests[i].lock()) {
+    for (size_t i = 0; i < requestsToCancel.size(); i++) {
+        if (auto& request = requestsToCancel[i]) {
+            // For hedged operations, we send `_killOperations` out-of-band, and the following may
+            // close the connection (used to send the original command) before it receives the
+            // response from the `_killOperations`.
             LOGV2_DEBUG(4646301,
                         2,
                         "Cancelling request",
                         "requestId"_attr = cmdState->requestOnAny.id,
                         "index"_attr = i);
-            requestState->cancel();
+            request->cancel();
+            request.reset();
         }
     }
 }
 
 void NetworkInterfaceTL::RequestManager::killOperationsForPendingRequests() {
+    // Send `_killOperation` out of band to all targets with initialized requests (i.e., those who
+    // acquired a connection), regardless of their state so long as they are not used to fulfill the
+    // operation. The following will hold indices for targets in the initial remote command request.
+    std::vector<size_t> indices;
     {
         stdx::lock_guard<Latch> lk(mutex);
         isLocked = true;
 
-        if (sentIdx == 0) {
-            // We've canceled before any connections were acquired.
-            return;
+        for (size_t i = 0; i < sentIdx; i++) {
+            auto& context = requests[i];
+            invariant(context.initialized);
+            if (auto requestState = context.request.lock();
+                requestState && requestState->fulfilledPromise) {
+                continue;  // This request is used to fulfill the promise.
+            }
+            indices.push_back(context.idx);
         }
     }
 
-    for (size_t i = 0; i < requests.size(); i++) {
-        auto requestState = requests[i].lock();
-        if (!requestState || requestState->fulfilledPromise) {
-            continue;
-        }
-
-        auto conn = requestState->weakConn.lock();
-        if (!conn) {
-            // If there is nothing from weakConn, the networking has already finished.
-            continue;
-        }
-
-        // If the request was sent, send a remote command request to the target host
-        // to kill the operation started by the request.
-
-        LOGV2_DEBUG(4664801,
-                    2,
-                    "Sending remote _killOperations request to cancel command",
-                    "operationKey"_attr = cmdState->operationKey,
-                    "target"_attr = requestState->request->target,
-                    "requestId"_attr = requestState->request->id);
-
-        auto status = requestState->interface()->_killOperation(requestState, this);
-        if (!status.isOK()) {
+    for (auto idx : indices) {
+        if (auto status = cmdState->interface->_killOperation(cmdState, idx); !status.isOK()) {
             LOGV2_DEBUG(4664810, 2, "Failed to send remote _killOperations", "error"_attr = status);
         }
     }
@@ -893,7 +881,7 @@ void NetworkInterfaceTL::RequestManager::trySend(
 
         auto currentSentIdx = sentIdx++;
 
-        requestState = std::make_shared<RequestState>(this, cmdState->shared_from_this(), idx);
+        requestState = std::make_shared<RequestState>(this, cmdState->shared_from_this());
         requestState->isHedge = currentSentIdx > 0;
 
         // Set conn/weakConn+request under the lock so they will always be observed during cancel.
@@ -931,7 +919,10 @@ void NetworkInterfaceTL::RequestManager::trySend(
             request->cmdObj = updatedCmdBuilder.obj();
         }
 
-        requests.at(currentSentIdx) = requestState;
+        auto& context = requests.at(currentSentIdx);
+        context.initialized = true;
+        context.idx = currentSentIdx;
+        context.request = requestState;
     }
 
     LOGV2_DEBUG(4646300,
@@ -1246,16 +1237,19 @@ void NetworkInterfaceTL::cancelCommand(const TaskExecutor::CallbackHandle& cbHan
                        << redact(cmdStateToCancel->requestOnAny.toString())});
 }
 
-Status NetworkInterfaceTL::_killOperation(std::shared_ptr<RequestState> requestStateToKill,
-                                          RequestManager* requestManager) try {
+Status NetworkInterfaceTL::_killOperation(CommandStateBase* cmdStateToKill, size_t idx) try {
     auto [target, sslMode] = [&] {
-        stdx::lock_guard<Latch> lk(requestManager->mutex);
-        invariant(requestStateToKill->request);
-        auto request = requestStateToKill->request.value();
-        return std::make_pair(request.target, request.sslMode);
+        const auto& request = cmdStateToKill->requestOnAny;
+        return std::make_pair(request.target[idx], request.sslMode);
     }();
-    auto cmdStateToKill = requestStateToKill->cmdState;
+
     auto operationKey = cmdStateToKill->operationKey.value();
+    LOGV2_DEBUG(4664801,
+                2,
+                "Sending remote _killOperations request to cancel command",
+                "operationKey"_attr = operationKey,
+                "target"_attr = target,
+                "requestId"_attr = idx);
 
     // Make a request state for _killOperations.
     executor::RemoteCommandRequest killOpRequest(
diff --git a/src/mongo/executor/network_interface_tl.h b/src/mongo/executor/network_interface_tl.h
index 0409303e1ca..ca0cb1b8a72 100644
--- a/src/mongo/executor/network_interface_tl.h
+++ b/src/mongo/executor/network_interface_tl.h
@@ -107,6 +107,15 @@ private:
     struct RequestState;
     struct RequestManager;
 
+    /**
+     * For each logical RPC, an instance of `CommandState` is created to capture the state of the
+     * remote command. As part of running a remote command, `NITL` sends out one or more requests
+     * to the specified targets, and `RequestState` represents the state of each request.
+     * `CommandState` owns a `RequestManager` that tracks individual requests. For each request sent
+     * over the wire, `RequestManager` creates a `Context` that holds a weak pointer to the
+     * `Request`, as well as the index of the target.
+     */
+
     struct CommandStateBase : public std::enable_shared_from_this<CommandStateBase> {
         CommandStateBase(NetworkInterfaceTL* interface_,
                          RemoteCommandRequestOnAny request_,
@@ -246,7 +255,17 @@ private:
         void killOperationsForPendingRequests();
 
         CommandStateBase* cmdState;
-        std::vector<std::weak_ptr<RequestState>> requests;
+
+        /**
+         * Holds context for individual requests, and is only valid if initialized.
+         * `idx` maps the request to its target in the corresponding `cmdState`.
+         */
+        struct Context {
+            bool initialized = false;
+            size_t idx;
+            std::weak_ptr<RequestState> request;
+        };
+        std::vector<Context> requests;
 
         Mutex mutex = MONGO_MAKE_LATCH("NetworkInterfaceTL::RequestManager::mutex");
 
@@ -263,8 +282,8 @@ private:
     struct RequestState final : public std::enable_shared_from_this<RequestState> {
         using ConnectionHandle = std::shared_ptr<ConnectionPool::ConnectionHandle::element_type>;
         using WeakConnectionHandle = std::weak_ptr<ConnectionPool::ConnectionHandle::element_type>;
-        RequestState(RequestManager* mgr, std::shared_ptr<CommandStateBase> cmdState_, size_t id)
-            : cmdState{std::move(cmdState_)}, requestManager(mgr), reqId(id) {}
+        RequestState(RequestManager* mgr, std::shared_ptr<CommandStateBase> cmdState_)
+            : cmdState{std::move(cmdState_)}, requestManager(mgr) {}
 
         ~RequestState();
 
@@ -305,9 +324,6 @@ private:
         ConnectionHandle conn;
         WeakConnectionHandle weakConn;
 
-        // Internal id of this request as tracked by the RequestManager.
-        size_t reqId;
-
         // True if this request is an additional request sent to hedge the operation.
         bool isHedge{false};
 
@@ -340,8 +356,7 @@ private:
 
     void _run();
 
-    Status _killOperation(std::shared_ptr<RequestState> requestStateToKill,
-                          RequestManager* requestManager);
+    Status _killOperation(CommandStateBase* cmdStateToKill, size_t idx);
 
     std::string _instanceName;
     ServiceContext* _svcCtx = nullptr;