initialize shard connection with serverId, only setShardV on single/replsets SERVER-4020

5 files changed, 102 insertions, 30 deletions

diff --git a/s/d_state.cpp b/s/d_state.cpp
index a6fa2b1a3f8..bd6f6ff0121 100644
--- a/s/d_state.cpp
+++ b/s/d_state.cpp
@@ -484,6 +484,13 @@ namespace mongo {
                 shardingState.gotShardHost( cmdObj["shardHost"].String() );
             }
             
+
+            // Handle initial shard connection
+            if( cmdObj["version"].eoo() && cmdObj["init"].trueValue() ){
+                result.append( "initialized", true );
+                return true;
+            }
+
             // step 2
             
             string ns = cmdObj["setShardVersion"].valuestrsafe();
diff --git a/s/shard.cpp b/s/shard.cpp
index 32a1528630a..e3dc89caf3c 100644
--- a/s/shard.cpp
+++ b/s/shard.cpp
@@ -365,12 +365,26 @@ namespace mongo {
                     conn->auth("local", internalSecurity.user, internalSecurity.pwd, err, false));
         }
 
-        if ( _shardedConnections ) {
-            conn->simpleCommand( "admin" , 0 , "setShardVersion" );
+        if ( _shardedConnections && isVersionableCB( conn ) ) {
+
+            // We must initialize sharding on all connections, so that we get exceptions if sharding is enabled on
+            // the collection.
+            BSONObj result;
+            bool ok = initShardVersionCB( *conn, result );
+
+            // assert that we actually successfully setup sharding
+            uassert( 15903, str::stream() << "could not initialize sharding on connection " << (*conn).toString() <<
+                        ( result["errmsg"].type() == String ? causedBy( result["errmsg"].String() ) :
+                                                              causedBy( (string)"unknown failure : " + result.toString() ) ), ok );
+
         }
     }
 
     void ShardingConnectionHook::onDestroy( DBClientBase * conn ) {
-        resetShardVersionCB( conn );
+
+        if( _shardedConnections && isVersionableCB( conn ) ){
+            resetShardVersionCB( conn );
+        }
+
     }
 }
