SERVER-25266 Enhance NetworkASIO to send the client metadata document

6 files changed, 80 insertions, 0 deletions

diff --git a/src/mongo/rpc/metadata.cpp b/src/mongo/rpc/metadata.cpp
index 5c9a2a47a40..d286ece26ca 100644
--- a/src/mongo/rpc/metadata.cpp
+++ b/src/mongo/rpc/metadata.cpp
@@ -33,6 +33,7 @@
 #include "mongo/client/dbclientinterface.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/rpc/metadata/audit_metadata.h"
+#include "mongo/rpc/metadata/client_metadata_ismaster.h"
 #include "mongo/rpc/metadata/config_server_metadata.h"
 #include "mongo/rpc/metadata/server_selection_metadata.h"
 #include "mongo/rpc/metadata/sharding_metadata.h"
@@ -48,6 +49,7 @@ Status readRequestMetadata(OperationContext* txn, const BSONObj& metadataObj) {
     BSONElement ssmElem;
     BSONElement auditElem;
     BSONElement configSvrElem;
+    BSONElement clientElem;
 
     for (const auto& metadataElem : metadataObj) {
         auto fieldName = metadataElem.fieldNameStringData();
@@ -57,6 +59,8 @@ Status readRequestMetadata(OperationContext* txn, const BSONObj& metadataObj) {
             auditElem = metadataElem;
         } else if (fieldName == ConfigServerMetadata::fieldName()) {
             configSvrElem = metadataElem;
+        } else if (fieldName == ClientMetadata::fieldName()) {
+            clientElem = metadataElem;
         }
     }
 
@@ -72,6 +76,12 @@ Status readRequestMetadata(OperationContext* txn, const BSONObj& metadataObj) {
     }
     AuditMetadata::get(txn) = std::move(swAuditMetadata.getValue());
 
+    const auto statusClientMetadata =
+        ClientMetadataIsMasterState::readFromMetadata(txn, clientElem);
+    if (!statusClientMetadata.isOK()) {
+        return statusClientMetadata;
+    }
+
     auto configServerMetadata = ConfigServerMetadata::readFromMetadata(configSvrElem);
     if (!configServerMetadata.isOK()) {
         return configServerMetadata.getStatus();
diff --git a/src/mongo/rpc/metadata/client_metadata.cpp b/src/mongo/rpc/metadata/client_metadata.cpp
index ea03d5e0f84..845a315dd74 100644
--- a/src/mongo/rpc/metadata/client_metadata.cpp
+++ b/src/mongo/rpc/metadata/client_metadata.cpp
@@ -47,6 +47,7 @@
 namespace mongo {
 
 namespace {
+constexpr auto kClientMetadataFieldName = "$client"_sd;
 
 constexpr auto kApplication = "application"_sd;
 constexpr auto kDriver = "driver"_sd;
@@ -397,4 +398,8 @@ void ClientMetadata::logClientMetadata(Client* client) const {
           << client->desc() << ": " << getDocument();
 }
 
+StringData ClientMetadata::fieldName() {
+    return kClientMetadataFieldName;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/rpc/metadata/client_metadata.h b/src/mongo/rpc/metadata/client_metadata.h
index 92646e3a88d..a809484e20d 100644
--- a/src/mongo/rpc/metadata/client_metadata.h
+++ b/src/mongo/rpc/metadata/client_metadata.h
@@ -170,6 +170,11 @@ public:
      */
     void logClientMetadata(Client* client) const;
 
+    /**
+     * Field name for OP_Command metadata that contains client metadata.
+     */
+    static StringData fieldName();
+
 public:
     /**
      * Create a new client metadata document.
diff --git a/src/mongo/rpc/metadata/client_metadata_ismaster.cpp b/src/mongo/rpc/metadata/client_metadata_ismaster.cpp
index 53164ed33d2..2df2461dfa3 100644
--- a/src/mongo/rpc/metadata/client_metadata_ismaster.cpp
+++ b/src/mongo/rpc/metadata/client_metadata_ismaster.cpp
@@ -74,4 +74,41 @@ void ClientMetadataIsMasterState::setClientMetadata(
 }
 
 
+Status ClientMetadataIsMasterState::readFromMetadata(OperationContext* txn, BSONElement& element) {
+    if (element.eoo()) {
+        return Status::OK();
+    }
+
+    auto swParseClientMetadata = ClientMetadata::parse(element);
+
+    if (!swParseClientMetadata.getStatus().isOK()) {
+        return swParseClientMetadata.getStatus();
+    }
+
+    auto& clientMetadataIsMasterState = ClientMetadataIsMasterState::get(txn->getClient());
+
+    clientMetadataIsMasterState.setClientMetadata(txn->getClient(),
+                                                  std::move(swParseClientMetadata.getValue()));
+
+    return Status::OK();
+}
+
+void ClientMetadataIsMasterState::writeToMetadata(OperationContext* txn, BSONObjBuilder* builder) {
+    // We may be asked to write metadata on background threads that are not associated with an
+    // operation context
+    if (!txn) {
+        return;
+    }
+
+    const auto& clientMetadata =
+        ClientMetadataIsMasterState::get(txn->getClient()).getClientMetadata();
+
+    // Skip appending metadata if there is none
+    if (!clientMetadata || clientMetadata.get().getDocument().isEmpty()) {
+        return;
+    }
+
+    builder->append(ClientMetadata::fieldName(), clientMetadata.get().getDocument());
+}
+
 }  // namespace mongo
diff --git a/src/mongo/rpc/metadata/client_metadata_ismaster.h b/src/mongo/rpc/metadata/client_metadata_ismaster.h
index 9bd274ae4f6..eb1f92e5033 100644
--- a/src/mongo/rpc/metadata/client_metadata_ismaster.h
+++ b/src/mongo/rpc/metadata/client_metadata_ismaster.h
@@ -72,6 +72,25 @@ public:
      */
     void setSeenIsMaster();
 
+    /**
+     * Read from the $client section of OP_Command's metadata.
+     *
+     * Returns an error if the $client section is not valid. It is valid for it to not exist though.
+     *
+     * Thread-Safety:
+     *   None - must be only be read and written from the thread owning "Client".
+     */
+    static Status readFromMetadata(OperationContext* txn, BSONElement& elem);
+
+    /**
+     * Write the $client section to OP_Command's metadata if there is a non-empty client metadata
+     * connection with the current client.
+     *
+     * Thread-Safety:
+     *   None - must be only be read and written from the thread owning "Client".
+     */
+    static void writeToMetadata(OperationContext* txn, BSONObjBuilder* builder);
+
 private:
     // Optional client metadata document.
     // Set if client sees isMaster cmd or as part of OP_Command processing.
diff --git a/src/mongo/s/sharding_egress_metadata_hook.cpp b/src/mongo/s/sharding_egress_metadata_hook.cpp
index 49ee7c03d60..25fa7676946 100644
--- a/src/mongo/s/sharding_egress_metadata_hook.cpp
+++ b/src/mongo/s/sharding_egress_metadata_hook.cpp
@@ -35,6 +35,7 @@
 #include "mongo/base/status.h"
 #include "mongo/db/audit.h"
 #include "mongo/db/service_context.h"
+#include "mongo/rpc/metadata/client_metadata_ismaster.h"
 #include "mongo/rpc/metadata/config_server_metadata.h"
 #include "mongo/rpc/metadata/metadata_hook.h"
 #include "mongo/rpc/metadata/repl_set_metadata.h"
@@ -55,6 +56,8 @@ Status ShardingEgressMetadataHook::writeRequestMetadata(bool shardedConnection,
                                                         BSONObjBuilder* metadataBob) {
     try {
         audit::writeImpersonatedUsersToMetadata(txn, metadataBob);
+
+        ClientMetadataIsMasterState::writeToMetadata(txn, metadataBob);
         if (!shardedConnection) {
             return Status::OK();
         }
@@ -70,6 +73,7 @@ Status ShardingEgressMetadataHook::writeRequestMetadata(OperationContext* txn,
                                                         BSONObjBuilder* metadataBob) {
     try {
         audit::writeImpersonatedUsersToMetadata(txn, metadataBob);
+        ClientMetadataIsMasterState::writeToMetadata(txn, metadataBob);
         rpc::ConfigServerMetadata(_getConfigServerOpTime()).writeToMetadata(metadataBob);
         return Status::OK();
     } catch (...) {