SERVER-3627: ability to assign shards to chunks

5 files changed, 54 insertions, 18 deletions

diff --git a/s/chunk.cpp b/s/chunk.cpp
index b2e67e0e6bc..f2ec063ae3a 100644
--- a/s/chunk.cpp
+++ b/s/chunk.cpp
@@ -664,38 +664,50 @@ namespace mongo {
         return _key.hasShardKey( obj );
     }
 
-    void ChunkManager::createFirstChunks( const Shard& shard ) const {
+    void ChunkManager::createFirstChunks( const Shard& shard , const vector<BSONObj>& chunkSplitPoints , const vector<Shard>& chunkShards ) const {
         // TODO distlock?
         assert( _chunkMap.size() == 0 );
 
+        vector<BSONObj> splitPoints;
+        vector<Shard> shards;
         unsigned long long numObjects = 0;
-        {
-            // get stats to see if there is any data
-            ScopedDbConnection shardConn( shard.getConnString() );
-            numObjects = shardConn->count( getns() );
-            shardConn.done();
+        Chunk c(this, _key.globalMin(), _key.globalMax(), shard);
+
+        if ( chunkSplitPoints.size() == 0 ) {
+            // discover split points
+            {
+                // get stats to see if there is any data
+                ScopedDbConnection shardConn( shard.getConnString() );
+                numObjects = shardConn->count( getns() );
+                shardConn.done();
+            }
+
+            if ( numObjects > 0 )
+                c.pickSplitVector( splitPoints , Chunk::MaxChunkSize );
+        } else {
+            splitPoints = chunkSplitPoints;
+            shards = chunkShards;
         }
 
+        // make sure there is at least 1 shard, the primary
+        if ( !shards.size() )
+            shards.push_back( shard );
+
         // this is the first chunk; start the versioning from scratch
         ShardChunkVersion version;
         version.incMajor();
-
-        Chunk c(this, _key.globalMin(), _key.globalMax(), shard);
-
-        vector<BSONObj> splitPoints;
-        if ( numObjects > 0 )
-            c.pickSplitVector( splitPoints , Chunk::MaxChunkSize );
         
         log() << "going to create " << splitPoints.size() + 1 << " chunk(s) for: " << _ns << endl;
         
 
         ScopedDbConnection conn( configServer.modelServer() );        
 
+        unsigned si = 0;
         for ( unsigned i=0; i<=splitPoints.size(); i++ ) {
             BSONObj min = i == 0 ? _key.globalMin() : splitPoints[i-1];
             BSONObj max = i < splitPoints.size() ? splitPoints[i] : _key.globalMax();
             
-            Chunk temp( this , min , max , shard );
+            Chunk temp( this , min , max , shards[si] );
         
             BSONObjBuilder chunkBuilder;
             temp.serialize( chunkBuilder , version );
@@ -704,6 +716,8 @@ namespace mongo {
             conn->update( Chunk::chunkMetadataNS, QUERY( "_id" << temp.genID() ), chunkObj,  true, false );
 
             version.incMinor();
+            if ( ++si == shards.size() )
+                si = 0;
         }
 
         string errmsg = conn->getLastError();
diff --git a/s/chunk.h b/s/chunk.h
index f451f74bc79..00e1e4010f8 100644
--- a/s/chunk.h
+++ b/s/chunk.h
@@ -305,7 +305,7 @@ namespace mongo {
         int numChunks() const { return _chunkMap.size(); }
         bool hasShardKey( const BSONObj& obj ) const;
 
-        void createFirstChunks( const Shard& shard ) const; // only call from DBConfig::shardCollection
+        void createFirstChunks( const Shard& shard , const vector<BSONObj>& chunkSplitPoints , const vector<Shard>& chunkShards ) const; // only call from DBConfig::shardCollection
         ChunkPtr findChunk( const BSONObj& obj ) const;
         ChunkPtr findChunkOnServer( const Shard& shard ) const;
 
diff --git a/s/commands_admin.cpp b/s/commands_admin.cpp
index 4aa82c74ac8..f72ff4b9eb7 100644
--- a/s/commands_admin.cpp
+++ b/s/commands_admin.cpp
@@ -529,7 +529,23 @@ namespace mongo {
 
                 tlog() << "CMD: shardcollection: " << cmdObj << endl;
 
-                config->shardCollection( ns , key , careAboutUnique );
+                vector<BSONObj> pts;
+                if (cmdObj.hasField("splitPoints")) {
+                    if ( cmdObj.getField("splitPoints").type() != Array ) {
+                        errmsg = "Value of splitPoints must be an array of objects";
+                        return false;
+                    }
+
+                    vector<BSONElement> elmts = cmdObj.getField("splitPoints").Array();
+                    for ( unsigned i = 0 ; i < elmts.size() ; ++i) {
+                        if ( elmts[i].type() != Object ) {
+                            errmsg = "Elements in the splitPoints array must be objects";
+                            return false;
+                        }
+                        pts.push_back( elmts[i].Obj() );
+                    }
+                }
+                config->shardCollection( ns , key , careAboutUnique, pts );
 
                 result << "collectionsharded" << ns;
                 return true;
diff --git a/s/config.cpp b/s/config.cpp
index a410bc75255..21638d62a86 100644
--- a/s/config.cpp
+++ b/s/config.cpp
@@ -141,7 +141,7 @@ namespace mongo {
         _save();
     }
 
-    ChunkManagerPtr DBConfig::shardCollection( const string& ns , ShardKeyPattern fieldsAndOrder , bool unique ) {
+    ChunkManagerPtr DBConfig::shardCollection( const string& ns , ShardKeyPattern fieldsAndOrder , bool unique , const vector<BSONObj>& splitPoints ) {
         uassert( 8042 , "db doesn't have sharding enabled" , _shardingEnabled );
         uassert( 13648 , str::stream() << "can't shard collection because not all config servers are up" , configServer.allUp() );
 
@@ -157,7 +157,13 @@ namespace mongo {
             ci.shard( ns , fieldsAndOrder , unique );
             ChunkManagerPtr cm = ci.getCM();
             uassert( 13449 , "collections already sharded" , (cm->numChunks() == 0) );
-            cm->createFirstChunks( getPrimary() );
+
+            vector<Shard> shards;
+            if ( splitPoints.size() ) {
+                getPrimary().getAllShards( shards );
+            }
+
+            cm->createFirstChunks( getPrimary() , splitPoints , shards );
             _save();
         }
 
diff --git a/s/config.h b/s/config.h
index 996c9318ad9..dbcbaf7d502 100644
--- a/s/config.h
+++ b/s/config.h
@@ -132,7 +132,7 @@ namespace mongo {
         }
 
         void enableSharding();
-        ChunkManagerPtr shardCollection( const string& ns , ShardKeyPattern fieldsAndOrder , bool unique );
+        ChunkManagerPtr shardCollection( const string& ns , ShardKeyPattern fieldsAndOrder , bool unique , const vector<BSONObj>& splitPoints );
 
         /**
            @return true if there was sharding info to remove