SERVER-25266 Enhance NetworkASIO to send the client metadata document

4 files changed, 66 insertions, 0 deletions

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.