SERVER-6754 Create a mock for testing replica set connections

Final step: Modify the replica set client to be able to use custom hooks for creating connections.

19 files changed, 488 insertions, 187 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index da38faaded6..382747a2665 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -453,7 +453,7 @@ mongosLibraryFiles = [
 env.Library( "mongoscore" , mongosLibraryFiles, LIBDEPS=['db/auth/authmongos'] )
 
 env.CppUnitTest( "balancer_policy_test" , [ "s/balancer_policy_tests.cpp" ] ,
-                 LIBDEPS=["mongoscore", "coreshard","mongocommon","coreserver","coredb","dbcmdline","mongodandmongos"] ,
+                 LIBDEPS=["mongoscore", "coreshard", "mongocommon","coreserver","coredb","dbcmdline","mongodandmongos"] ,
                  NO_CRUTCH=True)
 
 env.CppUnitTest("scoped_db_conn_test", [ "client/scoped_db_conn_test.cpp" ],
@@ -603,14 +603,15 @@ if has_option( "sharedclient" ):
 
 # dbtests test binary
 env.StaticLibrary('testframework', ['dbtests/framework.cpp'], LIBDEPS=['unittest/unittest'])
+
 env.StaticLibrary('mocklib', [
+        'dbtests/mock/mock_conn_registry.cpp',
         'dbtests/mock/mock_dbclient_connection.cpp',
+        'dbtests/mock/mock_dbclient_cursor.cpp',
         'dbtests/mock/mock_remote_db_server.cpp',
         'dbtests/mock/mock_replica_set.cpp'
     ],
-    LIBDEPS=['unittest/unittest', 'mongocommon'])
-
-
+    LIBDEPS=['mongocommon'])
 
 test = testEnv.Install(
     '#/',
@@ -644,7 +645,7 @@ if darwin or env["_HAVEPCAP"]:
         sniffEnv.Append( LIBS=[ "wpcap" ] )
 
     sniffEnv.Install( '#/', sniffEnv.Program( "mongosniff", "tools/sniffer.cpp",
-                                              LIBDEPS=["gridfs", "serveronly", "coreserver", "coredb", "notmongodormongos"] ) )
+                                              LIBDEPS=["gridfs", "serveronly", "coreserver", "coredb", "notmongodormongos"]))
 
 # --- shell ---
 
diff --git a/src/mongo/client/dbclient.cpp b/src/mongo/client/dbclient.cpp
index 03373f9f473..98eb72067e5 100644
--- a/src/mongo/client/dbclient.cpp
+++ b/src/mongo/client/dbclient.cpp
@@ -911,6 +911,10 @@ namespace mongo {
         return n;
     }
 
