SERVER-39234 Ensure `shardCollection` initial split works the same between config server and shard primary

5 files changed, 102 insertions, 87 deletions

diff --git a/jstests/multiVersion/shard_collection_between_mixed_version_mongods.js b/jstests/multiVersion/shard_collection_between_mixed_version_mongods.js
index 2fb86b9b8eb..9cae2839266 100644
--- a/jstests/multiVersion/shard_collection_between_mixed_version_mongods.js
+++ b/jstests/multiVersion/shard_collection_between_mixed_version_mongods.js
@@ -19,8 +19,6 @@ load("./jstests/multiVersion/libs/verify_versions.js");
     };
 
     var st = new ShardingTest(options);
-    st.stopBalancer();
-
     assert.binVersion(st.shard0, "latest");
     assert.binVersion(st.shard1, "4.0.1");
     assert.binVersion(st.shard2, "4.0.1");
@@ -28,37 +26,37 @@ load("./jstests/multiVersion/libs/verify_versions.js");
 
     var mongos = st.s0;
     var admin = mongos.getDB('admin');
-    var shards = mongos.getCollection('config.shards').find().toArray();
-    const fooDB = "fooTest";
-    const fooNS = fooDB + ".foo";
-    const barDB = "barTest";
-    const barNS = barDB + ".foo";
 
-    assert.commandWorked(admin.runCommand({enableSharding: fooDB}));
-    assert.commandWorked(admin.runCommand({enableSharding: barDB}));
-    st.ensurePrimaryShard(fooDB, shards[0]._id);
-    st.ensurePrimaryShard(barDB, shards[1]._id);
+    const kDBOnShardWithLatestBinary = "DBWithPrimaryOnLatestBinary";
+    const kNSOnLatestShard = kDBOnShardWithLatestBinary + ".Coll";
+    const kDBOnShardWithOldBinary = "DBWithPrimaryOnOldBinary";
+    const kNSOnOldShard = kDBOnShardWithOldBinary + ".Coll";
+
+    assert.commandWorked(admin.runCommand({enableSharding: kDBOnShardWithLatestBinary}));
+    assert.commandWorked(admin.runCommand({enableSharding: kDBOnShardWithOldBinary}));
+    st.ensurePrimaryShard(kDBOnShardWithLatestBinary, st.shard0.shardName);
+    st.ensurePrimaryShard(kDBOnShardWithOldBinary, st.shard1.shardName);
 
     // Test that shardCollection succeeds when both the config server and primary shard are
     // running with latest binVersion, but other shards are running with 4.0.1 which does not
     // have the new shardCollection protocol.
-    assert.commandWorked(admin.runCommand({shardCollection: fooNS, key: {a: 1}}));
+    assert.commandWorked(admin.runCommand({shardCollection: kNSOnLatestShard, key: {a: 1}}));
 
     // Test that shardCollection succeeds when the config server is running with the latest
     // binVersion, but the primary is running with 4.0.1.
-    assert.commandWorked(admin.runCommand({shardCollection: barNS, key: {a: 1}}));
+    assert.commandWorked(admin.runCommand({shardCollection: kNSOnOldShard, key: {a: 1}}));
 
-    mongos.getDB(fooDB).foo.drop();
-    mongos.getDB(barDB).foo.drop();
+    mongos.getDB(kDBOnShardWithLatestBinary).Coll.drop();
+    mongos.getDB(kDBOnShardWithOldBinary).Coll.drop();
 
     // Test that shardCollection with a hashed shard key succeeds when both the config server and
     // primary shard are running with latest binVersion, but other shards are running with 4.0.1
     // which does not have the new shardCollection protocol.
-    assert.commandWorked(admin.runCommand({shardCollection: fooNS, key: {a: "hashed"}}));
+    assert.commandWorked(admin.runCommand({shardCollection: kNSOnLatestShard, key: {a: "hashed"}}));
 
     // Test that shardCollection with a hashed shard key succeeds when the config server is running
     // with the latest binVersion, but the primary is running with 4.0.1.
-    assert.commandWorked(admin.runCommand({shardCollection: barNS, key: {a: "hashed"}}));
+    assert.commandWorked(admin.runCommand({shardCollection: kNSOnOldShard, key: {a: "hashed"}}));
 
     st.stop();
