DBClientReplicaSet::slaveConn hooked up again - final prep for SEVER-1634

3 files changed, 109 insertions, 13 deletions

diff --git a/client/dbclient_rs.cpp b/client/dbclient_rs.cpp
index 730d92b5e73..c73a7476985 100644
--- a/client/dbclient_rs.cpp
+++ b/client/dbclient_rs.cpp
@@ -93,6 +93,51 @@ namespace mongo {
         scoped_lock lk( _lock );
         return _nodes[_master].addr;
     }
+    
+    HostAndPort ReplicaSetMonitor::getSlave(){
+        int x = rand() % _nodes.size();
+        {
+            scoped_lock lk( _lock );
+            for ( int i=0; i<_nodes.size(); i++ ){
+                int p = ( i + x ) % _nodes.size();
+                if ( p == _master )
+                    continue;
+                if ( _nodes[p].ok )
+                    return _nodes[p].addr;
+            }
+        }
+
+        {
+            scoped_lock lk( _lock );
+            for ( int i=0; i<_nodes.size(); i++ ){
+                _nodes[i].ok = true;
+            }
+        }
+        
+        {
+            scoped_lock lk( _lock );
+            for ( int i=0; i<_nodes.size(); i++ ){
+                int p = ( i + x ) % _nodes.size();
+                if ( p == _master )
+                    continue;
+                if ( _nodes[p].ok )
+                    return _nodes[p].addr;
+            }
+        }
+        
+        return _nodes[0].addr;
+    }
+
+    /**
+     * notify the monitor that server has faild
+     */
+    void ReplicaSetMonitor::notifySlaveFailure( const HostAndPort& server ){
+        int x = _find( server );
+        if ( x >= 0 ){
+            scoped_lock lk( _lock );
+            _nodes[x].ok = false;
+        }
+    }
 
     bool ReplicaSetMonitor::_checkConnection( DBClientConnection * c , string& maybePrimary , bool verbose ) {
         bool good = false;
@@ -186,6 +231,15 @@ namespace mongo {
         return -1;
     }
 
+    int ReplicaSetMonitor::_find( const HostAndPort& server ) const {
+        scoped_lock lk( _lock );
+        for ( unsigned i=0; i<_nodes.size(); i++ )
+            if ( _nodes[i].addr == server )
+                return i;
+        return -1;
+    }
+
+
     mongo::mutex ReplicaSetMonitor::_setsLock( "ReplicaSetMonitor" );
     map<string,ReplicaSetMonitorPtr> ReplicaSetMonitor::_sets;
 
@@ -207,20 +261,44 @@ namespace mongo {
                 return _master.get();
             _monitor->notifyFailure( _masterHost );
         }
-        
-        _masterHost = _monitor->getMaster();
-        _master.reset( new DBClientConnection( true ) );
-        _master->connect( _masterHost );
-        
+
+        HostAndPort h = _monitor->getMaster();
+        if ( h != _masterHost ){
+            _masterHost = h;
+            _master.reset( new DBClientConnection( true ) );
+            _master->connect( _masterHost );
+            _auth( _master.get() );
+        }
+        return _master.get();
+    }
+
+    DBClientConnection * DBClientReplicaSet::checkSlave() {
+        if ( _slave ){
+            if ( ! _slave->isFailed() )
+                return _slave.get();
+            _monitor->notifySlaveFailure( _slaveHost );
+        }
+
+        HostAndPort h = _monitor->getSlave();
+        if ( h != _slaveHost ){
+            _slaveHost = h;
+            _slave.reset( new DBClientConnection( true ) );
+            _slave->connect( _slaveHost );
+            _auth( _slave.get() );
+        }
+        return _slave.get();
+    }
+
+
+    void DBClientReplicaSet::_auth( DBClientConnection * conn ){
         for ( list<AuthInfo>::iterator i=_auths.begin(); i!=_auths.end(); ++i ){
             const AuthInfo& a = *i;
             string errmsg;
-            if ( ! _master->auth( a.dbname , a.username , a.pwd , errmsg, a.digestPassword ) )
+            if ( ! conn->auth( a.dbname , a.username , a.pwd , errmsg, a.digestPassword ) )
                 warning() << "cached auth failed for set: " << _monitor->setName() << " db: " << a.dbname << " user: " << a.username << endl;
 
         }
 
-        return _master.get();
     }
 
     DBClientConnection& DBClientReplicaSet::masterConn(){
@@ -228,9 +306,7 @@ namespace mongo {
     }
 
     DBClientConnection& DBClientReplicaSet::slaveConn(){
-        // TODO
-        assert(0);
-        return masterConn();
+        return *checkSlave();
     }
 
     bool DBClientReplicaSet::connect(){
diff --git a/client/dbclient_rs.h b/client/dbclient_rs.h
index 91bde3b498c..0fc41aa8d0a 100644
--- a/client/dbclient_rs.h
+++ b/client/dbclient_rs.h
@@ -50,6 +50,14 @@ namespace mongo {
          */
         void notifyFailure( const HostAndPort& server );
 
+        /** @return a random slave that is ok for reads */
+        HostAndPort getSlave();
+
+        /**
+         * notify the monitor that server has faild
+         */
+        void notifySlaveFailure( const HostAndPort& server );        
+
         string setName() const { return _name; }
 
         string getServerAddress() const;
@@ -67,14 +75,20 @@ namespace mongo {
         bool _checkConnection( DBClientConnection * c , string& maybePrimary , bool verbose );
 
         int _find( const string& server ) const ;
+        int _find( const HostAndPort& server ) const ;
 
         mutable mongo::mutex _lock; // protects _nodes
 
         string _name;
         struct Node {
-            Node( const HostAndPort& a , DBClientConnection* c ) : addr( a ) , conn(c){}
+            Node( const HostAndPort& a , DBClientConnection* c ) : addr( a ) , conn(c) , ok(true){}
             HostAndPort addr;
             DBClientConnection* conn;
+
+            // if this node is in a failure state
+            // used for slave routing
+            // this is too simple, should make it better
+            bool ok; 
         };
         vector<Node> _nodes;
 
@@ -170,12 +184,17 @@ namespace mongo {
     private:
 
         DBClientConnection * checkMaster();
-        void _checkMaster();
+        DBClientConnection * checkSlave();
+        
+        void _auth( DBClientConnection * conn );
 
         ReplicaSetMonitorPtr _monitor;
 
         HostAndPort _masterHost;        
         scoped_ptr<DBClientConnection> _master;
+
+        HostAndPort _slaveHost;
+        scoped_ptr<DBClientConnection> _slave;
         
         /**
          * for storing authentication info
diff --git a/client/examples/rs.cpp b/client/examples/rs.cpp
index 9a3e4a2fee0..b024d030fa3 100644
--- a/client/examples/rs.cpp
+++ b/client/examples/rs.cpp
@@ -33,7 +33,7 @@ int main( int argc , const char ** argv ){
         return 1;
     }
     
-    DBClientBase * conn = cs.connect( errmsg );
+    DBClientReplicaSet * conn = (DBClientReplicaSet*)cs.connect( errmsg );
     if ( ! conn ){
         cout << "error connecting: " << errmsg << endl;
         return 2;
@@ -46,6 +46,7 @@ int main( int argc , const char ** argv ){
         try {
             conn->update( collName , BSONObj() , BSON( "$inc" << BSON( "x" << 1 ) ) , true );
             cout << conn->findOne( collName , BSONObj() ) << endl;
+            cout << "\t" << conn->slaveConn().findOne( collName , BSONObj() , 0 , QueryOption_SlaveOk ) << endl;
         }
         catch ( std::exception& e ){
             cout << "ERROR: " << e.what() << endl;