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.

1 files changed, 54 insertions, 2 deletions

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};