SERVER-44167 Added ability to kill operations by key

There are two patches here really. One of which makes killOp fast to use
and visible. The other adds OperationKey to various places and maps it
to an internal OpId.

18 files changed, 511 insertions, 90 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 1e1ad7b72dd..18f5e8e6795 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -402,6 +402,7 @@ env.Library(
         'default_baton.cpp',
         'operation_context.cpp',
         'operation_context_group.cpp',
+        'operation_key_manager.cpp',
         'service_context.cpp',
         'server_recovery.cpp',
         'unclean_shutdown.cpp',
@@ -1385,6 +1386,21 @@ env.Library(
 )
 
 env.Library(
+    target='operation_killer',
+    source=[
+        'operation_killer.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+        'service_context',
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/auth/auth',
+        '$BUILD_DIR/mongo/db/auth/authprivilege',
+    ],
+)
+
+env.Library(
     target='periodic_runner_job_abort_expired_transactions',
     source=[
         'periodic_runner_job_abort_expired_transactions.cpp',
diff --git a/src/mongo/db/command_generic_argument.cpp b/src/mongo/db/command_generic_argument.cpp
index 5969b9e1e0f..d03b4e64eae 100644
--- a/src/mongo/db/command_generic_argument.cpp
+++ b/src/mongo/db/command_generic_argument.cpp
@@ -52,7 +52,7 @@ struct SpecialArgRecord {
 // If that changes, it should be added. When you add to this list, consider whether you
 // should also change the filterCommandRequestForPassthrough() function.
 // clang-format off
-static constexpr std::array<SpecialArgRecord, 27> specials{{
+static constexpr std::array<SpecialArgRecord, 28> specials{{
     //                                       /-isGeneric
     //                                       |  /-stripFromRequest
     //                                       |  |  /-stripFromReply
@@ -73,6 +73,7 @@ static constexpr std::array<SpecialArgRecord, 27> specials{{
     {"tracking_info"_sd,                     1, 1, 0},
     {"writeConcern"_sd,                      1, 0, 0},
     {"lsid"_sd,                              1, 0, 0},
+    {"clientOperationKey"_sd,                1, 0, 0},
     {"txnNumber"_sd,                         1, 0, 0},
     {"autocommit"_sd,                        1, 0, 0},
     {"coordinator"_sd,                       1, 0, 0},
diff --git a/src/mongo/db/commands/SConscript b/src/mongo/db/commands/SConscript
index f309978cb67..ee4954c8654 100644
--- a/src/mongo/db/commands/SConscript
+++ b/src/mongo/db/commands/SConscript
@@ -438,6 +438,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/auth/authprivilege',
         '$BUILD_DIR/mongo/db/audit',
         '$BUILD_DIR/mongo/db/commands',
+        '$BUILD_DIR/mongo/db/operation_killer',
         '$BUILD_DIR/mongo/db/auth/authorization_manager_global',
         '$BUILD_DIR/mongo/db/query/command_request_response',
     ],
diff --git a/src/mongo/db/commands/kill_op_cmd_base.cpp b/src/mongo/db/commands/kill_op_cmd_base.cpp
index 71b23d52424..020d740bc92 100644
--- a/src/mongo/db/commands/kill_op_cmd_base.cpp
+++ b/src/mongo/db/commands/kill_op_cmd_base.cpp
@@ -32,93 +32,44 @@
 #include "mongo/db/commands/kill_op_cmd_base.h"
 
 #include "mongo/bson/util/bson_extract.h"
-#include "mongo/db/audit.h"
-#include "mongo/db/auth/authorization_session.h"
-#include "mongo/db/client.h"
-#include "mongo/db/operation_context.h"
+#include "mongo/db/operation_killer.h"
 #include "mongo/util/str.h"
 
 namespace mongo {
 
-Status KillOpCmdBase::checkAuthForCommand(Client* client,
+Status KillOpCmdBase::checkAuthForCommand(Client* worker,
                                           const std::string& dbname,
                                           const BSONObj& cmdObj) const {
-    AuthorizationSession* authzSession = AuthorizationSession::get(client);
+    auto opKiller = OperationKiller(worker);
 
-    if (authzSession->isAuthorizedForActionsOnResource(ResourcePattern::forClusterResource(),
-                                                       ActionType::killop)) {
-        // If we have administrative permission to run killop, we don't need to traverse the
-        // Client list to figure out if we own the operation which will be terminated.
+    if (opKiller.isGenerallyAuthorizedToKill()) {
         return Status::OK();
     }
 
-    if (authzSession->isAuthenticated() && isKillingLocalOp(cmdObj.getField("op"))) {
+    if (isKillingLocalOp(cmdObj.getField("op"))) {
         // Look up the OperationContext and see if we have permission to kill it. This is done once
         // here and again in the command body. The check here in the checkAuthForCommand() function
         // is necessary because if the check fails, it will be picked up by the auditing system.
         long long opId = parseOpId(cmdObj);
-        auto lkAndOp = KillOpCmdBase::findOpForKilling(client, opId);
-        if (lkAndOp) {
-            // We were able to find the Operation, and we were authorized to interact with it.
+        auto target = worker->getServiceContext()->getLockedClient(opId);
+
+        if (OperationKiller(worker).isAuthorizedToKill(target)) {
+            // We were authorized to interact with the target Client
             return Status::OK();
         }
     }
+
     return Status(ErrorCodes::Unauthorized, "Unauthorized");
 }
 
+void KillOpCmdBase::killLocalOperation(OperationContext* opCtx, OperationId opToKill) {
+    OperationKiller(opCtx->getClient()).killOperation(opToKill);
+}
 
 bool KillOpCmdBase::isKillingLocalOp(const BSONElement& opElem) {
     return opElem.isNumber();
 }
 
-boost::optional<std::tuple<stdx::unique_lock<Client>, OperationContext*>>
-KillOpCmdBase::findOperationContext(ServiceContext* serviceContext, unsigned int opId) {
-    for (ServiceContext::LockedClientsCursor cursor(serviceContext);
-         Client* opClient = cursor.next();) {
-        stdx::unique_lock<Client> lk(*opClient);
-
-        OperationContext* opCtx = opClient->getOperationContext();
-        if (opCtx && opCtx->getOpID() == opId) {
-            return {std::make_tuple(std::move(lk), opCtx)};
-        }
-    }
-
-    return boost::none;
-}
-
-boost::optional<std::tuple<stdx::unique_lock<Client>, OperationContext*>>
-KillOpCmdBase::findOpForKilling(Client* client, unsigned int opId) {
-    AuthorizationSession* authzSession = AuthorizationSession::get(client);
-
-    auto lockAndOpCtx = findOperationContext(client->getServiceContext(), opId);
-    if (lockAndOpCtx) {
-        OperationContext* opToKill = std::get<1>(*lockAndOpCtx);
-        if (authzSession->isAuthorizedForActionsOnResource(ResourcePattern::forClusterResource(),
-                                                           ActionType::killop) ||
-            authzSession->isCoauthorizedWithClient(opToKill->getClient(),
-                                                   std::get<0>(*lockAndOpCtx))) {
-            return lockAndOpCtx;
-        }
-    }
-
-    return boost::none;
-}
-
-void KillOpCmdBase::killLocalOperation(OperationContext* opCtx, unsigned int opToKill) {
-    stdx::unique_lock<Client> lk;
-    OperationContext* opCtxToKill;
-    auto lockAndOpCtx = findOpForKilling(opCtx->getClient(), opToKill);
-    if (!lockAndOpCtx) {
-        // killOp always reports success past the auth check.
-        return;
-    }
-
-    std::tie(lk, opCtxToKill) = std::move(*lockAndOpCtx);
-
-    invariant(lk);
-    opCtx->getServiceContext()->killOperation(lk, opCtxToKill);
-}
-
 unsigned int KillOpCmdBase::parseOpId(const BSONObj& cmdObj) {
     long long op;
     uassertStatusOK(bsonExtractIntegerField(cmdObj, "op", &op));
diff --git a/src/mongo/db/commands/kill_op_cmd_base.h b/src/mongo/db/commands/kill_op_cmd_base.h
index 57073480d7d..a68def842d4 100644
--- a/src/mongo/db/commands/kill_op_cmd_base.h
+++ b/src/mongo/db/commands/kill_op_cmd_base.h
@@ -60,25 +60,9 @@ public:
 
 protected:
     /**
-     * Given an operation ID, search for an OperationContext with that ID. Returns either
-     * boost::none if no operation with the given ID exists, or the OperationContext along with the
-     * (acquired) lock for the associated Client.
-     */
-    static boost::optional<std::tuple<stdx::unique_lock<Client>, OperationContext*>>
-    findOperationContext(ServiceContext* serviceContext, unsigned int opId);
-
-    /**
-     * Find the given operation, and check if we're authorized to kill it. If the operation is
-     * found, and we're allowed to kill it, this returns the OperationContext as well as the
-     * acquired lock for the associated Client. Otherwise boost::none is returned.
-     */
-    static boost::optional<std::tuple<stdx::unique_lock<Client>, OperationContext*>>
-    findOpForKilling(Client* client, unsigned int opId);
-
-    /**
      * Kill an operation running on this instance of mongod or mongos.
      */
-    static void killLocalOperation(OperationContext* opCtx, unsigned int opToKill);
+    static void killLocalOperation(OperationContext* opCtx, OperationId opToKill);
 
     /**
      * Extract the "op" field from 'cmdObj' and convert the value to unsigned int. Since BSON only
diff --git a/src/mongo/db/curop.cpp b/src/mongo/db/curop.cpp
index e6927ed392d..2148d07cfc0 100644
--- a/src/mongo/db/curop.cpp
+++ b/src/mongo/db/curop.cpp
@@ -293,6 +293,11 @@ void CurOp::reportCurrentOpForClient(OperationContext* opCtx,
 
     if (clientOpCtx) {
         infoBuilder->append("opid", static_cast<int>(clientOpCtx->getOpID()));
+
+        if (auto opKey = clientOpCtx->getOperationKey()) {
+            opKey->appendToBuilder(infoBuilder, "operationKey");
+        }
+
         if (clientOpCtx->isKillPending()) {
             infoBuilder->append("killPending", true);
         }
diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index 0a281d3b3ba..282612c54c8 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/operation_context.h"
 
 #include "mongo/db/client.h"
+#include "mongo/db/operation_key_manager.h"
 #include "mongo/db/service_context.h"
 #include "mongo/platform/mutex.h"
 #include "mongo/platform/random.h"
@@ -79,13 +80,17 @@ const auto kNoWaiterThread = stdx::thread::id();
 
 }  // namespace
 
-OperationContext::OperationContext(Client* client, unsigned int opId)
+OperationContext::OperationContext(Client* client, OperationId opId)
     : _client(client),
       _opId(opId),
       _elapsedTime(client ? client->getServiceContext()->getTickSource()
                           : SystemTickSource::get()) {}
 
-OperationContext::~OperationContext() = default;
+OperationContext::~OperationContext() {
+    if (_opKey) {
+        OperationKeyManager::get(_client).remove(*_opKey);
+    }
+}
 
 void OperationContext::setDeadlineAndMaxTime(Date_t when,
                                              Microseconds maxTime,
@@ -359,6 +364,14 @@ void OperationContext::setLogicalSessionId(LogicalSessionId lsid) {
     _lsid = std::move(lsid);
 }
 
+void OperationContext::setOperationKey(OperationKey opKey) {
+    // Only set the opKey once
+    invariant(!_opKey);
+
+    _opKey.emplace(std::move(opKey));
+    OperationKeyManager::get(_client).add(*_opKey, _opId);
+}
+
 void OperationContext::setTxnNumber(TxnNumber txnNumber) {
     invariant(_lsid);
     _txnNumber = txnNumber;
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index 8fc158a7b81..c1e224bfee7 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -79,7 +79,7 @@ class OperationContext : public Interruptible, public Decorable<OperationContext
     OperationContext& operator=(const OperationContext&) = delete;
 
 public:
-    OperationContext(Client* client, unsigned int opId);
+    OperationContext(Client* client, OperationId opId);
     virtual ~OperationContext();
 
     bool shouldParticipateInFlowControl() const {
@@ -166,11 +166,25 @@ public:
     /**
      * Returns the operation ID associated with this operation.
      */
-    unsigned int getOpID() const {
+    OperationId getOpID() const {
         return _opId;
     }
 
     /**
+     * Returns the operation UUID associated with this operation or boost::none.
+     */
+    const boost::optional<OperationKey>& getOperationKey() const {
+        return _opKey;
+    }
+
+    /**
+     * Sets the operation UUID associated with this operation.
+     *
+     * This function may only be called once per OperationContext.
+     */
+    void setOperationKey(OperationKey opKey);
+
+    /**
      * Returns the session ID associated with this operation, if there is one.
      */
     const boost::optional<LogicalSessionId>& getLogicalSessionId() const {
@@ -490,7 +504,8 @@ private:
 
     Client* const _client;
 
-    const unsigned int _opId;
+    const OperationId _opId;
+    boost::optional<OperationKey> _opKey;
 
     boost::optional<LogicalSessionId> _lsid;
     boost::optional<TxnNumber> _txnNumber;
diff --git a/src/mongo/db/operation_key_manager.cpp b/src/mongo/db/operation_key_manager.cpp
new file mode 100644
index 00000000000..2f545083098
--- /dev/null
+++ b/src/mongo/db/operation_key_manager.cpp
@@ -0,0 +1,89 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kCommand
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/operation_key_manager.h"
+
+#include <fmt/format.h>
+
+#include "mongo/base/error_codes.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/log.h"
+
+namespace mongo {
+namespace {
+auto getOpKeyManager = ServiceContext::declareDecoration<OperationKeyManager>();
+}
+
+OperationKeyManager::~OperationKeyManager() {
+    invariant(_idByOperationKey.empty(),
+              "Every associated OperationContext with an OperationKey must be destroyed before an "
+              "OperationKeyManager can be destroyed.");
+}
+
+OperationKeyManager& OperationKeyManager::get(ServiceContext* serviceContext) {
+    return getOpKeyManager(serviceContext);
+}
+
+void OperationKeyManager::add(const OperationKey& key, OperationId id) {
+    using namespace fmt::literals;
+
+    LOG(2) << "Mapping OperationKey {} to OperationId {}"_format(key.toString(), id);
+
+    stdx::lock_guard lk(_mutex);
+    auto result = _idByOperationKey.emplace(key, id).second;
+
+    uassert(
+        ErrorCodes::BadValue, "OperationKey currently '{}' in use"_format(key.toString()), result);
+}
+
+bool OperationKeyManager::remove(const OperationKey& key) {
+    stdx::lock_guard lk(_mutex);
+    return _idByOperationKey.erase(key);
+}
+
+boost::optional<OperationId> OperationKeyManager::at(const OperationKey& key) const {
+    stdx::lock_guard lk(_mutex);
+    auto it = _idByOperationKey.find(key);
+    if (it == _idByOperationKey.end()) {
+        return boost::none;
+    }
+
+    return it->second;
+}
+
+size_t OperationKeyManager::size() const {
+    stdx::lock_guard lk(_mutex);
+    return _idByOperationKey.size();
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_key_manager.h b/src/mongo/db/operation_key_manager.h
new file mode 100644
index 00000000000..865b618f5ed
--- /dev/null
+++ b/src/mongo/db/operation_key_manager.h
@@ -0,0 +1,87 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+#include "mongo/platform/mutex.h"
+#include "mongo/stdx/unordered_map.h"
+
+namespace mongo {
+
+/**
+ * OperationKeyManager maps OperationKeys to OperationIds
+ *
+ * To make this simpler to reason about, this class is a decoration on a ServiceContext with its own
+ * synchronization. This class is currently expected to be used at 3 points:
+ * - Before Command execution, the OperationKey from the client application is added.
+ * - During a killop-like operation, it is used to find the OperationContext to kill.
+ * - During OperationContext destruction, the OperationKey from the client application is removed.
+ */
+class OperationKeyManager {
+public:
+    static OperationKeyManager& get(ServiceContext* serviceContext = getCurrentServiceContext());
+    static OperationKeyManager& get(Client* client) {
+        return get(client->getServiceContext());
+    }
+    static OperationKeyManager& get(OperationContext* opCtx) {
+        return get(opCtx->getServiceContext());
+    }
+
+    ~OperationKeyManager();
+
+    /**
+     * Add a mapping from OperationKey to OperationId
+     */
+    void add(const OperationKey& key, OperationId id);
+
+    /**
+     * Remove any mapping from OperationKey
+     */
+    bool remove(const OperationKey& key);
+
+    /**
+     * Get the OperationId for OperationKey, or boost::none
+     */
+    boost::optional<OperationId> at(const OperationKey& key) const;
+
+    /**
+     * Get the total count of OperationKeys currently being tracked
+     */
+    size_t size() const;
+
+private:
+    mutable Mutex _mutex = MONGO_MAKE_LATCH("OperationKeyManager::_mutex");
+
+    stdx::unordered_map<OperationKey, OperationId, OperationKey::Hash> _idByOperationKey;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_killer.cpp b/src/mongo/db/operation_killer.cpp
new file mode 100644
index 00000000000..742a9d142e3
--- /dev/null
+++ b/src/mongo/db/operation_killer.cpp
@@ -0,0 +1,99 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kCommand
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/operation_killer.h"
+
+#include "mongo/db/audit.h"
+#include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/client.h"
+#include "mongo/db/operation_key_manager.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/log.h"
+
+namespace mongo {
+
+OperationKiller::OperationKiller(Client* myClient) : _myClient(myClient) {
+    invariant(_myClient);
+}
+
+bool OperationKiller::isGenerallyAuthorizedToKill() const {
+    AuthorizationSession* authzSession = AuthorizationSession::get(_myClient);
+    if (authzSession->isAuthorizedForActionsOnResource(ResourcePattern::forClusterResource(),
+                                                       ActionType::killop)) {
+        return true;
+    }
+
+    return false;
+}
+
+bool OperationKiller::isAuthorizedToKill(const LockedClient& target) const {
+    AuthorizationSession* authzSession = AuthorizationSession::get(_myClient);
+
+    if (target && authzSession->isCoauthorizedWithClient(target.client(), target)) {
+        return true;
+    }
+
+    return false;
+}
+
+void OperationKiller::killOperation(OperationId opId) {
+    auto serviceContext = _myClient->getServiceContext();
+
+    auto target = serviceContext->getLockedClient(opId);
+    if (!target) {
+        // There is no operation for opId
+        return;
+    }
+
+    if (!isGenerallyAuthorizedToKill() && !isAuthorizedToKill(target)) {
+        // The client is not authotized to kill this operation
+        return;
+    }
+
+    serviceContext->killOperation(target, target->getOperationContext());
+
+    log() << "Killed operation: " << opId;
+}
+
+void OperationKiller::killOperation(const OperationKey& opKey) {
+    auto opId = OperationKeyManager::get(_myClient).at(opKey);
+
+    if (!opId) {
+        // There is no operation for opKey
+        return;
+    }
+
+    killOperation(*opId);
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/operation_killer.h b/src/mongo/db/operation_killer.h
new file mode 100644
index 00000000000..d9856a8bcea
--- /dev/null
+++ b/src/mongo/db/operation_killer.h
@@ -0,0 +1,70 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/client.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+
+namespace mongo {
+
+/**
+ * OperationKiller is a utility class to facilitate interrupting other operations
+ *
+ * This class is constructed with a pointer to the Client which is attempting to perform the
+ * interruption. This Client must have a role with the killop action type or own the targeted
+ * OperationContext. The Client must also be a member of the same ServiceContext as the targeted
+ * OperationContext.
+ */
+class OperationKiller {
+public:
+    explicit OperationKiller(Client* myClient);
+
+    /**
+     * Verify that myClient has permission to kill operations it does not own
+     */
+    bool isGenerallyAuthorizedToKill() const;
+
+    /**
+     * Verify that the target exists and the myClient is allowed to kill it
+     */
+    bool isAuthorizedToKill(const LockedClient& target) const;
+
+    /**
+     * Kill an operation running on this instance of mongod or mongos.
+     */
+    void killOperation(OperationId opId);
+    void killOperation(const OperationKey& opKey);
+
+private:
+    Client* const _myClient;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/service_context.cpp b/src/mongo/db/service_context.cpp
index c77d19aaabf..511208ddc94 100644
--- a/src/mongo/db/service_context.cpp
+++ b/src/mongo/db/service_context.cpp
@@ -65,6 +65,8 @@ AtomicWord<int> _numCurrentOps{0};
 
 }  // namespace
 
+LockedClient::LockedClient(Client* client) : _lk{*client}, _client{client} {}
+
 bool hasGlobalServiceContext() {
     return globalServiceContext;
 }
@@ -266,6 +268,12 @@ ServiceContext::UniqueOperationContext ServiceContext::makeOperationContext(Clie
         stdx::lock_guard<Client> lk(*client);
         client->setOperationContext(opCtx.get());
     }
+
+    {
+        stdx::lock_guard lk(_mutex);
+        _clientByOperationId.emplace(opCtx->getOpID(), client);
+    }
+
     return UniqueOperationContext(opCtx.release());
 };
 
@@ -275,6 +283,12 @@ void ServiceContext::OperationContextDeleter::operator()(OperationContext* opCtx
         _numCurrentOps.subtractAndFetch(1);
     }
     auto service = client->getServiceContext();
+
+    {
+        stdx::lock_guard lk(service->_mutex);
+        service->_clientByOperationId.erase(opCtx->getOpID());
+    }
+
     {
         stdx::lock_guard<Client> lk(*client);
         client->resetOperationContext();
@@ -285,6 +299,16 @@ void ServiceContext::OperationContextDeleter::operator()(OperationContext* opCtx
     delete opCtx;
 }
 
+LockedClient ServiceContext::getLockedClient(OperationId id) {
+    stdx::lock_guard lk(_mutex);
+    auto it = _clientByOperationId.find(id);
+    if (it == _clientByOperationId.end()) {
+        return {};
+    }
+
+    return LockedClient(it->second);
+}
+
 void ServiceContext::registerClientObserver(std::unique_ptr<ClientObserver> observer) {
     _clientObservers.emplace_back(std::move(observer));
 }
diff --git a/src/mongo/db/service_context.h b/src/mongo/db/service_context.h
index a50e359e1d2..19e58e82e53 100644
--- a/src/mongo/db/service_context.h
+++ b/src/mongo/db/service_context.h
@@ -50,6 +50,7 @@
 #include "mongo/util/hierarchical_acquisition.h"
 #include "mongo/util/periodic_runner.h"
 #include "mongo/util/tick_source.h"
+#include "mongo/util/uuid.h"
 
 #include <iostream>
 
@@ -91,6 +92,53 @@ protected:
 };
 
 /**
+ * A simple container type to pass around a client and a lock on said client
+ */
+class LockedClient {
+public:
+    LockedClient() = default;
+    explicit LockedClient(Client* client);
+
+    Client* client() const noexcept {
+        return _client;
+    }
+
+    Client* operator->() const noexcept {
+        return client();
+    }
+
+    explicit operator bool() const noexcept {
+        return client();
+    }
+
+    operator WithLock() const noexcept {
+        return WithLock(_lk);
+    }
+
+private:
+    // Technically speaking, _lk holds a Client* and _client is a superfluous variable. That said,
+    // LockedClients will likely be optimized away and the extra variable is a cheap price to pay
+    // for better developer comprehension.
+    stdx::unique_lock<Client> _lk;
+    Client* _client = nullptr;
+};
+
+/**
+ * Every OperationContext is expected to have a unique OperationId within the domain of its
+ * ServiceContext. Generally speaking, OperationId is used for forming maps of OperationContexts and
+ * directing metaoperations like killop.
+ */
+using OperationId = uint32_t;
+
+/**
+ * Users may provide an OperationKey when sending a command request as a stable token by which to
+ * refer to an operation (and thus an OperationContext). An OperationContext is not required to have
+ * an OperationKey. The presence of an OperationKey implies that the client is either closely
+ * tracking or speculative executing its command.
+ */
+using OperationKey = UUID;
+
+/**
  * Class representing the context of a service, such as a MongoD database service or
  * a MongoS routing service.
  *
@@ -519,6 +567,8 @@ public:
         _catalogGeneration.fetchAndAdd(1);
     }
 
+    LockedClient getLockedClient(OperationId id);
+
 private:
     class ClientObserverHolder {
     public:
@@ -541,7 +591,7 @@ private:
         std::unique_ptr<ClientObserver> _observer;
     };
 
-    Mutex _mutex = MONGO_MAKE_LATCH(HierarchicalAcquisitionLevel(2), "ServiceContext::_mutex");
+    Mutex _mutex = MONGO_MAKE_LATCH(/*HierarchicalAcquisitionLevel(2), */ "ServiceContext::_mutex");
 
     /**
      * The periodic runner.
@@ -574,6 +624,8 @@ private:
     std::vector<ClientObserverHolder> _clientObservers;
     ClientSet _clients;
 
+    stdx::unordered_map<OperationId, Client*> _clientByOperationId;
+
     /**
      * The registered OpObserver.
      */
@@ -599,7 +651,7 @@ private:
     std::vector<KillOpListenerInterface*> _killOpListeners;
 
     // Counter for assigning operation ids.
-    AtomicWord<unsigned> _nextOpId{1};
+    AtomicWord<OperationId> _nextOpId{1};
 
     // When the catalog is restarted, the generation goes up by one each time.
     AtomicWord<uint64_t> _catalogGeneration{0};
diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 4570d44bb74..50398eba2e0 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -781,8 +781,6 @@ void execCommandDatabase(OperationContext* opCtx,
             }
         });
 
-        // TODO: move this back to runCommands when mongos supports OperationContext
-        // see SERVER-18515 for details.
         rpc::readRequestMetadata(opCtx, request.body, command->requiresAuth());
         rpc::TrackingMetadata::get(opCtx).initWithOperName(command->getName());
 
diff --git a/src/mongo/rpc/SConscript b/src/mongo/rpc/SConscript
index aa2cf7ea5f2..1dd4f810414 100644
--- a/src/mongo/rpc/SConscript
+++ b/src/mongo/rpc/SConscript
@@ -107,6 +107,7 @@ env.Library(
         '$BUILD_DIR/mongo/base',
         '$BUILD_DIR/mongo/bson/util/bson_extract',
         '$BUILD_DIR/mongo/client/read_preference',
+        '$BUILD_DIR/mongo/db/commands/test_commands_enabled',
         '$BUILD_DIR/mongo/db/logical_time_validator',
         '$BUILD_DIR/mongo/db/repl/optime',
         '$BUILD_DIR/mongo/db/signed_logical_time',
diff --git a/src/mongo/rpc/metadata.cpp b/src/mongo/rpc/metadata.cpp
index e3ed093a693..ab65025a725 100644
--- a/src/mongo/rpc/metadata.cpp
+++ b/src/mongo/rpc/metadata.cpp
@@ -33,6 +33,7 @@
 
 #include "mongo/client/read_preference.h"
 #include "mongo/db/auth/authorization_session.h"
+#include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/dbmessage.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/logical_clock.h"
@@ -59,6 +60,7 @@ void readRequestMetadata(OperationContext* opCtx, const BSONObj& metadataObj, bo
     BSONElement clientElem;
     BSONElement logicalTimeElem;
     BSONElement impersonationElem;
+    BSONElement clientOperationKeyElem;
 
     for (const auto& metadataElem : metadataObj) {
         auto fieldName = metadataElem.fieldNameStringData();
@@ -74,9 +76,21 @@ void readRequestMetadata(OperationContext* opCtx, const BSONObj& metadataObj, bo
             logicalTimeElem = metadataElem;
         } else if (fieldName == kImpersonationMetadataSectionName) {
             impersonationElem = metadataElem;
+        } else if (fieldName == "clientOperationKey"_sd) {
+            clientOperationKeyElem = metadataElem;
         }
     }
 
+    AuthorizationSession* authSession = AuthorizationSession::get(opCtx->getClient());
+
+    if (clientOperationKeyElem &&
+        (getTestCommandsEnabled() ||
+         authSession->isAuthorizedForActionsOnResource(ResourcePattern::forClusterResource(),
+                                                       ActionType::internal))) {
+        auto opKey = uassertStatusOK(UUID::parse(clientOperationKeyElem));
+        opCtx->setOperationKey(std::move(opKey));
+    }
+
     if (readPreferenceElem) {
         ReadPreferenceSetting::get(opCtx) =
             uassertStatusOK(ReadPreferenceSetting::fromInnerBSON(readPreferenceElem));
@@ -103,7 +117,6 @@ void readRequestMetadata(OperationContext* opCtx, const BSONObj& metadataObj, bo
             AuthorizationManager::get(opCtx->getServiceContext())->isAuthEnabled() &&
             (!signedTime.getProof() || *signedTime.getProof() == TimeProofService::TimeProof())) {
 
-            AuthorizationSession* authSession = AuthorizationSession::get(opCtx->getClient());
             // The client is not authenticated and is not using localhost auth bypass.
             if (authSession && !authSession->isAuthenticated() &&
                 !authSession->isUsingLocalhostBypass()) {
diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index 61f692f315a..cca60108d5d 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -396,6 +396,8 @@ void runCommand(OperationContext* opCtx,
 
     boost::optional<RouterOperationContextSession> routerSession;
     try {
+        rpc::readRequestMetadata(opCtx, request.body, command->requiresAuth());
+
         CommandHelpers::evaluateFailCommandFailPoint(opCtx, commandName, invocation->ns());
         bool startTransaction = false;
         if (osi.getAutocommit()) {