SERVER-18236 make runCommandHook and postRunCommand hook operate on the metadata object

8 files changed, 155 insertions, 124 deletions

diff --git a/src/mongo/client/dbclient.cpp b/src/mongo/client/dbclient.cpp
index 07459d0b02a..eca99e1614a 100644
--- a/src/mongo/client/dbclient.cpp
+++ b/src/mongo/client/dbclient.cpp
@@ -244,14 +244,6 @@ namespace {
         return QueryOptions(0);
     }
 
-    void DBClientWithCommands::setRunCommandHook(RunCommandHookFunc func) {
-        _runCommandHook = func;
-    }
-
-    void DBClientWithCommands::setPostRunCommandHook(PostRunCommandHookFunc func) {
-        _postRunCommandHook = func;
-    }
-
     rpc::ProtocolSet DBClientWithCommands::getClientRPCProtocols() const {
         return _clientRPCProtocols;
     }
@@ -268,23 +260,21 @@ namespace {
         _serverRPCProtocols = std::move(protocols);
     }
 
-    bool DBClientWithCommands::runCommand(const string& dbname,
-                                          const BSONObj& cmd,
-                                          BSONObj &info,
-                                          int options) {
+    void DBClientWithCommands::setRequestMetadataWriter(rpc::RequestMetadataWriter writer) {
+        _metadataWriter = std::move(writer);
+    }
 
-        uassert(ErrorCodes::InvalidNamespace, str::stream() << "Database name '" << dbname
-                                                            << "' is not valid.",
-                NamespaceString::validDBName(dbname));
+    const rpc::RequestMetadataWriter& DBClientWithCommands::getRequestMetadataWriter() {
+        return _metadataWriter;
+    }
 
-        BSONObj maybeInterposedCommand = cmd;
+    void DBClientWithCommands::setReplyMetadataReader(rpc::ReplyMetadataReader reader) {
+        _metadataReader = std::move(reader);
+    }
 
-        if (_runCommandHook) {
-            BSONObjBuilder hookInterposedBob;
-            hookInterposedBob.appendElements(cmd);
-            _runCommandHook(&hookInterposedBob);
-            maybeInterposedCommand = hookInterposedBob.obj();
-        }
+    const rpc::ReplyMetadataReader& DBClientWithCommands::getReplyMetadataReader() {
+        return _metadataReader;
+    }
 
     rpc::UniqueReply DBClientWithCommands::runCommandWithMetadata(StringData database,
                                                                   StringData command,
@@ -293,10 +283,9 @@ namespace {
         BSONObjBuilder metadataBob;
         metadataBob.appendElements(metadata);
 
-        // upconvert command and metadata to new format
-        // right now this only handles slaveOk
-        BSONObj upconvertedCmd;
-        BSONObj upconvertedMetadata;
+        if (_metadataWriter) {
+            uassertStatusOK(_metadataWriter(&metadataBob));
+        }
 
         uassert(ErrorCodes::InvalidNamespace, str::stream() << "Database name '" << database
                                                             << "' is not valid.",
@@ -334,10 +323,8 @@ namespace {
                                            commandReply->getCommandReply());
         }
 
-        info = std::move(commandReply->getCommandReply().getOwned());
-
-        if (_postRunCommandHook) {
-            _postRunCommandHook(info, host);
+        if (_metadataReader) {
+            uassertStatusOK(_metadataReader(commandReply->getMetadata(), host));
         }
 
         return rpc::UniqueReply(std::move(replyMsg), std::move(commandReply));
@@ -511,25 +498,25 @@ namespace {
     }
 
     namespace {
-        class RunCommandHookOverrideGuard {
-            MONGO_DISALLOW_COPYING(RunCommandHookOverrideGuard);
+        class ScopedMetadataWriterRemover {
+            MONGO_DISALLOW_COPYING(ScopedMetadataWriterRemover);
         public:
-            RunCommandHookOverrideGuard(DBClientWithCommands* cli,
-                                        const DBClientWithCommands::RunCommandHookFunc& hookFunc)
-                : _cli(cli), _oldHookFunc(cli->getRunCommandHook()) {
-                cli->setRunCommandHook(hookFunc);
+            ScopedMetadataWriterRemover(DBClientWithCommands* cli)
+                : _cli(cli), _oldWriter(cli->getRequestMetadataWriter()) {
+                _cli->setRequestMetadataWriter(rpc::RequestMetadataWriter{});
             }
-            ~RunCommandHookOverrideGuard() {
-                _cli->setRunCommandHook(_oldHookFunc);
+            ~ScopedMetadataWriterRemover() {
+                _cli->setRequestMetadataWriter(_oldWriter);
             }
         private:
             DBClientWithCommands* const _cli;
-            DBClientWithCommands::RunCommandHookFunc const _oldHookFunc;
+            rpc::RequestMetadataWriter _oldWriter;
         };
     }  // namespace
 
     void DBClientWithCommands::_auth(const BSONObj& params) {
-        RunCommandHookOverrideGuard hookGuard(this, RunCommandHookFunc());
+        ScopedMetadataWriterRemover{this};
+
         std::string mechanism;
 
         uassertStatusOK(bsonExtractStringField(params,
diff --git a/src/mongo/client/dbclient_rs.cpp b/src/mongo/client/dbclient_rs.cpp
index 980e2f4fea6..2d482d4ae63 100644
--- a/src/mongo/client/dbclient_rs.cpp
+++ b/src/mongo/client/dbclient_rs.cpp
@@ -175,27 +175,26 @@ namespace {
         return _master->getServerHostAndPort();
     }
 
-    void DBClientReplicaSet::setRunCommandHook(DBClientWithCommands::RunCommandHookFunc func) {
+    void DBClientReplicaSet::setRequestMetadataWriter(rpc::RequestMetadataWriter writer) {
         // Set the hooks in both our sub-connections and in ourselves.
         if (_master) {
-            _master->setRunCommandHook(func);
+            _master->setRequestMetadataWriter(writer);
         }
         if (_lastSlaveOkConn.get()) {
-           _lastSlaveOkConn->setRunCommandHook(func);
+            _lastSlaveOkConn->setRequestMetadataWriter(writer);
         }
-        _runCommandHook = func;
+        DBClientWithCommands::setRequestMetadataWriter(std::move(writer));
     }
 
-    void DBClientReplicaSet::setPostRunCommandHook
-    (DBClientWithCommands::PostRunCommandHookFunc func) {
+    void DBClientReplicaSet::setReplyMetadataReader(rpc::ReplyMetadataReader reader) {
         // Set the hooks in both our sub-connections and in ourselves.
         if (_master) {
-            _master->setPostRunCommandHook(func);
+            _master->setReplyMetadataReader(reader);
         }
         if (_lastSlaveOkConn.get()) {
-           _lastSlaveOkConn->setPostRunCommandHook(func);
+            _lastSlaveOkConn->setReplyMetadataReader(reader);
         }
-        _postRunCommandHook = func;
+        DBClientWithCommands::setReplyMetadataReader(std::move(reader));
     }
 
     // A replica set connection is never disconnected, since it controls its own reconnection
@@ -319,8 +318,8 @@ namespace {
         _masterHost = h;
         _master.reset(newConn);
         _master->setParentReplSetName(_setName);
-        _master->setRunCommandHook(_runCommandHook);
-        _master->setPostRunCommandHook(_postRunCommandHook);
+        _master->setRequestMetadataWriter(getRequestMetadataWriter());
+        _master->setReplyMetadataReader(getReplyMetadataReader());
 
         _auth( _master.get() );
         return _master.get();
@@ -723,8 +722,8 @@ namespace {
 
         _lastSlaveOkConn.reset(newConn);
         _lastSlaveOkConn->setParentReplSetName(_setName);
-        _lastSlaveOkConn->setRunCommandHook(_runCommandHook);
-        _lastSlaveOkConn->setPostRunCommandHook(_postRunCommandHook);
+        _lastSlaveOkConn->setRequestMetadataWriter(getRequestMetadataWriter());
+        _lastSlaveOkConn->setReplyMetadataReader(getReplyMetadataReader());
 
         if (_authPooledSecondaryConn) {
             _auth(_lastSlaveOkConn.get());
diff --git a/src/mongo/client/dbclient_rs.h b/src/mongo/client/dbclient_rs.h
index e88320a99fa..fb29559c80a 100644
--- a/src/mongo/client/dbclient_rs.h
+++ b/src/mongo/client/dbclient_rs.h
@@ -161,6 +161,9 @@ namespace mongo {
                                                 const BSONObj& metadata,
                                                 const BSONObj& commandArgs) final;
 
+        void setRequestMetadataWriter(rpc::RequestMetadataWriter writer) final;
+
+        void setReplyMetadataReader(rpc::ReplyMetadataReader reader) final;
         // ---- low level ------
 
         virtual bool call( Message &toSend, Message &response, bool assertOk=true , std::string * actualServer = 0 );
@@ -180,9 +183,6 @@ namespace mongo {
                                       const BSONObj& queryObj,
                                       int queryOptions );
 
-        virtual void setRunCommandHook(DBClientWithCommands::RunCommandHookFunc func);
-        virtual void setPostRunCommandHook(DBClientWithCommands::PostRunCommandHookFunc func);
-
         /**
          * Performs a "soft reset" by clearing all states relating to secondary nodes and
          * returning secondary connections to the pool.
diff --git a/src/mongo/client/dbclientcursor.cpp b/src/mongo/client/dbclientcursor.cpp
index 60b3c676b74..f30cd35b963 100644
--- a/src/mongo/client/dbclientcursor.cpp
+++ b/src/mongo/client/dbclientcursor.cpp
@@ -169,15 +169,6 @@ namespace {
     
     void DBClientCursor::initLazy( bool isRetry ) {
         massert( 15875 , "DBClientCursor::initLazy called on a client that doesn't support lazy" , _client->lazySupported() );
-        if (DBClientWithCommands::RunCommandHookFunc hook = _client->getRunCommandHook()) {
-            if (NamespaceString(ns).isCommand()) {
-                BSONObjBuilder bob;
-                bob.appendElements(query);
-                hook(&bob);
-                query = bob.obj();
-            }
-        }
-        
         Message toSend;
         _assembleInit( toSend );
         _client->say( toSend, isRetry, &_originalHost );
@@ -203,13 +194,6 @@ namespace {
 
         dataReceived( retry, _lazyHost );
 
-        if (DBClientWithCommands::PostRunCommandHookFunc hook = _client->getPostRunCommandHook()) {
-            if (NamespaceString(ns).isCommand()) {
-                BSONObj cmdResponse = peekFirst();
-                hook(cmdResponse, _lazyHost);
-            }
-        }
-
         return ! retry;
     }
 
diff --git a/src/mongo/client/dbclientinterface.h b/src/mongo/client/dbclientinterface.h
index 7ac0ef96646..a41d6cac357 100644
--- a/src/mongo/client/dbclientinterface.h
+++ b/src/mongo/client/dbclientinterface.h
@@ -437,6 +437,33 @@ namespace mongo {
 
         void setClientRPCProtocols(rpc::ProtocolSet clientProtocols);
 
+        /**
+         * Sets a RequestMetadataWriter on this connection.
+         *
+         * TODO: support multiple metadata writers.
+         */
+        virtual void setRequestMetadataWriter(rpc::RequestMetadataWriter writer);
+
+        /**
+         * Gets the RequestMetadataWriter that is set on this connection. This may
+         * be an uninitialized stdx::function, so it should be checked for validity
+         * with operator bool() first.
+         */
+        const rpc::RequestMetadataWriter& getRequestMetadataWriter();
+
+        /**
+         * Sets a ReplyMetadataReader on this connection.
+         *
+         * TODO: support multiple metadata readers.
+         */
+        virtual void setReplyMetadataReader(rpc::ReplyMetadataReader reader);
+
+        /**
+         * Gets the ReplyMetadataReader that is set on this connection. This may
+         * be an uninitialized stdx::function, so it should be checked for validity
+         * with operator bool() first.
+         */
+        const rpc::ReplyMetadataReader& getReplyMetadataReader();
 
         /**
          * Runs a database command. This variant allows the caller to manually specify the metadata
@@ -757,29 +784,6 @@ namespace mongo {
         virtual std::string toString() const = 0;
 
         /**
-         * A function type for runCommand hooking; the function takes a pointer
-         * to a BSONObjBuilder and returns nothing.  The builder contains a
-         * runCommand BSON object.
-         * Once such a function is set as the runCommand hook, every time the DBClient
-         * processes a runCommand, the hook will be called just prior to sending it to the server. 
-         */
-        typedef stdx::function<void(BSONObjBuilder*)> RunCommandHookFunc;
-        virtual void setRunCommandHook(RunCommandHookFunc func);
-        RunCommandHookFunc getRunCommandHook() const {
-            return _runCommandHook;
-        }
-
-        /** 
-         * Similar to above, but for running a function on a command response after a command
-         * has been run.
-         */
-        typedef stdx::function<void(const BSONObj&, const std::string&)> PostRunCommandHookFunc;
-        virtual void setPostRunCommandHook(PostRunCommandHookFunc func);
-        PostRunCommandHookFunc getPostRunCommandHook() const {
-            return _postRunCommandHook;
-        }
-
-        /**
          * Run a pseudo-command such as sys.inprog/currentOp, sys.killop/killOp
          * or sys.unlock/fsyncUnlock
          *
@@ -837,12 +841,6 @@ namespace mongo {
                        const std::string &username,
                        BSONObj *info);
 
-        /**
-         * These functions will be executed by the driver on runCommand calls.
-         */
-        RunCommandHookFunc _runCommandHook;
-        PostRunCommandHookFunc _postRunCommandHook;
-
         // should be set by subclasses during connection.
         void _setServerRPCProtocols(rpc::ProtocolSet serverProtocols);
 
@@ -864,6 +862,9 @@ namespace mongo {
          */
         rpc::ProtocolSet _serverRPCProtocols{rpc::supports::kAll};
 
+        rpc::RequestMetadataWriter _metadataWriter;
+        rpc::ReplyMetadataReader _metadataReader;
+
         enum QueryOptions _cachedAvailableOptions;
         bool _haveCachedAvailableOptions;
     };
diff --git a/src/mongo/client/syncclusterconnection.cpp b/src/mongo/client/syncclusterconnection.cpp
index 792cd530e22..851dd0e9b11 100644
--- a/src/mongo/client/syncclusterconnection.cpp
+++ b/src/mongo/client/syncclusterconnection.cpp
@@ -190,8 +190,8 @@ namespace mongo {
     void SyncClusterConnection::_connect( const std::string& host ) {
         log() << "SyncClusterConnection connecting to [" << host << "]" << endl;
         DBClientConnection * c = new DBClientConnection( true );
-        c->setRunCommandHook(_runCommandHook);
-        c->setPostRunCommandHook(_postRunCommandHook);
+        c->setRequestMetadataWriter(getRequestMetadataWriter());
+        c->setReplyMetadataReader(getReplyMetadataReader());
         c->setSoTimeout( _socketTimeout );
         string errmsg;
         if ( ! c->connect( HostAndPort(host), errmsg ) )
@@ -213,19 +213,45 @@ namespace mongo {
         std::string ns = dbname + ".$cmd";
         BSONObj interposedCmd = cmd;
 
-        if (_runCommandHook) {
-            BSONObjBuilder cmdObjBob;
-            cmdObjBob.appendElements(cmd);
-            _runCommandHook(&cmdObjBob);
-            interposedCmd = cmdObjBob.obj();
+        if (getRequestMetadataWriter()) {
+            // We have a metadata writer. We need to upconvert the metadata, write to it,
+            // Then downconvert it again. This unfortunate, but this code is going to be
+            // removed anyway as part of CSRS.
+
+            BSONObj upconvertedCommand;
+            BSONObj upconvertedMetadata;
+
+            std::tie(upconvertedCommand, upconvertedMetadata) =  uassertStatusOK(
+                rpc::upconvertRequestMetadata(cmd, options)
+            );
+
+            BSONObjBuilder metadataBob;
+            metadataBob.appendElements(upconvertedMetadata);
+
+            uassertStatusOK(getRequestMetadataWriter()(&metadataBob));
+
+            std::tie(interposedCmd, options) = uassertStatusOK(
+                rpc::downconvertRequestMetadata(std::move(upconvertedCommand), metadataBob.done())
+            );
         }
 
-        info = findOne(ns, Query(interposedCmd), 0, options);
+        BSONObj legacyResult = findOne(ns, Query(interposedCmd), 0, options);
+
+        BSONObj upconvertedMetadata;
+        BSONObj upconvertedReply;
 
-        if (_postRunCommandHook) {
-            _postRunCommandHook(info, getServerAddress());
+        std::tie(upconvertedReply, upconvertedMetadata) = uassertStatusOK(
+            rpc::upconvertReplyMetadata(legacyResult)
+        );
+
+        if (getReplyMetadataReader()) {
+            // TODO: what does getServerAddress() actually mean here as this connection
+            // represents a connection to 1 or 3 config servers...
+            uassertStatusOK(getReplyMetadataReader()(upconvertedReply, getServerAddress()));
         }
 
+        info = upconvertedReply;
+
         return isOk(info);
     }
 
@@ -600,24 +626,23 @@ namespace mongo {
             if( _conns[i] ) _conns[i]->setSoTimeout( socketTimeout );
     }
 
-    void SyncClusterConnection::setRunCommandHook(DBClientWithCommands::RunCommandHookFunc func) {
+    void SyncClusterConnection::setRequestMetadataWriter(rpc::RequestMetadataWriter writer) {
         // Set the hooks in both our sub-connections and in ourselves.
         for (size_t i = 0; i < _conns.size(); ++i) {
             if (_conns[i]) { 
-                _conns[i]->setRunCommandHook(func);
+                _conns[i]->setRequestMetadataWriter(writer);
             }
         }
-        _runCommandHook = func;
+        DBClientWithCommands::setRequestMetadataWriter(std::move(writer));
     }
 
-    void SyncClusterConnection::setPostRunCommandHook
-    (DBClientWithCommands::PostRunCommandHookFunc func) {
+    void SyncClusterConnection::setReplyMetadataReader(rpc::ReplyMetadataReader reader) {
         // Set the hooks in both our sub-connections and in ourselves.
         for (size_t i = 0; i < _conns.size(); ++i) {
             if (_conns[i]) { 
-                _conns[i]->setPostRunCommandHook(func);
+                _conns[i]->setReplyMetadataReader(reader);
             }
         }
-        _postRunCommandHook = func;
+        DBClientWithCommands::setReplyMetadataReader(std::move(reader));
     }
 }
diff --git a/src/mongo/client/syncclusterconnection.h b/src/mongo/client/syncclusterconnection.h
index 86ad8421ff1..0147cec1711 100644
--- a/src/mongo/client/syncclusterconnection.h
+++ b/src/mongo/client/syncclusterconnection.h
@@ -126,13 +126,13 @@ namespace mongo {
         // OP_QUERY command. The reason for this is that delicate logic for targeting/locking
         // config servers is in SyncClusterConnection::findOne, and refactoring that logic
         // is both risky and of dubious value as we move to config server replica sets (CSRS).
-        virtual bool runCommand(const std::string& dbname,
-                                const BSONObj& cmd,
-                                BSONObj& info,
-                                int options) final;
+        bool runCommand(const std::string& dbname,
+                        const BSONObj& cmd,
+                        BSONObj& info,
+                        int options) final;
 
-        virtual void setRunCommandHook(DBClientWithCommands::RunCommandHookFunc func);
-        virtual void setPostRunCommandHook(DBClientWithCommands::PostRunCommandHookFunc func);
+        void setRequestMetadataWriter(rpc::RequestMetadataWriter writer) final;
+        void setReplyMetadataReader(rpc::ReplyMetadataReader reader) final;
 
         /**
          * Allow custom query processing through an external (e.g. mongos-only) service.
diff --git a/src/mongo/rpc/metadata.h b/src/mongo/rpc/metadata.h
index ca281066cb9..ea0277059a3 100644
--- a/src/mongo/rpc/metadata.h
+++ b/src/mongo/rpc/metadata.h
@@ -31,6 +31,7 @@
 #include <tuple>
 
 #include "mongo/base/status_with.h"
+#include "mongo/stdx/functional.h"
 
 namespace mongo {
     class BSONObj;
@@ -94,5 +95,39 @@ namespace rpc {
      */
     StatusWith<LegacyCommandAndFlags> downconvertRequestMetadata(BSONObj cmdObj, BSONObj metadata);
 
+    /**
+     * A command reply and associated metadata object.
+     */
+    using CommandReplyWithMetadata = std::tuple<BSONObj, BSONObj>;
+
+    /**
+     * Given a legacy command reply, attempts to strip the metadata from the reply and construct
+     * a metadata object.
+     */
+    StatusWith<CommandReplyWithMetadata> upconvertReplyMetadata(BSONObj legacyReply);
+
+    /**
+     * Given a command reply object and an associated metadata object,
+     * attempts to construct a legacy command object.
+     */
+    StatusWith<BSONObj> downconvertReplyMetadata(BSONObj commandReply, BSONObj replyMetadata);
+
+    /**
+     * A function type for writing request metadata. The function takes a pointer to a
+     * BSONObjBuilder used to construct the metadata object and returns a Status indicating
+     * if the metadata was written successfully.
+     */
+    using RequestMetadataWriter = stdx::function<Status(BSONObjBuilder*)>;
+
+    /**
+     * A function type for reading reply metadata. The function takes a a reference to a
+     * metadata object received in a command reply and a string containing the server address of the
+     * host that executed the command and returns a Status indicating if the
+     * metadata was read successfully.
+     *
+     * TODO: would it be a layering violation if this hook took an OperationContext* ?
+     */
+    using ReplyMetadataReader = stdx::function<Status(const BSONObj&, StringData)>;
+
 }  // namespace rpc
 }  // namespace mongo