-})();
\ No newline at end of file
+})();
diff --git a/jstests/sharding/add_and_remove_shard_from_zone.js b/jstests/sharding/add_and_remove_shard_from_zone.js
new file mode 100644
index 00000000000..d4773597259
--- /dev/null
+++ b/jstests/sharding/add_and_remove_shard_from_zone.js
@@ -0,0 +1,40 @@
+/**
+ * Basic integration tests for addShardToZone command. More detailed tests can be found
+ * in sharding_catalog_add_shard_to_zone_test.cpp.
+ */
+(function() {
+    'use strict';
+
+    let st = new ShardingTest({shards: 1});
+    let mongos = st.s0;
+
+    let config = mongos.getDB('config');
+    var shardName = st.shard0.shardName;
+
+    // Test adding shard with no zone to a new zone.
+    assert.commandWorked(mongos.adminCommand({addShardToZone: shardName, zone: 'x'}));
+    var shardDoc = config.shards.findOne();
+    assert.eq(['x'], shardDoc.tags);
+
+    // Test adding zone to a shard with existing zones.
+    assert.commandWorked(mongos.adminCommand({addShardToZone: shardName, zone: 'y'}));
+    shardDoc = config.shards.findOne();
+    assert.eq(['x', 'y'], shardDoc.tags);
+
+    // Test removing shard from existing zone.
+    assert.commandWorked(mongos.adminCommand({removeShardFromZone: shardName, zone: 'x'}));
+    shardDoc = config.shards.findOne();
+    assert.eq(['y'], shardDoc.tags);
+
+    // Test removing shard from zone that no longer exists.
+    assert.commandWorked(mongos.adminCommand({removeShardFromZone: shardName, zone: 'x'}));
+    shardDoc = config.shards.findOne();
+    assert.eq(['y'], shardDoc.tags);
+
+    // Test removing the last zone from a shard
+    assert.commandWorked(mongos.adminCommand({removeShardFromZone: shardName, zone: 'y'}));
+    shardDoc = config.shards.findOne();
+    assert.eq([], shardDoc.tags);
+
+    st.stop();
+})();
diff --git a/jstests/sharding/add_shard_to_zone.js b/jstests/sharding/add_shard_to_zone.js
deleted file mode 100644
index bb142585cc2..00000000000
--- a/jstests/sharding/add_shard_to_zone.js
+++ /dev/null
@@ -1,22 +0,0 @@
-/**
- * Basic integration tests for addShardToZone command. More detailed tests can be found
- * in sharding_catalog_add_shard_to_zone_test.cpp.
- */
-(function() {
-    var st = new ShardingTest({shards: 1});
-
-    var configDB = st.s.getDB('config');
-    var shardName = configDB.shards.findOne()._id;
-
-    // Test adding shard with no zone to a new zone.
-    assert.commandWorked(st.s.adminCommand({addShardToZone: shardName, zone: 'x'}));
-    var shardDoc = configDB.shards.findOne();
-    assert.eq(['x'], shardDoc.tags);
-
-    // Test adding zone to a shard with existing zones.
-    assert.commandWorked(st.s.adminCommand({addShardToZone: shardName, zone: 'y'}));
-    shardDoc = configDB.shards.findOne();
-    assert.eq(['x', 'y'], shardDoc.tags);
-
-    st.stop();
-})();
diff --git a/jstests/sharding/remove_shard_from_zone.js b/jstests/sharding/remove_shard_from_zone.js
deleted file mode 100644
index 273d597d038..00000000000
--- a/jstests/sharding/remove_shard_from_zone.js
+++ /dev/null
@@ -1,28 +0,0 @@
-/**
- * Basic integration tests for removeShardFromZone command. More detailed tests can be found
- * in sharding_catalog_remove_shard_from_zone_test.cpp.
- */
-(function() {
-    "use strict";
-
-    var st = new ShardingTest({shards: 1});
-
-    var configDB = st.s.getDB('config');
-    var shardName = st.shard0.shardName;
-
-    assert.commandWorked(st.s.adminCommand({addShardToZone: shardName, zone: 'x'}));
-    var shardDoc = configDB.shards.findOne();
-    assert.eq(['x'], shardDoc.tags);
-
-    // Test removing shard from existing zone.
-    assert.commandWorked(st.s.adminCommand({removeShardFromZone: shardName, zone: 'x'}));
-    shardDoc = configDB.shards.findOne();
-    assert.eq([], shardDoc.tags);
-
-    // Test removing shard from zone that no longer exists.
-    assert.commandWorked(st.s.adminCommand({removeShardFromZone: shardName, zone: 'x'}));
-    shardDoc = configDB.shards.findOne();
-    assert.eq([], shardDoc.tags);
-
-    st.stop();
-})();
diff --git a/jstests/sharding/shard_collection_verify_initial_chunks.js b/jstests/sharding/shard_collection_verify_initial_chunks.js
index 0538bed1b34..e7072132b11 100644
--- a/jstests/sharding/shard_collection_verify_initial_chunks.js
+++ b/jstests/sharding/shard_collection_verify_initial_chunks.js
@@ -1,30 +1,57 @@
-//
-// Verify numInitialChunks can not be set for non hashed key or nonempty collections
-//
-
+/**
+ * Verify initial chunks are properly created and distributed in various combinations of shard key
+ * and empty/non-empty collections.
+ */
 (function() {
     'use strict';
 
-    var st = new ShardingTest({mongos: 1, shards: 2});
-    var kDbName = 'db';
-    var mongos = st.s0;
+    let st = new ShardingTest({mongos: 1, shards: 2});
+    let mongos = st.s0;
 
-    assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
+    let config = mongos.getDB("config");
+    let db = mongos.getDB('TestDB');
 
-    assert.commandFailed(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: 1}, numInitialChunks: 5}));
+    assert.commandWorked(mongos.adminCommand({enableSharding: 'TestDB'}));
+    st.ensurePrimaryShard('TestDB', st.shard1.shardName);
 
