SERVER-41362 Attach diagnostic captures to OperationContexts

20 files changed, 154 insertions, 36 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index ea364f80446..e52b3d792ea 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -121,6 +121,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/service_context',
         '$BUILD_DIR/mongo/db/query/command_request_response',
         '$BUILD_DIR/mongo/rpc/client_metadata',
+        '$BUILD_DIR/mongo/util/diagnostic_info',
         '$BUILD_DIR/mongo/util/fail_point',
         '$BUILD_DIR/mongo/util/net/network',
         '$BUILD_DIR/mongo/util/progress_meter',
diff --git a/src/mongo/db/curop.cpp b/src/mongo/db/curop.cpp
index d36b95bbbd9..ce9e60af38d 100644
--- a/src/mongo/db/curop.cpp
+++ b/src/mongo/db/curop.cpp
@@ -228,6 +228,7 @@ CurOp* CurOp::get(const OperationContext& opCtx) {
 void CurOp::reportCurrentOpForClient(OperationContext* opCtx,
                                      Client* client,
                                      bool truncateOps,
+                                     bool backtraceMode,
                                      BSONObjBuilder* infoBuilder) {
     invariant(client);
     OperationContext* clientOpCtx = client->getOperationContext();
@@ -294,6 +295,13 @@ void CurOp::reportCurrentOpForClient(OperationContext* opCtx,
 
         CurOp::get(clientOpCtx)->reportState(infoBuilder, truncateOps);
     }
+    if (backtraceMode) {
+        DiagnosticInfo& diagnostic = DiagnosticInfo::Diagnostic::get(client);
+        BSONObjBuilder waitingForLatchBuilder;
+        waitingForLatchBuilder.append("timestamp", diagnostic.getTimestamp());
+        waitingForLatchBuilder.append("captureName", diagnostic.getCaptureName());
+        infoBuilder->append("waitingForLatch", waitingForLatchBuilder.obj());
+    }
 }
 
 void CurOp::setGenericCursor_inlock(GenericCursor gc) {
diff --git a/src/mongo/db/curop.h b/src/mongo/db/curop.h
index 7444316ad33..37985a76865 100644
--- a/src/mongo/db/curop.h
+++ b/src/mongo/db/curop.h
@@ -36,6 +36,7 @@
 #include "mongo/db/operation_context.h"
 #include "mongo/db/server_options.h"
 #include "mongo/platform/atomic_word.h"
+#include "mongo/util/diagnostic_info.h"
 #include "mongo/util/progress_meter.h"
 #include "mongo/util/time_support.h"
 
@@ -275,6 +276,7 @@ public:
     static void reportCurrentOpForClient(OperationContext* opCtx,
                                          Client* client,
                                          bool truncateOps,
+                                         bool backtraceMode,
                                          BSONObjBuilder* infoBuilder);
 
     /**
diff --git a/src/mongo/db/pipeline/document_source_current_op.cpp b/src/mongo/db/pipeline/document_source_current_op.cpp
index 134bb58145a..ba9c11d9e98 100644
--- a/src/mongo/db/pipeline/document_source_current_op.cpp
+++ b/src/mongo/db/pipeline/document_source_current_op.cpp
@@ -42,6 +42,7 @@ const StringData kIdleSessionsFieldName = "idleSessions"_sd;
 const StringData kLocalOpsFieldName = "localOps"_sd;
 const StringData kTruncateOpsFieldName = "truncateOps"_sd;
 const StringData kIdleCursorsFieldName = "idleCursors"_sd;
+const StringData kBacktraceFieldName = "backtrace"_sd;
 
 const StringData kOpIdFieldName = "opid"_sd;
 const StringData kClientFieldName = "client"_sd;
@@ -116,7 +117,8 @@ DocumentSource::GetNextResult DocumentSourceCurrentOp::getNext() {
                                                              _includeIdleSessions,
                                                              _includeOpsFromAllUsers,
                                                              _truncateOps,
-                                                             _idleCursors);
+                                                             _idleCursors,
+                                                             _backtrace);
 
         _opsIter = _ops.begin();
 
@@ -192,6 +194,7 @@ intrusive_ptr<DocumentSource> DocumentSourceCurrentOp::createFromBson(
     LocalOpsMode showLocalOpsOnMongoS = LocalOpsMode::kRemoteShardOps;
     TruncationMode truncateOps = TruncationMode::kNoTruncation;
     CursorMode idleCursors = CursorMode::kExcludeCursors;
+    BacktraceMode backtrace = BacktraceMode::kExcludeBacktrace;
 
     for (auto&& elem : spec.embeddedObject()) {
         const auto fieldName = elem.fieldNameStringData();
@@ -245,6 +248,14 @@ intrusive_ptr<DocumentSource> DocumentSourceCurrentOp::createFromBson(
                     elem.type() == BSONType::Bool);
             idleCursors =
                 (elem.boolean() ? CursorMode::kIncludeCursors : CursorMode::kExcludeCursors);
+        } else if (fieldName == kBacktraceFieldName) {
+            uassert(ErrorCodes::FailedToParse,
+                    str::stream() << "The 'backtrace' parameter of the $currentOp stage must be "
+                                     "a boolean value, but found: "
+                                  << typeName(elem.type()),
+                    elem.type() == BSONType::Bool);
+            backtrace = (elem.boolean() ? BacktraceMode::kIncludeBacktrace
+                                        : BacktraceMode::kExcludeBacktrace);
         } else {
             uasserted(ErrorCodes::FailedToParse,
                       str::stream() << "Unrecognized option '" << fieldName
@@ -258,7 +269,8 @@ intrusive_ptr<DocumentSource> DocumentSourceCurrentOp::createFromBson(
                                        includeOpsFromAllUsers,
                                        showLocalOpsOnMongoS,
                                        truncateOps,
-                                       idleCursors);
+                                       idleCursors,
+                                       backtrace);
 }
 
 intrusive_ptr<DocumentSourceCurrentOp> DocumentSourceCurrentOp::create(
@@ -268,14 +280,16 @@ intrusive_ptr<DocumentSourceCurrentOp> DocumentSourceCurrentOp::create(
     UserMode includeOpsFromAllUsers,
     LocalOpsMode showLocalOpsOnMongoS,
     TruncationMode truncateOps,
-    CursorMode idleCursors) {
+    CursorMode idleCursors,
+    BacktraceMode backtrace) {
     return new DocumentSourceCurrentOp(pExpCtx,
                                        includeIdleConnections,
                                        includeIdleSessions,
                                        includeOpsFromAllUsers,
                                        showLocalOpsOnMongoS,
                                        truncateOps,
-                                       idleCursors);
+                                       idleCursors,
+                                       backtrace);
 }
 
 Value DocumentSourceCurrentOp::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
@@ -292,6 +306,8 @@ Value DocumentSourceCurrentOp::serialize(boost::optional<ExplainOptions::Verbosi
                   {kTruncateOpsFieldName,
                    _truncateOps == TruncationMode::kTruncateOps ? Value(true) : Value()},
                   {kIdleCursorsFieldName,
-                   _idleCursors == CursorMode::kIncludeCursors ? Value(true) : Value()}}}});
+                   _idleCursors == CursorMode::kIncludeCursors ? Value(true) : Value()},
+                  {kBacktraceFieldName,
+                   _backtrace == BacktraceMode::kIncludeBacktrace ? Value(true) : Value()}}}});
 }
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_current_op.h b/src/mongo/db/pipeline/document_source_current_op.h
index 44055dcb5ad..537f4662fee 100644
--- a/src/mongo/db/pipeline/document_source_current_op.h
+++ b/src/mongo/db/pipeline/document_source_current_op.h
@@ -41,6 +41,7 @@ public:
     using SessionMode = MongoProcessInterface::CurrentOpSessionsMode;
     using UserMode = MongoProcessInterface::CurrentOpUserMode;
     using CursorMode = MongoProcessInterface::CurrentOpCursorMode;
+    using BacktraceMode = MongoProcessInterface::CurrentOpBacktraceMode;
 
     static constexpr StringData kStageName = "$currentOp"_sd;
 
@@ -99,7 +100,8 @@ public:
         UserMode includeOpsFromAllUsers = UserMode::kExcludeOthers,
         LocalOpsMode showLocalOpsOnMongoS = LocalOpsMode::kRemoteShardOps,
         TruncationMode truncateOps = TruncationMode::kNoTruncation,
-        CursorMode idleCursors = CursorMode::kExcludeCursors);
+        CursorMode idleCursors = CursorMode::kExcludeCursors,
+        BacktraceMode backtrace = BacktraceMode::kExcludeBacktrace);
 
     GetNextResult getNext() final;
 
@@ -137,14 +139,16 @@ private:
                             UserMode includeOpsFromAllUsers,
                             LocalOpsMode showLocalOpsOnMongoS,
                             TruncationMode truncateOps,
-                            CursorMode idleCursors)
+                            CursorMode idleCursors,
+                            BacktraceMode backtrace)
         : DocumentSource(pExpCtx),
           _includeIdleConnections(includeIdleConnections),
           _includeIdleSessions(includeIdleSessions),
           _includeOpsFromAllUsers(includeOpsFromAllUsers),
           _showLocalOpsOnMongoS(showLocalOpsOnMongoS),
           _truncateOps(truncateOps),
-          _idleCursors(idleCursors) {}
+          _idleCursors(idleCursors),
+          _backtrace(backtrace) {}
 
     ConnMode _includeIdleConnections = ConnMode::kExcludeIdle;
     SessionMode _includeIdleSessions = SessionMode::kIncludeIdle;
@@ -152,6 +156,7 @@ private:
     LocalOpsMode _showLocalOpsOnMongoS = LocalOpsMode::kRemoteShardOps;
     TruncationMode _truncateOps = TruncationMode::kNoTruncation;
     CursorMode _idleCursors = CursorMode::kExcludeCursors;
+    BacktraceMode _backtrace = BacktraceMode::kExcludeBacktrace;
 
     std::string _shardName;
 
diff --git a/src/mongo/db/pipeline/document_source_current_op_test.cpp b/src/mongo/db/pipeline/document_source_current_op_test.cpp
index 7f72327a51c..f9cd2a91b14 100644
--- a/src/mongo/db/pipeline/document_source_current_op_test.cpp
+++ b/src/mongo/db/pipeline/document_source_current_op_test.cpp
@@ -71,7 +71,8 @@ public:
                                        CurrentOpSessionsMode sessionMode,
                                        CurrentOpUserMode userMode,
                                        CurrentOpTruncateMode truncateMode,
-                                       CurrentOpCursorMode cursorMode) const {
+                                       CurrentOpCursorMode cursorMode,
+                                       CurrentOpBacktraceMode backtraceMode) const {
         return _ops;
     }
 
diff --git a/src/mongo/db/pipeline/mongo_process_common.cpp b/src/mongo/db/pipeline/mongo_process_common.cpp
index a2527aded86..5cbd1501280 100644
--- a/src/mongo/db/pipeline/mongo_process_common.cpp
+++ b/src/mongo/db/pipeline/mongo_process_common.cpp
@@ -51,7 +51,8 @@ std::vector<BSONObj> MongoProcessCommon::getCurrentOps(
     CurrentOpSessionsMode sessionMode,
     CurrentOpUserMode userMode,
     CurrentOpTruncateMode truncateMode,
-    CurrentOpCursorMode cursorMode) const {
+    CurrentOpCursorMode cursorMode,
+    CurrentOpBacktraceMode backtraceMode) const {
     OperationContext* opCtx = expCtx->opCtx;
     AuthorizationSession* ctxAuth = AuthorizationSession::get(opCtx->getClient());
 
@@ -76,7 +77,7 @@ std::vector<BSONObj> MongoProcessCommon::getCurrentOps(
         }
 
         // Delegate to the mongoD- or mongoS-specific implementation of _reportCurrentOpForClient.
-        ops.emplace_back(_reportCurrentOpForClient(opCtx, client, truncateMode));
+        ops.emplace_back(_reportCurrentOpForClient(opCtx, client, truncateMode, backtraceMode));
     }
 
     // If 'cursorMode' is set to include idle cursors, retrieve them and add them to ops.
diff --git a/src/mongo/db/pipeline/mongo_process_common.h b/src/mongo/db/pipeline/mongo_process_common.h
index cef5cb8c7db..3ba69ba20a2 100644
--- a/src/mongo/db/pipeline/mongo_process_common.h
+++ b/src/mongo/db/pipeline/mongo_process_common.h
@@ -57,7 +57,8 @@ public:
                                        CurrentOpSessionsMode sessionMode,
                                        CurrentOpUserMode userMode,
                                        CurrentOpTruncateMode truncateMode,
-                                       CurrentOpCursorMode cursorMode) const final;
+                                       CurrentOpCursorMode cursorMode,
+                                       CurrentOpBacktraceMode backtraceMode) const final;
 
     virtual std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
         OperationContext*, const NamespaceString&) const override;
@@ -81,7 +82,8 @@ protected:
      */
     virtual BSONObj _reportCurrentOpForClient(OperationContext* opCtx,
                                               Client* client,
-                                              CurrentOpTruncateMode truncateOps) const = 0;
+                                              CurrentOpTruncateMode truncateOps,
+                                              CurrentOpBacktraceMode backtraceMode) const = 0;
 
     /**
      * Iterates through all entries in the local SessionCatalog, and adds an entry to the 'ops'
diff --git a/src/mongo/db/pipeline/mongo_process_interface.h b/src/mongo/db/pipeline/mongo_process_interface.h
index 4e4165d29f3..503a4065cb6 100644
--- a/src/mongo/db/pipeline/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/mongo_process_interface.h
@@ -91,6 +91,7 @@ public:
     enum class CurrentOpLocalOpsMode { kLocalMongosOps, kRemoteShardOps };
     enum class CurrentOpSessionsMode { kIncludeIdle, kExcludeIdle };
     enum class CurrentOpCursorMode { kIncludeCursors, kExcludeCursors };
+    enum class CurrentOpBacktraceMode { kIncludeBacktrace, kExcludeBacktrace };
 
     /**
      * Factory function to create MongoProcessInterface of the right type. The implementation will
@@ -283,7 +284,8 @@ public:
         CurrentOpSessionsMode sessionMode,
         CurrentOpUserMode userMode,
         CurrentOpTruncateMode,
-        CurrentOpCursorMode) const = 0;
+        CurrentOpCursorMode,
+        CurrentOpBacktraceMode) const = 0;
 
     /**
      * Returns the name of the local shard if sharding is enabled, or an empty string.
diff --git a/src/mongo/db/pipeline/mongos_process_interface.cpp b/src/mongo/db/pipeline/mongos_process_interface.cpp
index 331b02bcbcf..d91233c52db 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.cpp
+++ b/src/mongo/db/pipeline/mongos_process_interface.cpp
@@ -243,11 +243,15 @@ boost::optional<Document> MongoSInterface::lookupSingleDocument(
 
 BSONObj MongoSInterface::_reportCurrentOpForClient(OperationContext* opCtx,
                                                    Client* client,
-                                                   CurrentOpTruncateMode truncateOps) const {
+                                                   CurrentOpTruncateMode truncateOps,
+                                                   CurrentOpBacktraceMode backtraceMode) const {
     BSONObjBuilder builder;
 
-    CurOp::reportCurrentOpForClient(
-        opCtx, client, (truncateOps == CurrentOpTruncateMode::kTruncateOps), &builder);
+    CurOp::reportCurrentOpForClient(opCtx,
+                                    client,
+                                    (truncateOps == CurrentOpTruncateMode::kTruncateOps),
+                                    (backtraceMode == CurrentOpBacktraceMode::kIncludeBacktrace),
+                                    &builder);
 
     return builder.obj();
 }
diff --git a/src/mongo/db/pipeline/mongos_process_interface.h b/src/mongo/db/pipeline/mongos_process_interface.h
index fcfe82d6321..eaf418a1153 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/mongos_process_interface.h
@@ -242,7 +242,8 @@ public:
 protected:
     BSONObj _reportCurrentOpForClient(OperationContext* opCtx,
                                       Client* client,
-                                      CurrentOpTruncateMode truncateOps) const final;
+                                      CurrentOpTruncateMode truncateOps,
+                                      CurrentOpBacktraceMode backtraceMode) const final;
 
     void _reportCurrentOpsForIdleSessions(OperationContext* opCtx,
                                           CurrentOpUserMode userMode,
diff --git a/src/mongo/db/pipeline/process_interface_standalone.cpp b/src/mongo/db/pipeline/process_interface_standalone.cpp
index 6e13e969ac4..b81ee5a435d 100644
--- a/src/mongo/db/pipeline/process_interface_standalone.cpp
+++ b/src/mongo/db/pipeline/process_interface_standalone.cpp
@@ -572,11 +572,17 @@ bool MongoInterfaceStandalone::fieldsHaveSupportingUniqueIndex(
 }
 
 BSONObj MongoInterfaceStandalone::_reportCurrentOpForClient(
-    OperationContext* opCtx, Client* client, CurrentOpTruncateMode truncateOps) const {
+    OperationContext* opCtx,
+    Client* client,
+    CurrentOpTruncateMode truncateOps,
+    CurrentOpBacktraceMode backtraceMode) const {
     BSONObjBuilder builder;
 
-    CurOp::reportCurrentOpForClient(
-        opCtx, client, (truncateOps == CurrentOpTruncateMode::kTruncateOps), &builder);
+    CurOp::reportCurrentOpForClient(opCtx,
+                                    client,
+                                    (truncateOps == CurrentOpTruncateMode::kTruncateOps),
+                                    (backtraceMode == CurrentOpBacktraceMode::kIncludeBacktrace),
+                                    &builder);
 
     OperationContext* clientOpCtx = client->getOperationContext();
 
diff --git a/src/mongo/db/pipeline/process_interface_standalone.h b/src/mongo/db/pipeline/process_interface_standalone.h
index 8d7cf4693c4..392af0beb5e 100644
--- a/src/mongo/db/pipeline/process_interface_standalone.h
+++ b/src/mongo/db/pipeline/process_interface_standalone.h
@@ -156,7 +156,8 @@ public:
 protected:
     BSONObj _reportCurrentOpForClient(OperationContext* opCtx,
                                       Client* client,
-                                      CurrentOpTruncateMode truncateOps) const final;
+                                      CurrentOpTruncateMode truncateOps,
+                                      CurrentOpBacktraceMode backtraceMode) const final;
 
     void _reportCurrentOpsForIdleSessions(OperationContext* opCtx,
                                           CurrentOpUserMode userMode,
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index 82950fdc427..147be45fc46 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -152,7 +152,8 @@ public:
                                        CurrentOpSessionsMode sessionMode,
                                        CurrentOpUserMode userMode,
                                        CurrentOpTruncateMode truncateMode,
-                                       CurrentOpCursorMode cursorMode) const override {
+                                       CurrentOpCursorMode cursorMode,
+                                       CurrentOpBacktraceMode backtraceMode) const override {
         MONGO_UNREACHABLE;
     }
 
diff --git a/src/mongo/platform/SConscript b/src/mongo/platform/SConscript
index 171a183fa45..6ccbf7f77ce 100644
--- a/src/mongo/platform/SConscript
+++ b/src/mongo/platform/SConscript
@@ -17,7 +17,4 @@ env.CppUnitTest(
         'decimal128_bson_test.cpp',
         'overflow_arithmetic_test.cpp'
     ],
-    LIBDEPS=[
-        '$BUILD_DIR/mongo/db/service_context'
-    ]
 )
\ No newline at end of file
diff --git a/src/mongo/platform/mutex.cpp b/src/mongo/platform/mutex.cpp
index c81783f4bd7..b5661cf6f53 100644
--- a/src/mongo/platform/mutex.cpp
+++ b/src/mongo/platform/mutex.cpp
@@ -31,16 +31,35 @@
 
 namespace mongo {
 
+namespace {
+std::unique_ptr<LockActions> gLockActions;
+}
+
 void Mutex::lock() {
-    auto hasLock = _mutex.try_lock_for(_lockTimeout.toSystemDuration());
+    auto hasLock = _mutex.try_lock_for(kContendedLockTimeout.toSystemDuration());
+    if (hasLock) {
+        return;
+    }
+    if (gLockActions) {
+        gLockActions->onContendedLock(_name);
+    }
+    hasLock = _mutex.try_lock_for(_lockTimeout.toSystemDuration() -
+                                  kContendedLockTimeout.toSystemDuration());
     uassert(
         ErrorCodes::InternalError, "Unable to take latch, wait time exceeds set timeout", hasLock);
 }
 void Mutex::unlock() {
+    if (gLockActions) {
+        gLockActions->onUnlock();
+    }
     _mutex.unlock();
 }
 bool Mutex::try_lock() {
     return _mutex.try_lock();
 }
 
+void Mutex::setLockActions(std::unique_ptr<LockActions> actions) {
+    gLockActions = std::move(actions);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/platform/mutex.h b/src/mongo/platform/mutex.h
index a0be7ae971d..bb885d0cbe1 100644
--- a/src/mongo/platform/mutex.h
+++ b/src/mongo/platform/mutex.h
@@ -36,6 +36,13 @@
 
 namespace mongo {
 
+class LockActions {
+public:
+    virtual ~LockActions() = default;
+    virtual void onContendedLock(const StringData& name) = 0;
+    virtual void onUnlock() = 0;
+};
+
 class Mutex {
 public:
     Mutex() : Mutex("AnonymousMutex"_sd) {}
@@ -48,9 +55,12 @@ public:
         return _name;
     }
 
+    static void setLockActions(std::unique_ptr<LockActions> actions);
+
 private:
     const StringData _name;
     const Seconds _lockTimeout = Seconds(60);
+    static constexpr Milliseconds kContendedLockTimeout = Milliseconds(100);
     stdx::timed_mutex _mutex;
 };
 
diff --git a/src/mongo/platform/mutex_test.cpp b/src/mongo/platform/mutex_test.cpp
index 895914c529b..6c6674aa197 100644
--- a/src/mongo/platform/mutex_test.cpp
+++ b/src/mongo/platform/mutex_test.cpp
@@ -29,15 +29,10 @@
 
 #include "mongo/unittest/unittest.h"
 
-#include "mongo/db/service_context.h"
 #include "mongo/platform/mutex.h"
-#include "mongo/stdx/thread.h"
 
 namespace mongo {
 TEST(MongoMutexTest, BasicSingleThread) {
-    auto serviceContext = ServiceContext::make();
-    setGlobalServiceContext(std::move(serviceContext));
-
     Mutex m;
     m.lock();
     ASSERT(!m.try_lock());
diff --git a/src/mongo/util/diagnostic_info.cpp b/src/mongo/util/diagnostic_info.cpp
index e74b0968233..97718dfe22c 100644
--- a/src/mongo/util/diagnostic_info.cpp
+++ b/src/mongo/util/diagnostic_info.cpp
@@ -27,7 +27,10 @@
  *    it in the license file.
  */
 
+#include "mongo/base/init.h"
+#include "mongo/db/client.h"
 #include "mongo/platform/basic.h"
+#include "mongo/platform/mutex.h"
 
 #include "mongo/util/diagnostic_info.h"
 
@@ -35,7 +38,45 @@
 
 namespace mongo {
 
+namespace {
+const auto gDiagnosticHandle = Client::declareDecoration<DiagnosticInfo::Diagnostic>();
+
+MONGO_INITIALIZER(LockActions)(InitializerContext* context) {
+    class LockActionsSubclass : public LockActions {
+        void onContendedLock(const StringData& name) override {
+            if (haveClient()) {
+                DiagnosticInfo::Diagnostic::set(
+                    Client::getCurrent(),
+                    std::make_shared<DiagnosticInfo>(takeDiagnosticInfo(name)));
+            }
+        }
+        void onUnlock() override {
+            if (haveClient()) {
+                DiagnosticInfo::Diagnostic::set(Client::getCurrent(), nullptr);
+            }
+        }
+    };
+    std::unique_ptr<LockActions> myPointer = std::make_unique<LockActionsSubclass>();
+    Mutex::setLockActions(std::move(myPointer));
+
+    return Status::OK();
+}
+}  // namespace
+
+auto DiagnosticInfo::Diagnostic::get(Client* const client) -> DiagnosticInfo& {
+    auto& handle = gDiagnosticHandle(client);
+    stdx::lock_guard lk(handle.m);
+    return *handle.diagnostic;
+}
+
+void DiagnosticInfo::Diagnostic::set(Client* const client,
+                                     std::shared_ptr<DiagnosticInfo> newDiagnostic) {
+    auto& handle = gDiagnosticHandle(client);
+    stdx::lock_guard lk(handle.m);
+    handle.diagnostic = newDiagnostic;
+}
+
 DiagnosticInfo takeDiagnosticInfo(const StringData& captureName) {
     return DiagnosticInfo(getGlobalServiceContext()->getFastClockSource()->now(), captureName);
 }
-}  // namespace mongo
+}  // namespace mongo
\ No newline at end of file
diff --git a/src/mongo/util/diagnostic_info.h b/src/mongo/util/diagnostic_info.h
index 9685135b0d2..8eba872bb20 100644
--- a/src/mongo/util/diagnostic_info.h
+++ b/src/mongo/util/diagnostic_info.h
@@ -31,6 +31,7 @@
 
 #include "mongo/base/string_data.h"
 #include "mongo/db/service_context.h"
+#include "mongo/platform/mutex.h"
 #include "mongo/util/time_support.h"
 
 namespace mongo {
@@ -41,6 +42,13 @@ namespace mongo {
  */
 class DiagnosticInfo {
 public:
+    struct Diagnostic {
+        static DiagnosticInfo& get(Client*);
+        static void set(Client*, std::shared_ptr<DiagnosticInfo>);
+        Mutex m;
+        std::shared_ptr<DiagnosticInfo> diagnostic;
+    };
+
     virtual ~DiagnosticInfo() = default;
     DiagnosticInfo(const DiagnosticInfo&) = delete;
     DiagnosticInfo& operator=(const DiagnosticInfo&) = delete;
@@ -61,7 +69,6 @@ private:
     Date_t _timestamp;
     StringData _captureName;
 
-
     DiagnosticInfo(const Date_t& timestamp, const StringData& captureName)
         : _timestamp(timestamp), _captureName(captureName) {}
 };
@@ -70,6 +77,4 @@ private:
  * Captures the diagnostic information based on the caller's context.
  */
 DiagnosticInfo takeDiagnosticInfo(const StringData& captureName);
-
-
 }  // namespace monogo