some refactoring for SERVER-946

1 files changed, 36 insertions, 13 deletions

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;
+    }
 }