some refactoring for SERVER-946

7 files changed, 64 insertions, 44 deletions

diff --git a/s/balance.cpp b/s/balance.cpp
index 3493f09c88f..3675cba250a 100644
--- a/s/balance.cpp
+++ b/s/balance.cpp
@@ -111,9 +111,9 @@ namespace mongo {
             return;
         
         {
-            list<Shard> all;
+            vector<Shard> all;
             Shard::getAllShards( all );
-            for ( list<Shard>::iterator i=all.begin(); i!=all.end(); ++i ){
+            for ( vector<Shard>::iterator i=all.begin(); i!=all.end(); ++i ){
                 // this just makes sure there is an entry in the map for every shard
                 Shard s = *i;
                 shards[s.getName()].size();
diff --git a/s/chunk.cpp b/s/chunk.cpp
index 0a3058cb775..3e91c3a5b3d 100644
--- a/s/chunk.cpp
+++ b/s/chunk.cpp
@@ -302,7 +302,7 @@ namespace mongo {
         if ( ! toMove )
             return false;
         
-        Shard newLocation = grid.pickShardForNewDB();
+        Shard newLocation = Shard::pick();
         if ( getShard() == newLocation ){
             // if this is the best server, then we shouldn't do anything!
             log(1) << "not moving chunk: " << toString() << " b/c would move to same place  " << newLocation.toString() << " -> " << getShard().toString() << endl;
diff --git a/s/commands_admin.cpp b/s/commands_admin.cpp
index 577c6c418e6..faa95097223 100644
--- a/s/commands_admin.cpp
+++ b/s/commands_admin.cpp
@@ -654,6 +654,9 @@ namespace mongo {
                 }
                 result.append( "added" , shard["host"].valuestrsafe() );
                 conn.done();
+                
+                Shard::reloadShardInfo();
+                    
                 return true;
             }
         } addServer;
@@ -840,13 +843,13 @@ namespace mongo {
         virtual void help( stringstream& help ) const { help << "list databases on cluster"; }
         
         bool run(const char *ns, BSONObj& jsobj, string& errmsg, BSONObjBuilder& result, bool /*fromRepl*/) {
-            list<Shard> shards;
+            vector<Shard> shards;
             Shard::getAllShards( shards );
             
             map<string,long long> sizes;
             map< string,shared_ptr<BSONObjBuilder> > dbShardInfo;
 
-            for ( list<Shard>::iterator i=shards.begin(); i!=shards.end(); i++ ){
+            for ( vector<Shard>::iterator i=shards.begin(); i!=shards.end(); i++ ){
                 Shard s = *i;
                 BSONObj x = s.runCommand( "admin" , "listDatabases" );
                 cout << s.toString() << "\t" << x.jsonString() << endl;
diff --git a/s/config.cpp b/s/config.cpp
index 90613d13c57..732de6d5482 100644
--- a/s/config.cpp
+++ b/s/config.cpp
@@ -286,26 +286,6 @@ namespace mongo {
     
     /* --- Grid --- */
     
-    Shard Grid::pickShardForNewDB(){
-        ShardConnection conn( configServer.getPrimary() );
-        
-        // TODO: this is temporary
-        
-        vector<string> all;
-        auto_ptr<DBClientCursor> c = conn->query( "config.shards" , Query() );
-        while ( c->more() ){
-            BSONObj s = c->next();
-            all.push_back( s["host"].valuestrsafe() );
-            // look at s["maxSize"] if exists
-        }
-        conn.done();
-        
-        if ( all.size() == 0 )
-            return Shard::EMPTY;
-        
-        return Shard::make( all[ rand() % all.size() ] );
-    }
-
     bool Grid::knowAboutShard( string name ) const{
         ShardConnection conn( configServer.getPrimary() );
         BSONObj shard = conn->findOne( "config.shards" , BSON( "host" << name ) );
@@ -336,7 +316,7 @@ namespace mongo {
                     if ( database == "admin" )
                         cc->_primary = configServer.getPrimary();
                     else
-                        cc->_primary = pickShardForNewDB();
+                        cc->_primary = Shard::pick();
                     
                     if ( cc->_primary.ok() ){
                         cc->save();
diff --git a/s/config.h b/s/config.h
index 9d10a1370b1..1bcc5c1fb6e 100644
--- a/s/config.h
+++ b/s/config.h
@@ -159,8 +159,6 @@ namespace mongo {
          * on next getDBConfig call will fetch from db
          */
         void removeDB( string db );
-
-        Shard pickShardForNewDB();
         
         bool knowAboutShard( string name ) const;
         
diff --git a/s/shard.cpp b/s/shard.cpp
index ed8524149f5..7805a91bb45 100644
--- a/s/shard.cpp
+++ b/s/shard.cpp
@@ -26,16 +26,8 @@ namespace mongo {
     class StaticShardInfo {
     public:
         
-        const Shard& find( const string& ident ){
-            {
-                scoped_lock lk( _mutex );
-                map<string,Shard>::iterator i = _lookup.find( ident );
-                if ( i != _lookup.end() )
-                    return i->second;
-            }
-            
-            // not in our maps, re-load all
-            
+        void reload(){
+
             list<BSONObj> all;
             {
                 ShardConnection conn( configServer.getPrimary() );
@@ -46,7 +38,6 @@ namespace mongo {
                 conn.done();
             }
             
-            
             scoped_lock lk( _mutex );
             
             for ( list<BSONObj>::iterator i=all.begin(); i!=all.end(); ++i ){
@@ -57,6 +48,21 @@ namespace mongo {
                 _lookup[name] = s;
                 _lookup[host] = s;
             }
+
+        }
+        
+        const Shard& find( const string& ident ){
+            {
+                scoped_lock lk( _mutex );
+                map<string,Shard>::iterator i = _lookup.find( ident );
+                if ( i != _lookup.end() )
+                    return i->second;
+            }
+            
+            // not in our maps, re-load all
+            reload();
+
+            scoped_lock lk( _mutex );
             
             map<string,Shard>::iterator i = _lookup.find( ident );
             if ( i == _lookup.end() ) printStackTrace();
@@ -73,7 +79,7 @@ namespace mongo {
                 _lookup[addr] = s;
         }
         
-        void getAllShards( list<Shard>& all ){
+        void getAllShards( vector<Shard>& all ){
             scoped_lock lk( _mutex );
             std::set<string> seen;
             for ( map<string,Shard>::iterator i = _lookup.begin(); i!=_lookup.end(); ++i ){
@@ -107,7 +113,7 @@ namespace mongo {
         _addr = s._addr;
     }
     
-    void Shard::getAllShards( list<Shard>& all ){
+    void Shard::getAllShards( vector<Shard>& all ){
         staticShardInfo.getAllShards( all );
     }
 
@@ -125,4 +131,21 @@ namespace mongo {
         conn.done();
         return res;
     }
+    
+    void Shard::reloadShardInfo(){
+        staticShardInfo.reload();
+    }
+    
+    Shard Shard::pick(){
+        vector<Shard> all;
+        staticShardInfo.getAllShards( all );
+        if ( all.size() == 0 ){
+            staticShardInfo.reload();
+            staticShardInfo.getAllShards( all );
+            if ( all.size() == 0 )
+                return EMPTY;
+        }
+        Shard temp = all[rand()%all.size()];
+        return temp;
+    }
 }
diff --git a/s/shard.h b/s/shard.h
index 19df0cd2848..a4571f8e90c 100644
--- a/s/shard.h
+++ b/s/shard.h
@@ -24,7 +24,8 @@
 namespace mongo {
 
     class ShardConnection;
-    
+    class ShardStatus;
+
     class Shard {
     public:
         Shard()
@@ -108,8 +109,17 @@ namespace mongo {
             return runCommand( db , BSON( simple << 1 ) );
         }
         BSONObj runCommand( const string& db , const BSONObj& cmd );
-
-        static void getAllShards( list<Shard>& all );
+        
+        ShardStatus getStatus() const;
+        
+        static void getAllShards( vector<Shard>& all );
+        
+        /**
+         * picks a Shard for more load
+         */
+        static Shard pick();
+        
+        static void reloadShardInfo();
 
         static Shard EMPTY;
 
@@ -118,6 +128,12 @@ namespace mongo {
         string _addr;
     };
 
+    class ShardStatus {
+    private:
+        long long _mapped;
+        double _writeLock;
+    };
+
     class ShardConnection {
     public:
         ShardConnection( const Shard * s )