-    assert.writeOK(mongos.getDB(kDbName).foo.insert({aKey: 1}));
-    assert.commandWorked(mongos.getDB(kDbName).foo.createIndex({aKey: "hashed"}));
-    assert.commandFailed(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: "hashed"}, numInitialChunks: 5}));
+    function checkChunkCounts(collName, chunksOnShard0, chunksOnShard1) {
+        let counts = st.chunkCounts(collName, 'TestDB');
+        assert.eq(chunksOnShard0,
+                  counts[st.shard0.shardName],
+                  'Count mismatch on shard0: ' + tojson(counts));
+        assert.eq(chunksOnShard1,
+                  counts[st.shard1.shardName],
+                  'Count mismatch on shard1: ' + tojson(counts));
+    }
 
-    assert.writeOK(mongos.getDB(kDbName).foo.remove({}));
+    // Unsupported: Range sharding + numInitialChunks
+    assert.commandFailed(mongos.adminCommand(
+        {shardCollection: 'TestDB.RangeCollEmpty', key: {aKey: 1}, numInitialChunks: 6}));
+
+    // Unsupported: Hashed sharding + numInitialChunks + non-empty collection
+    assert.writeOK(db.HashedCollNotEmpty.insert({aKey: 1}));
+    assert.commandWorked(db.HashedCollNotEmpty.createIndex({aKey: "hashed"}));
+    assert.commandFailed(mongos.adminCommand({
+        shardCollection: 'TestDB.HashedCollNotEmpty',
+        key: {aKey: "hashed"},
+        numInitialChunks: 6
+    }));
+
+    // Supported: Hashed sharding + numInitialChunks + empty collection
+    // Expected: Even chunk distribution
+    assert.commandWorked(db.HashedCollEmpty.createIndex({aKey: "hashed"}));
     assert.commandWorked(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: "hashed"}, numInitialChunks: 5}));
-
-    mongos.getDB(kDbName).dropDatabase();
+        {shardCollection: 'TestDB.HashedCollEmpty', key: {aKey: "hashed"}, numInitialChunks: 6}));
+    checkChunkCounts('HashedCollEmpty', 3, 3);
+
+    // Supported: Hashed sharding + numInitialChunks + non-existent collection
+    // Expected: Even chunk distribution
+    assert.commandWorked(mongos.adminCommand({
+        shardCollection: 'TestDB.HashedCollNonExistent',
+        key: {aKey: "hashed"},
+        numInitialChunks: 6
+    }));
+    checkChunkCounts('HashedCollNonExistent', 3, 3);
 
     st.stop();
-
 })();