SERVER-23298 refactor ShardingConnectionHook and ShardingEgressMetadataHook to improve link graph

1 files changed, 11 insertions, 69 deletions

diff --git a/src/mongo/s/client/sharding_connection_hook.cpp b/src/mongo/s/client/sharding_connection_hook.cpp
index d3cbd242c04..60ee98b5557 100644
--- a/src/mongo/s/client/sharding_connection_hook.cpp
+++ b/src/mongo/s/client/sharding_connection_hook.cpp
@@ -41,11 +41,6 @@
 #include "mongo/db/auth/internal_user_auth.h"
 #include "mongo/db/client.h"
 #include "mongo/rpc/get_status_from_command_result.h"
-#include "mongo/rpc/metadata/audit_metadata.h"
-#include "mongo/rpc/metadata/config_server_metadata.h"
-#include "mongo/s/client/shard_registry.h"
-#include "mongo/s/cluster_last_error_info.h"
-#include "mongo/s/grid.h"
 #include "mongo/s/version_manager.h"
 #include "mongo/util/log.h"
 
@@ -53,65 +48,9 @@ namespace mongo {
 
 using std::string;
 
-namespace {
-
-// A hook that parses the reply metadata from every response to a command sent from a DBClient
-// created by mongos or a sharding aware mongod and being used for sharded operations.
-// Used by mongos to capture the GLE stats so that we can target the correct node when subsequent
-// getLastError calls are made, as well as by both mongod and mongos to update the stored config
-// server optime.
-Status _shardingReplyMetadataReader(const BSONObj& metadataObj, StringData hostString) {
-    saveGLEStats(metadataObj, hostString);
-
-    auto shard = grid.shardRegistry()->getShardNoReload(hostString.toString());
-    if (!shard) {
-        return Status::OK();
-    }
-    // If this host is a known shard of ours, look for a config server optime in the response
-    // metadata to use to update our notion of the current config server optime.
-    auto responseStatus = rpc::ConfigServerMetadata::readFromMetadata(metadataObj);
-    if (!responseStatus.isOK()) {
-        return responseStatus.getStatus();
-    }
-    auto opTime = responseStatus.getValue().getOpTime();
-    if (opTime.is_initialized()) {
-        grid.shardRegistry()->advanceConfigOpTime(opTime.get());
-    }
-    return Status::OK();
-}
-
-// A hook that will append impersonated users to the metadata of every runCommand run by a DBClient
-// created by mongos or a sharding aware mongod.  mongos uses this information to send information
-// to mongod so that the mongod can produce auditing records attributed to the proper authenticated
-// user(s).
-// Additionally, if the connection is sharding-aware, also appends the stored config server optime.
-Status _shardingRequestMetadataWriter(bool shardedConn,
-                                      BSONObjBuilder* metadataBob,
-                                      StringData hostStringData) {
-    audit::writeImpersonatedUsersToMetadata(metadataBob);
-    if (!shardedConn) {
-        return Status::OK();
-    }
-
-    // Add config server optime to metadata sent to shards.
-    std::string hostString = hostStringData.toString();
-    auto shard = grid.shardRegistry()->getShardNoReload(hostString);
-    if (!shard) {
-        return Status(ErrorCodes::ShardNotFound,
-                      str::stream() << "Shard not found for server: " << hostString);
-    }
-    if (shard->isConfig()) {
-        return Status::OK();
-    }
-    rpc::ConfigServerMetadata(grid.shardRegistry()->getConfigOpTime()).writeToMetadata(metadataBob);
-
-    return Status::OK();
-}
-
-}  // namespace
-
-ShardingConnectionHook::ShardingConnectionHook(bool shardedConnections)
-    : _shardedConnections(shardedConnections) {}
+ShardingConnectionHook::ShardingConnectionHook(
+    bool shardedConnections, std::unique_ptr<rpc::ShardingEgressMetadataHook> egressHook)
+    : _shardedConnections(shardedConnections), _egressHook(std::move(egressHook)) {}
 
 void ShardingConnectionHook::onCreate(DBClientBase* conn) {
     if (conn->type() == ConnectionString::INVALID) {
@@ -131,14 +70,17 @@ void ShardingConnectionHook::onCreate(DBClientBase* conn) {
                 result);
     }
 
+    // Delegate the metadata hook logic to the egress hook; use lambdas to pass the arguments in
+    // the order expected by the egress hook.
     if (_shardedConnections) {
-        conn->setReplyMetadataReader(_shardingReplyMetadataReader);
+        conn->setReplyMetadataReader([this](const BSONObj& metadataObj, StringData target) {
+            return _egressHook->readReplyMetadata(target, metadataObj);
+        });
     }
+    conn->setRequestMetadataWriter([this](BSONObjBuilder* metadataBob, StringData hostStringData) {
+        return _egressHook->writeRequestMetadata(_shardedConnections, hostStringData, metadataBob);
+    });
 
-    conn->setRequestMetadataWriter(
-        [this](BSONObjBuilder* metadataBob, StringData hostStringData) -> Status {
-            return _shardingRequestMetadataWriter(_shardedConnections, metadataBob, hostStringData);
-        });
 
     if (conn->type() == ConnectionString::MASTER) {
         BSONObj isMasterResponse;