From c77e4c8ee9f7696bfa5fc0da9b6c10284c2805a5 Mon Sep 17 00:00:00 2001
From: Eliot Horowitz <eliot@10gen.com>
Date: Thu, 17 Feb 2011 01:10:50 -0500
Subject: [PATCH] put all replica set nods in shard map

---
 jstests/slowNightly/sharding_rs2.js |  1 +
 s/shard.cpp                         | 55 ++++++++++++++++++++++++-----
 s/shard.h                           |  2 ++
 3 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/jstests/slowNightly/sharding_rs2.js b/jstests/slowNightly/sharding_rs2.js
index ebff712c4b4..cd7cf6888ad 100644
--- a/jstests/slowNightly/sharding_rs2.js
+++ b/jstests/slowNightly/sharding_rs2.js
@@ -157,6 +157,7 @@ assert.eq( 100 , ts.find().itcount() , "E5" )
 printjson( ts.find().batchSize(5).explain() )
 assert.eq( 100 , ts.find().batchSize(5).itcount() , "E6" )
 
+printjson( db.adminCommand( "getShardMap" ) );
 
 
 s.stop()
diff --git a/s/shard.cpp b/s/shard.cpp
index fa8798d4f0d..14ba61681a7 100644
--- a/s/shard.cpp
+++ b/s/shard.cpp
@@ -20,6 +20,7 @@
 #include "shard.h"
 #include "config.h"
 #include "request.h"
+#include "../db/commands.h"
 #include <set>
 
 namespace mongo {
@@ -76,13 +77,7 @@ namespace mongo {
 
                 Shard s( name , host , maxSize , isDraining );
                 _lookup[name] = s;
-                _lookup[host] = s;
-
-                // add rs name to lookup (if it exists)
-                size_t pos;
-                if ((pos = host.find('/', 0)) != string::npos) {
-                    _lookup[host.substr(0, pos)] = s;
-                }
+                _installHost( host , s );
             }
 
         }
@@ -119,7 +114,22 @@ namespace mongo {
             if ( setName )
                 _lookup[name] = s;
             if ( setAddr )
-	        _lookup[s.getConnString()] = s;
+                _installHost( s.getConnString() , s );
+        }
+
+        void _installHost( const string& host , const Shard& s ) {
+            _lookup[host] = s;
+            
+            const ConnectionString& cs = s.getAddress();
+            if ( cs.type() == ConnectionString::SET ) {
+                if ( cs.getSetName().size() )
+                    _lookup[ cs.getSetName() ] = s;
+                
+                vector<HostAndPort> servers = cs.getServers();
+                for ( unsigned i=0; i<servers.size(); i++ ) {
+                    _lookup[ servers[i].toString() ] = s;
+                }
+            }
         }
         
         void remove( const string& name ) {
@@ -169,12 +179,41 @@ namespace mongo {
 
             return false;
         }
+        
+        bool getShardMap( BSONObjBuilder& result , string& errmsg ) const {
+            scoped_lock lk( _mutex );
+
+            BSONObjBuilder b( _lookup.size() + 50 );
+
+            for ( map<string,Shard>::const_iterator i = _lookup.begin(); i!=_lookup.end(); ++i ) {
+                b.append( i->first , i->second.getConnString() );
+            }
+            
+            result.append( "map" , b.obj() );
+            
+            return true;
+        }
 
     private:
         map<string,Shard> _lookup;
         mutable mongo::mutex _mutex;
     } staticShardInfo;
 
+    
+    class CmdGetShardMap : public Command {
+    public:
+        CmdGetShardMap() : Command( "getShardMap" ){}
+        virtual void help( stringstream &help ) const { help<<"internal"; }
+        virtual LockType locktype() const { return NONE; }
+        virtual bool slaveOk() const { return true; }
+        virtual bool adminOnly() const { return true; }
+
+        virtual bool run(const string&, mongo::BSONObj&, std::string& errmsg , mongo::BSONObjBuilder& result, bool) {
+            return staticShardInfo.getShardMap( result , errmsg );
+        }
+    } cmdGetShardMap;
+    
+
     void Shard::_setAddr( const string& addr ) {
         _addr = addr;
         if ( _addr.size() ) {
diff --git a/s/shard.h b/s/shard.h
index b1e15b7850a..836ffe7753d 100644
--- a/s/shard.h
+++ b/s/shard.h
@@ -68,6 +68,8 @@ namespace mongo {
         void reset( const string& ident );
 
         void setAddress( const ConnectionString& cs );
+        
+        ConnectionString getAddress() const { return _cs; }
 
         string getName() const {
             assert( _name.size() );