+    void DBClientConnection::setReplSetClientCallback(DBClientReplicaSet* rsClient) {
+        clientSet = rsClient;
+    }
+
     unsigned long long DBClientConnection::query(
             boost::function<void(DBClientCursorBatchIterator &)> f,
             const string& ns,
diff --git a/src/mongo/client/dbclient_rs.cpp b/src/mongo/client/dbclient_rs.cpp
index 18e87f9d2f0..fc7a20adf88 100644
--- a/src/mongo/client/dbclient_rs.cpp
+++ b/src/mongo/client/dbclient_rs.cpp
@@ -260,6 +260,27 @@ namespace mongo {
         return new TagSet(arrayBuilder.arr());
     }
 
+    /**
+     * @return the connection associated with the monitor node. Will also attempt
+     *     to establish connection if NULL. Can still return NULL if reconnect failed.
+     */
+    shared_ptr<DBClientConnection> _getConnWithRefresh(ReplicaSetMonitor::Node& node) {
+        if (node.conn.get() == NULL) {
+            ConnectionString connStr(node.addr);
+            string errmsg;
+
+            try {
+                node.conn.reset(dynamic_cast<DBClientConnection*>(
+                        connStr.connect(errmsg, ReplicaSetMonitor::SOCKET_TIMEOUT_SECS)));
+            }
+            catch (const AssertionException& ex) {
+                node.ok = false;
+            }
+        }
+
+        return node.conn;
+    }
+
     // --------------------------------
     // ----- ReplicaSetMonitor ---------
     // --------------------------------
@@ -315,6 +336,8 @@ namespace mongo {
         return seedStr;
     }
 
+    const double ReplicaSetMonitor::SOCKET_TIMEOUT_SECS = 5;
+
     // Must already be in _setsLock when constructing a new ReplicaSetMonitor. This is why you
     // should only create ReplicaSetMonitors from ReplicaSetMonitor::get and
     // ReplicaSetMonitor::createIfNeeded.
@@ -756,21 +779,38 @@ namespace mongo {
             log() << "trying to add new host " << *i << " to replica set " << this->_name << endl;
 
             // Connect to new node
-            HostAndPort h( *i );
-            DBClientConnection * newConn = new DBClientConnection( true, 0, 5.0 );
+            HostAndPort host(*i);
+            ConnectionString connStr(host);
+
+            uassert(16530, str::stream() << "cannot create a replSet node connection that "
+                    "is not single: " << host.toString(true),
+                    connStr.type() == ConnectionString::MASTER ||
+                    connStr.type() == ConnectionString::CUSTOM);
 
+            DBClientConnection* newConn = NULL;
             string errmsg;
-            try{
-                if( ! newConn->connect( h , errmsg ) ){
-                    throw DBException( errmsg, 15927 );
-                }
-                log() << "successfully connected to new host " << *i << " in replica set " << this->_name << endl;
+            try {
+                // Needs to perform a dynamic_cast because we need to set the replSet
+                // callback. We should eventually not need this after we remove the
+                // callback.
+                newConn = dynamic_cast<DBClientConnection*>(
+                        connStr.connect(errmsg, SOCKET_TIMEOUT_SECS));
+            }
+            catch (const AssertionException& ex) {
+                errmsg = ex.toString();
             }
-            catch( DBException& e ){
-                warning() << "cannot connect to new host " << *i << " to replica set " << this->_name << causedBy( e ) << endl;
+
+            if (!errmsg.empty()) {
+                log() << "successfully connected to new host " << *i
+                        << " in replica set " << this->_name << endl;
+            }
+            else {
+                warning() << "cannot connect to new host " << *i
+                        << " to replica set " << this->_name
+                        << ", err: " << errmsg << endl;
             }
 
-            _nodes.push_back( Node( h , newConn ) );
+            _nodes.push_back(Node(host, newConn));
         }
 
         // Invalidate the cached _master index since the _nodes structure has
@@ -906,7 +946,8 @@ namespace mongo {
                 {
                     scoped_lock lk( _lock );
                     if ( i >= _nodes.size() ) break;
-                    nodeConn = _nodes[i].conn;
+                    nodeConn = _getConnWithRefresh(_nodes[i]);
+                    if (nodeConn.get() == NULL) continue;
                 }
 
                 string maybePrimary;
@@ -942,7 +983,9 @@ namespace mongo {
                         probablePrimaryIdx = _find_inlock( maybePrimary );
 
                         if (probablePrimaryIdx >= 0) {
-                            probablePrimaryConn = _nodes[probablePrimaryIdx].conn;
+                            probablePrimaryConn = _getConnWithRefresh(
+                                    _nodes[probablePrimaryIdx]);
+                            if (probablePrimaryConn.get() == NULL) continue;
                         }
                     }
 
@@ -1035,7 +1078,7 @@ namespace mongo {
             // first see if the current master is fine
             if ( _master >= 0 ) {
                 verify(_master < static_cast<int>(_nodes.size()));
-                masterConn = _nodes[_master].conn;
+                masterConn = _getConnWithRefresh(_nodes[_master]);
             }
         }
 
@@ -1104,9 +1147,10 @@ namespace mongo {
     
     bool ReplicaSetMonitor::_checkConnMatch_inlock( DBClientConnection* conn,
             size_t nodeOffset ) const {
-        
-        return ( nodeOffset < _nodes.size() &&
-            conn->getServerAddress() == _nodes[nodeOffset].conn->getServerAddress() );
+        return (nodeOffset < _nodes.size() &&
+                // Assumption: value for getServerAddress was extracted from
+                // HostAndPort::toString()
+                conn->getServerAddress() == _nodes[nodeOffset].addr.toString());
     }
 
     HostAndPort ReplicaSetMonitor::selectAndCheckNode(ReadPreference preference,
@@ -1247,25 +1291,37 @@ namespace mongo {
             // Don't check servers we have already
             if (_find(*iter) >= 0) continue;
 
-            scoped_ptr<DBClientConnection> conn(new DBClientConnection(true, 0, 5.0));
+            ConnectionString connStr(*iter);
+            scoped_ptr<DBClientConnection> conn;
 
-            try{
-                string errmsg;
-                if (!conn->connect(*iter, errmsg)) {
-                    throw DBException(errmsg, 15928);
-                }
+            uassert(16531, str::stream() << "cannot create a replSet node connection that "
+                    "is not single: " << iter->toString(true),
+                    connStr.type() == ConnectionString::MASTER ||
+                    connStr.type() == ConnectionString::CUSTOM);
 
+            string errmsg;
+            try {
+                // Needs to perform a dynamic_cast because we need to set the replSet
+                // callback. We should eventually not need this after we remove the
+                // callback.
+                conn.reset(dynamic_cast<DBClientConnection*>(
+                        connStr.connect(errmsg, SOCKET_TIMEOUT_SECS)));
+            }
+            catch (const AssertionException& ex) {
+                errmsg = ex.toString();
+            }
+
+            if (conn.get() != NULL && errmsg.empty()) {
                 log() << "successfully connected to seed " << *iter
                         << " for replica set " << _name << endl;
+
+                string maybePrimary;
+                _checkConnection(conn.get(), maybePrimary, false, -1);
             }
-            catch(const DBException& e){
-                log() << "error connecting to seed " << *iter << causedBy(e) << endl;
-                // skip seeds that don't work
-                continue;
+            else {
+                log() << "error connecting to seed " << *iter
+                      << ", err: " << errmsg << endl;
             }
-
-            string maybePrimary;
-            _checkConnection(conn.get(), maybePrimary, false, -1);
         }
 
         // Check everything to get the first data
@@ -1417,12 +1473,33 @@ namespace mongo {
         }
 
         _masterHost = monitor->getMaster();
-        _master.reset( new DBClientConnection( true , this , _so_timeout ) );
+
+        ConnectionString connStr(_masterHost);
+
         string errmsg;
-        if ( ! _master->connect( _masterHost , errmsg ) ) {
-            monitor->notifyFailure( _masterHost );
-            uasserted( 13639 , str::stream() << "can't connect to new replica set master [" << _masterHost.toString() << "] err: " << errmsg );
+        DBClientConnection* newConn;
+
+        try {
+            // Needs to perform a dynamic_cast because we need to set the replSet
+            // callback. We should eventually not need this after we remove the
+            // callback.
+            newConn = dynamic_cast<DBClientConnection*>(
+                    connStr.connect(errmsg, _so_timeout));
+        }
+        catch (const AssertionException& ex) {
+            errmsg = ex.toString();
         }
+
+        if (newConn == NULL || !errmsg.empty()) {
+            monitor->notifyFailure(_masterHost);
+            uasserted(13639, str::stream() << "can't connect to new replica set master ["
+                      << _masterHost.toString() << "]"
+                      << (errmsg.empty()? "" : ", err: ") << errmsg);
+        }
+
+        _master.reset(newConn);
+        _master->setReplSetClientCallback(this);
+
         _auth( _master.get() );
         return _master.get();
     }
@@ -1684,8 +1761,22 @@ namespace mongo {
             return _master.get();
         }
 
-        _lastSlaveOkConn.reset(new DBClientConnection(true , this , _so_timeout));
-        _lastSlaveOkConn->connect(_lastSlaveOkHost);
+        string errmsg;
+        ConnectionString connStr(_lastSlaveOkHost);
+        // Needs to perform a dynamic_cast because we need to set the replSet
+        // callback. We should eventually not need this after we remove the
+        // callback.
+        DBClientConnection* newConn = dynamic_cast<DBClientConnection*>(
+                connStr.connect(errmsg, _so_timeout));
+
+	// Assert here instead of returning NULL since the contract of this method is such
+	// that returning NULL means none of the nodes were good, which is not the case here.
+        uassert(16532, str::stream() << "Failed to connect to "
+                << _lastSlaveOkHost.toString(true),
+                newConn != NULL);
+
+        _lastSlaveOkConn.reset(newConn);
+        _lastSlaveOkConn->setReplSetClientCallback(this);
 
         _auth(_lastSlaveOkConn.get());
         return _lastSlaveOkConn.get();
diff --git a/src/mongo/client/dbclient_rs.h b/src/mongo/client/dbclient_rs.h
index 61a0e824800..2d2c0de7b36 100644
--- a/src/mongo/client/dbclient_rs.h
+++ b/src/mongo/client/dbclient_rs.h
@@ -43,7 +43,6 @@ namespace mongo {
      */
     class ReplicaSetMonitor {
     public:
-
         typedef boost::function1<void,const ReplicaSetMonitor*> ConfigChangeHook;
 
         /**
@@ -133,6 +132,8 @@ namespace mongo {
 
         };
 
+        static const double SOCKET_TIMEOUT_SECS;
+
         /**
          * Selects the right node given the nodes to pick from and the preference.
          *
diff --git a/src/mongo/client/dbclientinterface.h b/src/mongo/client/dbclientinterface.h
index 70dda6c958b..b561805c00c 100644
--- a/src/mongo/client/dbclientinterface.h
+++ b/src/mongo/client/dbclientinterface.h
@@ -172,6 +172,7 @@ namespace mongo {
     bool hasReadPreference(const BSONObj& queryObj);
 
     class DBClientBase;
+    class DBClientConnection;
 
     /**
      * ConnectionString handles parsing different ways to connect to mongo and determining method
@@ -289,6 +290,11 @@ namespace mongo {
             _connectHook = hook;
         }
 
+        static ConnectionHook* getConnectionHook() {
+            scoped_lock lk( _connectHookMutex );
+            return _connectHook;
+        }
+
     private:
 
         void _fillServers( string s );
@@ -1179,6 +1185,18 @@ namespace mongo {
             return _numConnections;
         }
 
+        /**
+         * Primarily used for notifying the replica set client that the server
+         * it is talking to is not primary anymore.
+         *
+         * @param rsClient caller is responsible for managing the life of rsClient
+         * and making sure that it lives longer than this object.
+         *
+         * Warning: This is only for internal use and will eventually be removed in
+         * the future.
+         */
+        void setReplSetClientCallback(DBClientReplicaSet* rsClient);
+
         static void setLazyKillCursor( bool lazy ) { _lazyKillCursor = lazy; }
         static bool getLazyKillCursor() { return _lazyKillCursor; }
 
diff --git a/src/mongo/dbtests/mock/mock_conn_registry.cpp b/src/mongo/dbtests/mock/mock_conn_registry.cpp
new file mode 100644
index 00000000000..7a1970102b5
--- /dev/null
+++ b/src/mongo/dbtests/mock/mock_conn_registry.cpp
@@ -0,0 +1,97 @@
+/*
+ *    Copyright (C) 2012 10gen Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "mongo/dbtests/mock/mock_conn_registry.h"
+
+#include "mongo/base/init.h"
+#include "mongo/dbtests/mock/mock_dbclient_connection.h"
+
+namespace mongo {
+    boost::scoped_ptr<MockConnRegistry> MockConnRegistry::_instance;
+
+    MONGO_INITIALIZER(MockConnRegistry)(InitializerContext* context) {
+        return MockConnRegistry::init();
+    }
+
+    Status MockConnRegistry::init() {
+        MockConnRegistry::_instance.reset(new MockConnRegistry());
+        return Status::OK();
+    }
+
+    MockConnRegistry::MockConnRegistry():
+            _mockConnStrHook(this),
+            _registryMutex("mockConnRegistryMutex") {
+    }
+
+    MockConnRegistry* MockConnRegistry::get() {
+        return _instance.get();
+    }
+
+    ConnectionString::ConnectionHook* MockConnRegistry::getConnStrHook() {
+        return &_mockConnStrHook;
+    }
+
+    void MockConnRegistry::addServer(MockRemoteDBServer* server) {
+        scoped_lock sl(_registryMutex);
+
+        const std::string hostName(server->getServerAddress());
+        fassert(16533, _registry.count(hostName) == 0);
+
+        _registry[hostName] = server;
+    }
+
+    bool MockConnRegistry::removeServer(const std::string& hostName) {
+        scoped_lock sl(_registryMutex);
+        return _registry.erase(hostName) == 1;
+    }
+
+    void MockConnRegistry::clear() {
+        scoped_lock sl(_registryMutex);
+        _registry.clear();
+    }
+
+    MockDBClientConnection* MockConnRegistry::connect(const std::string& connStr) {
+        scoped_lock sl(_registryMutex);
+        fassert(16534, _registry.count(connStr) == 1);
+        return new MockDBClientConnection(_registry[connStr]);
+    }
+
+    MockConnRegistry::MockConnHook::MockConnHook(MockConnRegistry* registry):
+            _registry(registry) {
+    }
+
+    MockConnRegistry::MockConnHook::~MockConnHook() {
+    }
+
+    mongo::DBClientBase* MockConnRegistry::MockConnHook::connect(
+                const ConnectionString& connString,
+                std::string& errmsg,
+                double socketTimeout) {
+        const string hostName(connString.toString());
+        MockDBClientConnection* conn = _registry->connect(hostName);
+
+        if (!conn->connect(hostName.c_str(), errmsg)) {
+            // Assumption: connect never throws, so no leak.
+            delete conn;
+
+            // mimic ConnectionString::connect for MASTER type connection to return NULL
+            // if the destination is unreachable.
+            return NULL;
+        }
+
+        return conn;
+    }
+}
diff --git a/src/mongo/dbtests/mock/mock_conn_registry.h b/src/mongo/dbtests/mock/mock_conn_registry.h
new file mode 100644
index 00000000000..fa4bfdb0e7e
--- /dev/null
+++ b/src/mongo/dbtests/mock/mock_conn_registry.h
@@ -0,0 +1,108 @@
+/*
+ *    Copyright (C) 2012 10gen Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#pragma once
+
+#include "mongo/base/status.h"
+#include "mongo/client/dbclientinterface.h"
+#include "mongo/dbtests/mock/mock_dbclient_connection.h"
+#include "mongo/dbtests/mock/mock_remote_db_server.h"
+#include "mongo/platform/unordered_map.h"
+#include "mongo/util/concurrency/mutex.h"
+
+namespace mongo {
+    /**
+     * Registry for storing mock servers and can create mock connections to these
+     * servers.
+     */
+    class MockConnRegistry {
+    public:
+        /**
+         * Initializes the static instance.
+         */
+        static Status init();
+
+        /**
+         * @return the singleton registry. If this needs to be called before main(),
+         *     then the initializer method should depend on "MockConnRegistry".
+         */
+        static MockConnRegistry* get();
+
+        /**
+         * Adds a server to this registry.
+         *
+         * @param server the server to add. Caller is responsible for disposing
+         *     the server.
+         */
+        void addServer(MockRemoteDBServer* server);
+
+        /**
+         * Removes the server from this registry.
+         *
+         * @param hostName the host name of the server to remove.
+         *
+         * @return true if the server is in the registry and was removed.
+         */
+        bool removeServer(const std::string& hostName);
+
+        /**
+         * Clears the registry.
+         */
+        void clear();
+
+        /**
+         * @return a new mocked connection to a server with the given hostName.
+         */
+        MockDBClientConnection* connect(const std::string& hostName);
+
+        /**
+         * @return the hook that can be used with ConnectionString.
+         */
+        ConnectionString::ConnectionHook* getConnStrHook();
+
+    private:
+        class MockConnHook: public ConnectionString::ConnectionHook {
+        public:
+            /**
+             * Creates a new connection hook for the ConnectionString class that
+             * can create mock connections to mock replica set members using their
+             * pseudo host names.
+             *
+             * @param replSet the mock replica set. Caller is responsible for managing
+             *     replSet and making sure that it lives longer than this object.
+             */
+            MockConnHook(MockConnRegistry* registry);
+            ~MockConnHook();
+
+            mongo::DBClientBase* connect(
+                    const mongo::ConnectionString& connString,
+                    std::string& errmsg, double socketTimeout);
+
+        private:
+            MockConnRegistry* _registry;
+        };
+
+        MockConnRegistry();
+
+        static boost::scoped_ptr<MockConnRegistry> _instance;
+
+        MockConnHook _mockConnStrHook;
+
+        // protects _registry
+        mongo::mutex _registryMutex;
+        unordered_map<std::string, MockRemoteDBServer*> _registry;
+    };
+}
diff --git a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
index a8d76723ed9..4817348eea1 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
+++ b/src/mongo/dbtests/mock/mock_dbclient_connection.cpp
@@ -24,7 +24,7 @@ using mongo::BSONObj;
 using std::string;
 using std::vector;
 
-namespace mongo_test {
+namespace mongo {
     MockDBClientConnection::MockDBClientConnection(MockRemoteDBServer* remoteServer,
             bool autoReconnect):
             _remoteServerInstanceID(remoteServer->getInstanceID()),
@@ -37,6 +37,16 @@ namespace mongo_test {
     MockDBClientConnection::~MockDBClientConnection() {
     }
 
+    bool MockDBClientConnection::connect(const char* hostName, std::string& errmsg) {
+        if (_remoteServer->isRunning()) {
+            _remoteServerInstanceID = _remoteServer->getInstanceID();
+            return true;
+        }
+
+        errmsg.assign("cannot connect to " + _remoteServer->getServerAddress());
+        return false;
+    }
+
     bool MockDBClientConnection::runCommand(const string& dbname, const BSONObj& cmdObj,
             BSONObj &info, int options, const mongo::AuthenticationTable* auth) {
         checkConnection();
diff --git a/src/mongo/dbtests/mock/mock_dbclient_connection.h b/src/mongo/dbtests/mock/mock_dbclient_connection.h
index 4fed28b36ca..6a534bf4192 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_connection.h
+++ b/src/mongo/dbtests/mock/mock_dbclient_connection.h
@@ -23,7 +23,7 @@
 #include "mongo/client/dbclientinterface.h"
 #include "mongo/dbtests/mock/mock_remote_db_server.h"
 
-namespace mongo_test {
+namespace mongo {
     /**
      * A simple class for mocking mongo::DBClientConnection.
      *
@@ -48,6 +48,12 @@ namespace mongo_test {
         // DBClientBase methods
         //
 
+        bool connect(const char* hostName, std::string& errmsg);
+
+        inline bool connect(const HostAndPort& host, std::string& errmsg) {
+            return connect(host.toString().c_str(), errmsg);
+        }
+
         bool runCommand(const std::string& dbname, const mongo::BSONObj& cmdObj,
                 mongo::BSONObj &info, int options = 0,
                 const mongo::AuthenticationTable* auth = NULL);
diff --git a/src/mongo/dbtests/mock/mock_dbclient_cursor.cpp b/src/mongo/dbtests/mock/mock_dbclient_cursor.cpp
index a18f63eba75..64705c5d7c4 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_cursor.cpp
+++ b/src/mongo/dbtests/mock/mock_dbclient_cursor.cpp
@@ -17,7 +17,7 @@
 
 #include "mongo/dbtests/mock/mock_dbclient_cursor.h"
 
-namespace mongo_test {
+namespace mongo {
     MockDBClientCursor::MockDBClientCursor(mongo::DBClientBase* client,
             const mongo::BSONArray& resultSet):
         mongo::DBClientCursor(client, "", 0, 0, 0) {
diff --git a/src/mongo/dbtests/mock/mock_dbclient_cursor.h b/src/mongo/dbtests/mock/mock_dbclient_cursor.h
index 3779592a014..463c0be016b 100644
--- a/src/mongo/dbtests/mock/mock_dbclient_cursor.h
+++ b/src/mongo/dbtests/mock/mock_dbclient_cursor.h
@@ -20,7 +20,7 @@
 #include "mongo/client/dbclientcursor.h"
 #include "mongo/client/dbclientmockcursor.h"
 
-namespace mongo_test {
+namespace mongo {
 
     /**
      * Simple adapter class for mongo::DBClientMockCursor to mongo::DBClientCursor.
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.cpp b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
index 8f4fb0953c8..0887b323e7a 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.cpp
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.cpp
@@ -27,7 +27,7 @@ using mongo::str::stream;
 using std::string;
 using std::vector;
 
-namespace mongo_test {
+namespace mongo {
     MockRemoteDBServer::CircularBSONIterator::CircularBSONIterator(
             const vector<BSONObj>& replyVector) {
         for (std::vector<mongo::BSONObj>::const_iterator iter = replyVector.begin();
@@ -57,16 +57,12 @@ namespace mongo_test {
             _delayMilliSec(0),
             _cmdCount(0),
             _queryCount(0),
-            _connStrHook(this) {
+            _instanceID(0) {
     }
 
     MockRemoteDBServer::~MockRemoteDBServer() {
     }
 
-    mongo::ConnectionString::ConnectionHook* MockRemoteDBServer::getConnectionHook() {
-        return &_connStrHook;
-    }
-
     void MockRemoteDBServer::setDelay(long long milliSec) {
         scoped_spinlock sLock(_lock);
         _delayMilliSec = milliSec;
@@ -209,25 +205,4 @@ namespace mongo_test {
             throw mongo::SocketException(mongo::SocketException::CLOSED, _hostName);
         }
     }
-
-    MockRemoteDBServer::MockDBClientConnStrHook::MockDBClientConnStrHook(
-            MockRemoteDBServer* mockServer): _mockServer(mockServer) {
-    }
-
-    MockRemoteDBServer::MockDBClientConnStrHook::~MockDBClientConnStrHook() {
-    }
-
-    mongo::DBClientBase* MockRemoteDBServer::MockDBClientConnStrHook::connect(
-            const mongo::ConnectionString& connString,
-            std::string& errmsg,
-            double socketTimeout) {
-        if (_mockServer->isRunning()) {
-            return new MockDBClientConnection(_mockServer);
-        }
-        else {
-            // mimic ConnectionString::connect for MASTER type connection to return NULL
-            // if the destination is unreachable.
-            return NULL;
-        }
-    }
 }
diff --git a/src/mongo/dbtests/mock/mock_remote_db_server.h b/src/mongo/dbtests/mock/mock_remote_db_server.h
index 6c050da8648..0b2823ca580 100644
--- a/src/mongo/dbtests/mock/mock_remote_db_server.h
+++ b/src/mongo/dbtests/mock/mock_remote_db_server.h
@@ -23,7 +23,7 @@
 #include "mongo/client/dbclientinterface.h"
 #include "mongo/util/concurrency/spin_lock.h"
 
-namespace mongo_test {
+namespace mongo {
     /**
      * A very simple mock that acts like a database server. Every object keeps track of its own
      * InstanceID, which initially starts at zero and increments every time it is restarted.
@@ -37,28 +37,26 @@ namespace mongo_test {
         typedef size_t InstanceID;
 
         /**
-         * Creates a new mock server. This also setups a hook to this server that can be used
-         * to allow clients using the ConnectionString::connect interface to create connections
-         * to this server. The requirements to make this hook fully functional are:
+         * Creates a new mock server. This can also be setup to work with the
+         * ConnectionString class by using mongo::MockConnRegistry as follows:
          *
-         * 1. hostName of this server should start with $.
-         * 2. No other instance has the same hostName as this.
+         * ConnectionString::setConnectionHook(MockConnRegistry::get()->getConnStrHook());
+         * MockRemoteDBServer server("$a");
+         * MockConnRegistry::get()->addServer(&server);
          *
-         * To register the hook, simply call setConnectionHook:
+         * This allows clients using the ConnectionString::connect interface to create
+         * connections to this server. The requirements to make this hook fully functional are:
          *
-         * MockRemoteDBServer mockServer;
-         * ConnectionString::setConnectionHook(mockServer.getConnectionHook());
+         * 1. hostName of this server should start with $.
+         * 2. No other instance has the same hostName as this.
          *
          * @param hostName the host name for this server.
+         *
+         * @see MockConnRegistry
          */
         MockRemoteDBServer(const std::string& hostName);
         virtual ~MockRemoteDBServer();
 
-        /**
-         * @return the hook that can be used with ConnectionString.
-         */
-        mongo::ConnectionString::ConnectionHook* getConnectionHook();
-
         //
         // Connectivity methods
         //
@@ -171,29 +169,6 @@ namespace mongo_test {
         };
 
         /**
-         * Custom connection hook that can be used with ConnectionString.
-         */
-        class MockDBClientConnStrHook: public mongo::ConnectionString::ConnectionHook {
-        public:
-            /**
-             * Creates a new connection hook for the ConnectionString class that
-             * can create mock connections to this server.
-             *
-             * @param replSet the mock replica set. Caller is responsible for managing
-             *     mockServer and making sure that it lives longer than this object.
-             */
-            MockDBClientConnStrHook(MockRemoteDBServer* mockServer);
-            ~MockDBClientConnStrHook();
-
-            mongo::DBClientBase* connect(
-                    const mongo::ConnectionString& connString,
-                    std::string& errmsg, double socketTimeout);
-
-        private:
-            MockRemoteDBServer* _mockServer;
-        };
-
-        /**
          * Checks whether the instance of the server is still up.
          *
          * @throws mongo::SocketException if this server is down
@@ -204,7 +179,7 @@ namespace mongo_test {
 
         bool _isRunning;
 
-        std::string _hostName;
+        const std::string _hostName;
         long long _delayMilliSec;
 
         //
@@ -225,7 +200,5 @@ namespace mongo_test {
 
         // protects this entire instance
         mutable mongo::SpinLock _lock;
-
-        MockDBClientConnStrHook _connStrHook;
     };
 }
diff --git a/src/mongo/dbtests/mock/mock_replica_set.cpp b/src/mongo/dbtests/mock/mock_replica_set.cpp
index ead132d9b29..331a0bc7331 100644
--- a/src/mongo/dbtests/mock/mock_replica_set.cpp
+++ b/src/mongo/dbtests/mock/mock_replica_set.cpp
@@ -16,7 +16,9 @@
 #include "mongo/dbtests/mock/mock_replica_set.h"
 
 #include "mongo/db/repl/rs_member.h"
+#include "mongo/dbtests/mock/mock_conn_registry.h"
 #include "mongo/dbtests/mock/mock_dbclient_connection.h"
+#include "mongo/util/map_util.h"
 
 #include <sstream>
 
@@ -26,38 +28,9 @@ using std::map;
 using std::string;
 using std::vector;
 
-namespace mongo_test {
-    MockReplicaSet::ReplSetConnHook::ReplSetConnHook(MockReplicaSet* replSet):
-            _replSet(replSet) {
-    }
-
-    MockReplicaSet::ReplSetConnHook::~ReplSetConnHook() {
-    }
-
-    mongo::DBClientBase* MockReplicaSet::ReplSetConnHook::connect(
-            const mongo::ConnectionString& connString,
-            std::string& errmsg,
-            double socketTimeout) {
-        vector<mongo::HostAndPort> servers = connString.getServers();
-        verify(servers.size() == 1);
-
-        const string serverName = servers.front().toString(true);
-
-        MockRemoteDBServer* node = _replSet->getNode(serverName);
-
-        if (node->isRunning()) {
-            // Follow ConnectionString::connect to set autoReconnect to true
-            return new MockDBClientConnection(node, true);
-        }
-
-        // mimic ConnectionString::connect for MASTER type connection to return NULL
-        // if the destination is unreachable
-        return NULL;
-    }
-
+namespace mongo {
     MockReplicaSet::MockReplicaSet(const string& setName, size_t nodes):
-            _setName(setName),
-            _connStringHook(this) {
+            _setName(setName) {
         std::vector<mongo::ReplSetConfig::MemberCfg> replConfig;
 
         for (size_t n = 0; n < nodes; n++) {
@@ -72,7 +45,10 @@ namespace mongo_test {
                 _secondaryHosts.push_back(hostName);
             }
 
-            _nodeMap[hostName] = new MockRemoteDBServer(hostName);
+            MockRemoteDBServer* mockServer = new MockRemoteDBServer(hostName);
+            _nodeMap[hostName] = mockServer;
+
+            MockConnRegistry::get()->addServer(mockServer);
 
             mongo::ReplSetConfig::MemberCfg config;
             config.h = mongo::HostAndPort(hostName);
@@ -85,6 +61,7 @@ namespace mongo_test {
     MockReplicaSet::~MockReplicaSet() {
         for (std::map<string, MockRemoteDBServer*>::iterator iter = _nodeMap.begin();
                 iter != _nodeMap.end(); ++iter) {
+            MockConnRegistry::get()->removeServer(iter->second->getServerAddress());
             delete iter->second;
         }
     }
@@ -130,12 +107,8 @@ namespace mongo_test {
         return _secondaryHosts;
     }
 
-    mongo::ConnectionString::ConnectionHook* MockReplicaSet::getConnectionHook() {
-        return &_connStringHook;
-    }
-
     MockRemoteDBServer* MockReplicaSet::getNode(const string& hostName) {
-        return _nodeMap[hostName];
+        return mapFindWithDefault(_nodeMap, hostName, static_cast<MockRemoteDBServer*>(NULL));
     }
 
     const vector<mongo::ReplSetConfig::MemberCfg>& MockReplicaSet::getReplConfig() const {
diff --git a/src/mongo/dbtests/mock/mock_replica_set.h b/src/mongo/dbtests/mock/mock_replica_set.h
index 8fb98255123..a628c61c747 100644
--- a/src/mongo/dbtests/mock/mock_replica_set.h
+++ b/src/mongo/dbtests/mock/mock_replica_set.h
@@ -22,7 +22,7 @@
 #include <map>
 #include <vector>
 
-namespace mongo_test {
+namespace mongo {
     /**
      * This is a helper class for managing a replica set consisting of
      * MockRemoteDBServer instances.
@@ -49,7 +49,6 @@ namespace mongo_test {
         std::string getSetName() const;
         std::string getConnectionString() const;
         std::vector<mongo::HostAndPort> getHosts() const;
-        mongo::ConnectionString::ConnectionHook* getConnectionHook();
         const std::vector<mongo::ReplSetConfig::MemberCfg>& getReplConfig() const;
         std::string getPrimary() const;
         const std::vector<std::string>& getSecondaries() const;
@@ -89,27 +88,6 @@ namespace mongo_test {
         void restore(const std::string& hostName);
 
     private:
-        class ReplSetConnHook: public mongo::ConnectionString::ConnectionHook {
-        public:
-            /**
-             * Creates a new connection hook for the ConnectionString class that
-             * can create mock connections to mock replica set members using their
-             * pseudo host names.
-             *
-             * @param replSet the mock replica set. Caller is responsible for managing
-             *     replSet and making sure that it lives longer than this object.
-             */
-            ReplSetConnHook(MockReplicaSet* replSet);
-            ~ReplSetConnHook();
-
-            mongo::DBClientBase* connect(
-                    const mongo::ConnectionString& connString,
-                    std::string& errmsg, double socketTimeout);
-
-        private:
-            MockReplicaSet* _replSet;
-        };
-
         /**
          * Mocks the ismaster command based on the information on the current
          * replica set configuration.
@@ -129,7 +107,6 @@ namespace mongo_test {
 
         const std::string _setName;
         std::map<std::string, MockRemoteDBServer*> _nodeMap;
-        ReplSetConnHook _connStringHook;
         std::vector<mongo::ReplSetConfig::MemberCfg> _replConfig;
 
         std::string _primaryHost;
diff --git a/src/mongo/dbtests/mock_dbclient_conn_test.cpp b/src/mongo/dbtests/mock_dbclient_conn_test.cpp
index 68b6f7d9e75..39e80353519 100644
--- a/src/mongo/dbtests/mock_dbclient_conn_test.cpp
+++ b/src/mongo/dbtests/mock_dbclient_conn_test.cpp
@@ -30,6 +30,8 @@
 
 using mongo::BSONObj;
 using mongo::ConnectionString;
+using mongo::MockDBClientConnection;
+using mongo::MockRemoteDBServer;
 
 using std::string;
 using std::vector;
diff --git a/src/mongo/dbtests/mock_replica_set_test.cpp b/src/mongo/dbtests/mock_replica_set_test.cpp
index abbf23ac19a..d82c7ff8d30 100644
--- a/src/mongo/dbtests/mock_replica_set_test.cpp
+++ b/src/mongo/dbtests/mock_replica_set_test.cpp
@@ -24,6 +24,8 @@ using mongo::BSONElement;
 using mongo::BSONObj;
 using mongo::BSONObjIterator;
 using mongo::ConnectionString;
+using mongo::MockRemoteDBServer;
+using mongo::MockReplicaSet;
 
 using std::set;
 using std::string;
diff --git a/src/mongo/dbtests/replica_set_monitor_test.cpp b/src/mongo/dbtests/replica_set_monitor_test.cpp
index d3525574061..db909a5c205 100644
--- a/src/mongo/dbtests/replica_set_monitor_test.cpp
+++ b/src/mongo/dbtests/replica_set_monitor_test.cpp
@@ -19,23 +19,30 @@
  * ReplicaSetMonitor::selectNode and TagSet
  */
 
-#include <vector>
-
+#include "mongo/client/dbclientinterface.h"
 #include "mongo/client/dbclient_rs.h"
+#include "mongo/dbtests/mock/mock_conn_registry.h"
+#include "mongo/dbtests/mock/mock_replica_set.h"
 #include "mongo/unittest/unittest.h"
 
-namespace {
-    using std::vector;
-    using boost::scoped_ptr;
-
-    using mongo::BSONObj;
-    using mongo::BSONArray;
-    using mongo::BSONArrayBuilder;
-    using mongo::ReplicaSetMonitor;
-    using mongo::HostAndPort;
-    using mongo::ReadPreference;
-    using mongo::TagSet;
+#include <vector>
 
+using std::vector;
+using std::string;
+using boost::scoped_ptr;
+
+using mongo::BSONObj;
+using mongo::BSONArray;
+using mongo::BSONArrayBuilder;
+using mongo::ConnectionString;
+using mongo::HostAndPort;
+using mongo::MockReplicaSet;
+using mongo::ReadPreference;
+using mongo::ReplicaSetMonitor;
+using mongo::ReplicaSetMonitorPtr;
+using mongo::TagSet;
+
+namespace mongo_test {
     const BSONObj SampleIsMasterDoc = BSON("tags"
                                             << BSON("dc" << "NYC"
                                                     << "p" << "2"
@@ -1475,4 +1482,54 @@ namespace {
         ASSERT_THROWS(tags.getCurrentTag(), mongo::AssertionException);
 #endif
     }
+
+    // TODO: Port these existing tests here: replmonitor_bad_seed.js, repl_monitor_refresh.js
+
+    /**
+     * Warning: Tests running this fixture cannot be run in parallel with other tests
+     * that uses ConnectionString::setConnectionHook
+     */
+    class ReplicaSetMonitorTest: public mongo::unittest::Test {
+    protected:
+        void setUp() {
+            _replSet.reset(new MockReplicaSet("test", 3));
+            _originalConnectionHook = ConnectionString::getConnectionHook();
+            ConnectionString::setConnectionHook(
+                    mongo::MockConnRegistry::get()->getConnStrHook());
+        }
+
+        void tearDown() {
+            ConnectionString::setConnectionHook(_originalConnectionHook);
+            ReplicaSetMonitor::remove(_replSet->getSetName(), true);
+        }
+
+        MockReplicaSet* getReplSet() {
+            return _replSet.get();
+        }
+
+    private:
+        ConnectionString::ConnectionHook* _originalConnectionHook;
+        boost::scoped_ptr<MockReplicaSet> _replSet;
+    };
+
+    TEST_F(ReplicaSetMonitorTest, SeedWithPriOnlySecDown) {
+        // Test to make sure that the monitor doesn't crash when
+        // ConnectionString::connect returns NULL
+        MockReplicaSet* replSet = getReplSet();
+        replSet->kill(replSet->getSecondaries());
+
+        // Create a monitor with primary as the only seed list and the two secondaries
+        // down so a NULL connection object will be stored for these secondaries in
+        // the _nodes vector.
+        const string replSetName(replSet->getSetName());
+        vector<HostAndPort> seedList;
+        seedList.push_back(HostAndPort(replSet->getPrimary()));
+        ReplicaSetMonitor::createIfNeeded(replSetName, seedList);
+
+        replSet->kill(replSet->getPrimary());
+
+        ReplicaSetMonitorPtr monitor = ReplicaSetMonitor::get(replSet->getSetName());
+        // Trigger calls to Node::getConnWithRefresh
+        monitor->check(true);
+    }
 }
diff --git a/src/mongo/s/shard_conn_test.cpp b/src/mongo/s/shard_conn_test.cpp
index 73b3f2008c5..060007b3bf4 100644
--- a/src/mongo/s/shard_conn_test.cpp
+++ b/src/mongo/s/shard_conn_test.cpp
@@ -15,6 +15,7 @@
 
 #include "mongo/base/init.h"
 #include "mongo/client/connpool.h"
+#include "mongo/dbtests/mock/mock_conn_registry.h"
 #include "mongo/dbtests/mock/mock_dbclient_connection.h"
 #include "mongo/platform/cstdint.h"
 #include "mongo/s/shard.h"
@@ -32,6 +33,7 @@
 
 using boost::scoped_ptr;
 using mongo::DBClientBase;
+using mongo::MockRemoteDBServer;
 using mongo::ShardConnection;
 using std::string;
 using std::vector;
@@ -68,12 +70,16 @@ namespace mongo_test {
         void setUp() {
             _maxPoolSizePerHost = mongo::PoolForHost::getMaxPerHost();
 
+            mongo::ConnectionString::setConnectionHook(
+                    mongo::MockConnRegistry::get()->getConnStrHook());
             _dummyServer = new MockRemoteDBServer(TARGET_HOST);
-            mongo::ConnectionString::setConnectionHook(_dummyServer->getConnectionHook());
+            mongo::MockConnRegistry::get()->addServer(_dummyServer);
         }
 
         void tearDown() {
             ShardConnection::clearPool();
+
+            mongo::MockConnRegistry::get()->removeServer(_dummyServer->getServerAddress());
             delete _dummyServer;
 
             mongo::PoolForHost::setMaxPerHost(_maxPoolSizePerHost);