SERVER-39150 markKillOnClientDisconnect

16 files changed, 93 insertions, 0 deletions

diff --git a/src/mongo/db/auth/sasl_commands.cpp b/src/mongo/db/auth/sasl_commands.cpp
index affabc7450d..19cf85ee004 100644
--- a/src/mongo/db/auth/sasl_commands.cpp
+++ b/src/mongo/db/auth/sasl_commands.cpp
@@ -276,6 +276,7 @@ bool CmdSaslStart::run(OperationContext* opCtx,
                        const std::string& db,
                        const BSONObj& cmdObj,
                        BSONObjBuilder& result) {
+    opCtx->markKillOnClientDisconnect();
     Client* client = opCtx->getClient();
     AuthenticationSession::set(client, std::unique_ptr<AuthenticationSession>());
 
@@ -310,6 +311,7 @@ bool CmdSaslContinue::run(OperationContext* opCtx,
                           const std::string& db,
                           const BSONObj& cmdObj,
                           BSONObjBuilder& result) {
+    opCtx->markKillOnClientDisconnect();
     Client* client = Client::getCurrent();
     std::unique_ptr<AuthenticationSession> sessionGuard;
     AuthenticationSession::swap(client, sessionGuard);
diff --git a/src/mongo/db/commands.cpp b/src/mongo/db/commands.cpp
index 90b16d671d1..c4323d6a36f 100644
--- a/src/mongo/db/commands.cpp
+++ b/src/mongo/db/commands.cpp
@@ -55,6 +55,7 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/server_parameters.h"
 #include "mongo/rpc/factory.h"
+#include "mongo/rpc/metadata/client_metadata_ismaster.h"
 #include "mongo/rpc/op_msg_rpc_impls.h"
 #include "mongo/rpc/protocol.h"
 #include "mongo/rpc/write_concern_error_detail.h"
@@ -470,6 +471,7 @@ Status CommandHelpers::canUseTransactions(StringData dbName, StringData cmdName)
 constexpr StringData CommandHelpers::kHelpFieldName;
 
 MONGO_FAIL_POINT_DEFINE(failCommand);
+MONGO_FAIL_POINT_DEFINE(waitInCommandMarkKillOnClientDisconnect);
 
 bool CommandHelpers::shouldActivateFailCommandFailPoint(const BSONObj& data,
                                                         StringData cmdName,
@@ -530,6 +532,28 @@ void CommandHelpers::evaluateFailCommandFailPoint(OperationContext* opCtx, Strin
     }
 }
 
+void CommandHelpers::handleMarkKillOnClientDisconnect(OperationContext* opCtx,
+                                                      bool shouldMarkKill) {
+    if (opCtx->getClient()->isInDirectClient()) {
+        return;
+    }
+
+    if (shouldMarkKill) {
+        opCtx->markKillOnClientDisconnect();
+    }
+
+    MONGO_FAIL_POINT_BLOCK_IF(
+        waitInCommandMarkKillOnClientDisconnect, options, [&](const BSONObj& obj) {
+            const auto& clientMetadata =
+                ClientMetadataIsMasterState::get(opCtx->getClient()).getClientMetadata();
+
+            return clientMetadata && (clientMetadata->getApplicationName() == obj["appName"].str());
+        }) {
+        MONGO_FAIL_POINT_PAUSE_WHILE_SET_OR_INTERRUPTED(opCtx,
+                                                        waitInCommandMarkKillOnClientDisconnect);
+    }
+}
+
 //////////////////////////////////////////////////////////////
 // CommandInvocation
 
diff --git a/src/mongo/db/commands.h b/src/mongo/db/commands.h
index ef827bf4904..c8eeaa2958d 100644
--- a/src/mongo/db/commands.h
+++ b/src/mongo/db/commands.h
@@ -56,6 +56,7 @@
 namespace mongo {
 
 MONGO_FAIL_POINT_DECLARE(failCommand);
+MONGO_FAIL_POINT_DECLARE(waitInCommandMarkKillOnClientDisconnect);
 
 class Command;
 class CommandInvocation;
@@ -228,6 +229,12 @@ struct CommandHelpers {
      * Possibly uasserts according to the "failCommand" fail point.
      */
     static void evaluateFailCommandFailPoint(OperationContext* opCtx, StringData commandName);
+
+    /**
+     * Handles marking kill on client disconnect.
+     */
+    static void handleMarkKillOnClientDisconnect(OperationContext* opCtx,
+                                                 bool shouldMarkKill = true);
 };
 
 /**
diff --git a/src/mongo/db/commands/authentication_commands.cpp b/src/mongo/db/commands/authentication_commands.cpp
index b5c6ac9c4c6..7aeaf8d279c 100644
--- a/src/mongo/db/commands/authentication_commands.cpp
+++ b/src/mongo/db/commands/authentication_commands.cpp
@@ -232,6 +232,7 @@ public:
              const std::string&,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) final {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         auto n = getNextNonce();
         std::stringstream ss;
         ss << std::hex << n;
@@ -253,6 +254,7 @@ bool CmdAuthenticate::run(OperationContext* opCtx,
                           const std::string& dbname,
                           const BSONObj& cmdObj,
                           BSONObjBuilder& result) {
+    CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
     if (!serverGlobalParams.quiet.load()) {
         mutablebson::Document cmdToLog(cmdObj, mutablebson::Document::kInPlaceDisabled);
         log() << " authenticate db: " << dbname << " " << cmdToLog;
diff --git a/src/mongo/db/commands/count_cmd.cpp b/src/mongo/db/commands/count_cmd.cpp
index 818ea1400e4..6b80f8016d7 100644
--- a/src/mongo/db/commands/count_cmd.cpp
+++ b/src/mongo/db/commands/count_cmd.cpp
@@ -180,6 +180,7 @@ public:
              const string& dbname,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) override {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         // Acquire locks and resolve possible UUID. The RAII object is optional, because in the case
         // of a view, the locks need to be released.
         boost::optional<AutoGetCollectionForReadCommand> ctx;
diff --git a/src/mongo/db/commands/distinct.cpp b/src/mongo/db/commands/distinct.cpp
index f26feb2cf80..580f15206a6 100644
--- a/src/mongo/db/commands/distinct.cpp
+++ b/src/mongo/db/commands/distinct.cpp
@@ -168,6 +168,7 @@ public:
              const std::string& dbname,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) override {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         // Acquire locks and resolve possible UUID. The RAII object is optional, because in the case
         // of a view, the locks need to be released.
         boost::optional<AutoGetCollectionForReadCommand> ctx;
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index e9a20684a6b..a73aef6375a 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -229,6 +229,7 @@ public:
          *   --Generate response to send to the client.
          */
         void run(OperationContext* opCtx, rpc::ReplyBuilderInterface* result) {
+            CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
             // Although it is a command, a find command gets counted as a query.
             globalOpCounters.gotQuery();
             ServerReadConcernMetrics::get(opCtx)->recordReadConcern(
diff --git a/src/mongo/db/commands/list_collections.cpp b/src/mongo/db/commands/list_collections.cpp
index a702a8f6807..df886abf938 100644
--- a/src/mongo/db/commands/list_collections.cpp
+++ b/src/mongo/db/commands/list_collections.cpp
@@ -247,6 +247,7 @@ public:
              const string& dbname,
              const BSONObj& jsobj,
              BSONObjBuilder& result) final {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         unique_ptr<MatchExpression> matcher;
         const auto as = AuthorizationSession::get(opCtx->getClient());
 
diff --git a/src/mongo/db/commands/list_databases.cpp b/src/mongo/db/commands/list_databases.cpp
index 7e2adf2ee1b..c212403535f 100644
--- a/src/mongo/db/commands/list_databases.cpp
+++ b/src/mongo/db/commands/list_databases.cpp
@@ -96,6 +96,7 @@ public:
              const string& dbname,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) final {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         IDLParserErrorContext ctx("listDatabases");
         auto cmd = ListDatabasesCommand::parse(ctx, cmdObj);
         auto* as = AuthorizationSession::get(opCtx->getClient());
diff --git a/src/mongo/db/commands/list_indexes.cpp b/src/mongo/db/commands/list_indexes.cpp
index 2210639210d..8c4b15606dc 100644
--- a/src/mongo/db/commands/list_indexes.cpp
+++ b/src/mongo/db/commands/list_indexes.cpp
@@ -128,6 +128,7 @@ public:
              const string& dbname,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         const long long defaultBatchSize = std::numeric_limits<long long>::max();
         long long batchSize;
         uassertStatusOK(
diff --git a/src/mongo/db/commands/pipeline_command.cpp b/src/mongo/db/commands/pipeline_command.cpp
index ef5dc99606b..06a91ad79f0 100644
--- a/src/mongo/db/commands/pipeline_command.cpp
+++ b/src/mongo/db/commands/pipeline_command.cpp
@@ -97,6 +97,9 @@ public:
         }
 
         void run(OperationContext* opCtx, rpc::ReplyBuilderInterface* reply) override {
+            CommandHelpers::handleMarkKillOnClientDisconnect(
+                opCtx, !Pipeline::aggSupportsWriteConcern(_request.body));
+
             const auto aggregationRequest = uassertStatusOK(
                 AggregationRequest::parseFromBSON(_dbName, _request.body, boost::none));
             uassertStatusOK(runAggregate(opCtx,
diff --git a/src/mongo/db/dbdirectclient.cpp b/src/mongo/db/dbdirectclient.cpp
index 0bd20642df9..c3702258c58 100644
--- a/src/mongo/db/dbdirectclient.cpp
+++ b/src/mongo/db/dbdirectclient.cpp
@@ -174,6 +174,7 @@ unique_ptr<DBClientCursor> DBDirectClient::query(const NamespaceStringOrUUID& ns
 
 unsigned long long DBDirectClient::count(
     const string& ns, const BSONObj& query, int options, int limit, int skip) {
+    DirectClientScope directClientScope(_opCtx);
     BSONObj cmdObj = _countCmd(ns, query, options, limit, skip);
 
     NamespaceString nsString(ns);
diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index 319b803288c..b5bc7821562 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -229,6 +229,20 @@ Status OperationContext::checkForInterruptNoAssert() noexcept {
         return Status(killStatus, "operation was interrupted");
     }
 
+    if (_markKillOnClientDisconnect) {
+        const auto now = getServiceContext()->getFastClockSource()->now();
+
+        if (now > _lastClientCheck + Milliseconds(500)) {
+            _lastClientCheck = now;
+
+            if (!getClient()->session()->isConnected()) {
+                markKilled(ErrorCodes::ClientDisconnect);
+                return Status(ErrorCodes::ClientDisconnect,
+                              "operation was interrupted because a client disconnected");
+            }
+        }
+    }
+
     return Status::OK();
 }
 
@@ -332,6 +346,11 @@ StatusWith<stdx::cv_status> OperationContext::waitForConditionOrInterruptNoAsser
 
 void OperationContext::markKilled(ErrorCodes::Error killCode) {
     invariant(killCode != ErrorCodes::OK);
+
+    if (killCode == ErrorCodes::ClientDisconnect) {
+        log() << "operation was interrupted because a client disconnected";
+    }
+
     stdx::unique_lock<stdx::mutex> lkWaitMutex;
 
     // If we have a _waitMutex, it means this opCtx is currently blocked in
@@ -362,6 +381,26 @@ void OperationContext::markKilled(ErrorCodes::Error killCode) {
     }
 }
 
+void OperationContext::markKillOnClientDisconnect() {
+    if (getClient()->isInDirectClient()) {
+        return;
+    }
+
+    if (_markKillOnClientDisconnect) {
+        return;
+    }
+
+    if (getClient() && getClient()->session()) {
+        _lastClientCheck = getServiceContext()->getFastClockSource()->now();
+
+        _markKillOnClientDisconnect = true;
+
+        if (_baton) {
+            _baton->markKillOnClientDisconnect();
+        }
+    }
+}
+
 void OperationContext::setLogicalSessionId(LogicalSessionId lsid) {
     invariant(!_lsid);
     _lsid = std::move(lsid);
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index f376a98e048..c600218da4f 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -237,6 +237,8 @@ public:
         return _writesAreReplicated;
     }
 
+    void markKillOnClientDisconnect();
+
     /**
      * Marks this operation as killed so that subsequent calls to checkForInterrupt and
      * checkForInterruptNoAssert by the thread executing the operation will start returning the
@@ -463,6 +465,8 @@ private:
     ErrorCodes::Error _timeoutError = ErrorCodes::ExceededTimeLimit;
     bool _ignoreInterrupts = false;
     bool _hasArtificialDeadline = false;
+    bool _markKillOnClientDisconnect = false;
+    Date_t _lastClientCheck;
 
     // Max operation time requested by the user or by the cursor in the case of a getMore with no
     // user-specified maxTime. This is tracked with microsecond granularity for the purpose of
diff --git a/src/mongo/db/repl/replication_info.cpp b/src/mongo/db/repl/replication_info.cpp
index c356d23e76f..bda88bab094 100644
--- a/src/mongo/db/repl/replication_info.cpp
+++ b/src/mongo/db/repl/replication_info.cpp
@@ -237,6 +237,7 @@ public:
              const string&,
              const BSONObj& cmdObj,
              BSONObjBuilder& result) final {
+        CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         /* currently request to arbiter is (somewhat arbitrarily) an ismaster request that is not
            authenticated.
         */
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index ee40ef54b4b..e9519e89b80 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -546,7 +546,9 @@ void execCommandDatabase(OperationContext* opCtx,
     CommandHelpers::uassertShouldAttemptParse(opCtx, command, request);
     BSONObjBuilder extraFieldsBuilder;
     auto startOperationTime = getClientOperationTime(opCtx);
+
     auto invocation = command->parse(opCtx, request);
+
     OperationSessionInfoFromClient sessionOptions;
 
     try {
@@ -1203,6 +1205,8 @@ DbResponse ServiceEntryPointCommon::handleRequest(OperationContext* opCtx,
         dbresponse = receivedCommands(opCtx, m, behaviors);
     } else if (op == dbQuery) {
         invariant(!isCommand);
+        opCtx->markKillOnClientDisconnect();
+
         dbresponse = receivedQuery(opCtx, nsString, c, m, behaviors);
     } else if (op == dbGetMore) {
         dbresponse = receivedGetMore(opCtx, m, currentOp, &forceLog);