diff --git a/s/shard_version.cpp b/s/shard_version.cpp
index 71a4bff3460..35a970cf28a 100644
--- a/s/shard_version.cpp
+++ b/s/shard_version.cpp
@@ -31,6 +31,8 @@ namespace mongo {
 
     // when running in sharded mode, use chunk shard version control
 
+    static bool isVersionable( DBClientBase * conn );
+    static bool initShardVersion( DBClientBase & conn, BSONObj& result );
     static bool checkShardVersion( DBClientBase & conn , const string& ns , bool authoritative = false , int tryNumber = 1 );
     static void resetShardVersion( DBClientBase * conn );
 
@@ -40,6 +42,8 @@ namespace mongo {
         //
         // TODO: Better encapsulate this mechanism.
         //
+        isVersionableCB = isVersionable;
+        initShardVersionCB = initShardVersion;
         checkShardVersionCB = checkShardVersion;
         resetShardVersionCB = resetShardVersion;
     }
@@ -79,6 +83,63 @@ namespace mongo {
         connectionShardStatus.reset( conn );
     }
 
+    bool isVersionable( DBClientBase* conn ){
+        return conn->type() == ConnectionString::MASTER || conn->type() == ConnectionString::SET;
+    }
+
+    DBClientBase* getVersionable( DBClientBase* conn ){
+
+        switch ( conn->type() ) {
+        case ConnectionString::INVALID:
+           massert( 15904, str::stream() << "cannot set version on invalid connection " << conn->toString(), false );
+           return NULL;
+        case ConnectionString::MASTER:
+           return conn;
+        case ConnectionString::PAIR:
+           massert( 15905, str::stream() << "cannot set version or shard on pair connection " << conn->toString(), false );
+           return NULL;
+        case ConnectionString::SYNC:
+           massert( 15906, str::stream() << "cannot set version or shard on sync connection " << conn->toString(), false );
+           return NULL;
+        case ConnectionString::SET:
+           DBClientReplicaSet* set = (DBClientReplicaSet*) conn;
+           return &( set->masterConn() );
+        }
+
+        assert( false );
+        return NULL;
+    }
+
+    extern OID serverID;
+
+    bool initShardVersion( DBClientBase& conn_in, BSONObj& result ){
+
+        WriteBackListener::init( conn_in );
+
+        DBClientBase* conn = getVersionable( &conn_in );
+        assert( conn ); // errors thrown above
+
+        BSONObjBuilder cmdBuilder;
+
+        cmdBuilder.append( "setShardVersion" , "" );
+        cmdBuilder.appendBool( "init", true );
+        cmdBuilder.append( "configdb" , configServer.modelServer() );
+        cmdBuilder.appendOID( "serverID" , &serverID );
+        cmdBuilder.appendBool( "authoritative" , true );
+
+        BSONObj cmd = cmdBuilder.obj();
+
+        LOG(1) << "initializing shard connection to " << conn->toString() << endl;
+        LOG(2) << "initial sharding settings : " << cmd << endl;
+
+        bool ok = conn->runCommand( "admin" , cmd , result );
+
+        LOG(3) << "initial sharding result : " << result << endl;
+
+        return ok;
+
+    }
+
     /**
      * @return true if had to do something
      */
@@ -91,30 +152,8 @@ namespace mongo {
         if ( ! conf )
             return false;
 
-        DBClientBase* conn = 0;
-
-        switch ( conn_in.type() ) {
-        case ConnectionString::INVALID:
-            assert(0);
-            break;
-        case ConnectionString::MASTER:
-            // great
-            conn = &conn_in;
-            break;
-        case ConnectionString::PAIR:
-            assert( ! "pair not support for sharding" );
-            break;
-        case ConnectionString::SYNC:
-            // TODO: we should check later that we aren't actually sharded on this
-            conn = &conn_in;
-            break;
-        case ConnectionString::SET:
-            DBClientReplicaSet* set = (DBClientReplicaSet*)&conn_in;
-            conn = &(set->masterConn());
-            break;
-        }
-
-        assert(conn);
+        DBClientBase* conn = getVersionable( &conn_in );
+        assert(conn); // errors thrown above
 
         unsigned long long officialSequenceNumber = 0;
 
diff --git a/s/shardconnection.cpp b/s/shardconnection.cpp
index 04b49f2868a..964229a97b8 100644
--- a/s/shardconnection.cpp
+++ b/s/shardconnection.cpp
@@ -29,6 +29,14 @@ namespace mongo {
     //
     // TODO: better encapsulate this mechanism.
 
+    bool defaultIsVersionable( DBClientBase * conn ){
+        return false;
+    }
+
+    bool defaultInitShardVersion( DBClientBase & conn, BSONObj& result ){
+        return false;
+    }
+
     bool defaultCheckShardVersion( DBClientBase & conn , const string& ns , bool authoritative , int tryNumber ) {
         // no-op in mongod
         return false;
@@ -38,6 +46,8 @@ namespace mongo {
         // no-op in mongod
     }
 
+    boost::function1<bool, DBClientBase* > isVersionableCB = defaultIsVersionable;
+    boost::function2<bool, DBClientBase&, BSONObj& > initShardVersionCB = defaultInitShardVersion;
     boost::function4<bool, DBClientBase&, const string&, bool, int> checkShardVersionCB = defaultCheckShardVersion;
     boost::function1<void, DBClientBase*> resetShardVersionCB = defaultResetShardVersion;
 
@@ -71,7 +81,7 @@ namespace mongo {
                     /* if we're shutting down, don't want to initiate release mechanism as it is slow,
                        and isn't needed since all connections will be closed anyway */
                     if ( inShutdown() ) {
-                        resetShardVersionCB( ss->avail );
+                        if( isVersionableCB( ss->avail ) ) resetShardVersionCB( ss->avail );
                         delete ss->avail;
                     }
                     else
@@ -226,7 +236,7 @@ namespace mongo {
             return;
         _finishedInit = true;
 
-        if ( _ns.size() ) {
+        if ( _ns.size() && isVersionableCB( _conn ) ) {
             _setVersion = checkShardVersionCB( *_conn , _ns , false , 1 );
         }
         else {
@@ -245,7 +255,7 @@ namespace mongo {
 
     void ShardConnection::kill() {
         if ( _conn ) {
-            resetShardVersionCB( _conn );
+            if( isVersionableCB( _conn ) ) resetShardVersionCB( _conn );
             delete _conn;
             _conn = 0;
             _finishedInit = true;
diff --git a/s/util.h b/s/util.h
index 2bf1c94a961..b1927130185 100644
--- a/s/util.h
+++ b/s/util.h
@@ -152,6 +152,8 @@ namespace mongo {
         string _ns;
     };
 
+    extern boost::function1<bool, DBClientBase* > isVersionableCB;
+    extern boost::function2<bool, DBClientBase&, BSONObj& > initShardVersionCB;
     extern boost::function4<bool, DBClientBase&, const string&, bool, int> checkShardVersionCB;
     extern boost::function1<void, DBClientBase*> resetShardVersionCB;