SERVER-56227 Add user-facing command for permitMigrations using DDL coordinator

1 files changed, 162 insertions, 0 deletions

diff --git a/jstests/sharding/move_chunk_permitMigrations.js b/jstests/sharding/move_chunk_permitMigrations.js
new file mode 100644
index 00000000000..9f21ba01673
--- /dev/null
+++ b/jstests/sharding/move_chunk_permitMigrations.js
@@ -0,0 +1,162 @@
+/**
+ * Tests that a collection with permitMigrations: false in config.collections prohibits committing a
+ * moveChunk and disables the balancer.
+ *
+ * @tags: [
+ *   does_not_support_stepdowns,
+ *   requires_fcv_52,
+ * ]
+ */
+(function() {
+'use strict';
+
+load('jstests/libs/fail_point_util.js');
+load('jstests/libs/parallel_shell_helpers.js');
+load("jstests/sharding/libs/find_chunks_util.js");
+load("jstests/sharding/libs/shard_versioning_util.js");
+
+const st = new ShardingTest({shards: 2});
+const configDB = st.s.getDB("config");
+const dbName = 'AllowMigrations';
+
+// Resets database dbName and enables sharding and establishes shard0 as primary, test case agnostic
+const setUpDb = function setUpDatabaseAndEnableSharding() {
+    assert.commandWorked(st.s.getDB(dbName).dropDatabase());
+    assert.commandWorked(
+        st.s.adminCommand({enableSharding: dbName, primaryShard: st.shard0.shardName}));
+};
+
+// Use the setAllowMigrations command to set the permitMigrations flag in the collection.
+const setAllowMigrations = function(ns, allow) {
+    assert.commandWorked(st.s.adminCommand({setAllowMigrations: ns, allowMigrations: allow}));
+};
+
+// Tests that moveChunk does not succeed when setAllowMigrations is called with a false value.
+(function testSetAllowMigrationsFalsePreventsMoveChunk() {
+    setUpDb();
+
+    const collName = "collA";
+    const ns = dbName + "." + collName;
+
+    assert.commandWorked(st.s.getDB(dbName).getCollection(collName).insert({_id: 0}));
+    assert.commandWorked(st.s.getDB(dbName).getCollection(collName).insert({_id: 1}));
+    assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {_id: 1}}));
+
+    setAllowMigrations(ns, false);
+
+    // setAllowMigrations was called, sending a new moveChunk command should fail.
+    assert.commandFailedWithCode(
+        st.s.adminCommand({moveChunk: ns, find: {_id: 0}, to: st.shard1.shardName}),
+        ErrorCodes.ConflictingOperationInProgress);
+})();
+
+// Tests setAllowMigrations disables balancing for collB and does not interfere with
+// balancing for collA.
+//
+// collBSetParams specify the field(s) that will be set on the collB in config.collections.
+const testBalancer = function(setAllowMigrations, collBSetNoBalanceParam) {
+    setUpDb();
+
+    const collAName = "collA";
+    const collBName = "collB";
+    const collA = st.s.getCollection(`${dbName}.${collAName}`);
+    const collB = st.s.getCollection(`${dbName}.${collBName}`);
+
+    assert.commandWorked(st.s.adminCommand({shardCollection: collA.getFullName(), key: {_id: 1}}));
+    assert.commandWorked(st.s.adminCommand({shardCollection: collB.getFullName(), key: {_id: 1}}));
+
+    // Split both collections into 4 chunks so balancing can occur.
+    for (let coll of [collA, collB]) {
+        coll.insert({_id: 1});
+        coll.insert({_id: 10});
+        coll.insert({_id: 20});
+        coll.insert({_id: 30});
+
+        assert.commandWorked(st.splitAt(coll.getFullName(), {_id: 10}));
+        assert.commandWorked(st.splitAt(coll.getFullName(), {_id: 20}));
+        assert.commandWorked(st.splitAt(coll.getFullName(), {_id: 30}));
+
+        // Confirm the chunks are initially unbalanced. All chunks should start out on shard0
+        // (primary shard for the database).
+        const balancerStatus = assert.commandWorked(
+            st.s0.adminCommand({balancerCollectionStatus: coll.getFullName()}));
+        assert.eq(balancerStatus.balancerCompliant, false);
+        assert.eq(balancerStatus.firstComplianceViolation, 'chunksImbalance');
+        assert.eq(4,
+                  findChunksUtil
+                      .findChunksByNs(configDB, coll.getFullName(), {shard: st.shard0.shardName})
+                      .count());
+    }
+
+    jsTestLog(`Disabling balancing of ${collB.getFullName()} with setAllowMigrations ${
+        setAllowMigrations} and parameters ${tojson(collBSetNoBalanceParam)}`);
+    assert.commandWorked(
+        configDB.collections.update({_id: collB.getFullName()}, {$set: collBSetNoBalanceParam}));
+
+    setAllowMigrations(collB.getFullName(), setAllowMigrations);
+
+    st.startBalancer();
+    assert.soon(() => {
+        st.awaitBalancerRound();
+        const shard0Chunks =
+            findChunksUtil
+                .findChunksByNs(configDB, collA.getFullName(), {shard: st.shard0.shardName})
+                .itcount();
+        const shard1Chunks =
+            findChunksUtil
+                .findChunksByNs(configDB, collA.getFullName(), {shard: st.shard1.shardName})
+                .itcount();
+        jsTestLog(`shard0 chunks ${shard0Chunks}, shard1 chunks ${shard1Chunks}`);
+        return shard0Chunks == 2 && shard1Chunks == 2;
+    }, `Balancer failed to balance ${collA.getFullName()}`, 1000 * 60 * 10);
+    st.stopBalancer();
+
+    const collABalanceStatus =
+        assert.commandWorked(st.s.adminCommand({balancerCollectionStatus: collA.getFullName()}));
+    assert.eq(collABalanceStatus.balancerCompliant, true);
+
+    // Test that collB remains unbalanced.
+    const collBBalanceStatus =
+        assert.commandWorked(st.s.adminCommand({balancerCollectionStatus: collB.getFullName()}));
+    assert.eq(collBBalanceStatus.balancerCompliant, false);
+    assert.eq(collBBalanceStatus.firstComplianceViolation, 'chunksImbalance');
+    assert.eq(
+        4,
+        findChunksUtil.findChunksByNs(configDB, collB.getFullName(), {shard: st.shard0.shardName})
+            .count());
+};
+
+const testSetAllowMigrationsCommand = function() {
+    setUpDb();
+
+    const collName = "foo";
+    const ns = dbName + "." + collName;
+
+    assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {x: 1}}));
+
+    ShardVersioningUtil.assertCollectionVersionEquals(st.shard0, ns, Timestamp(1, 0));
+
+    // Use setAllowMigrations to forbid migrations from happening
+    setAllowMigrations(ns, false);
+
+    // Check that allowMigrations has been set to 'false' on the configsvr config.collections.
+    assert.eq(false, configDB.collections.findOne({_id: ns}).permitMigrations);
+
+    // Use setAllowMigrations to allow migrations to happen
+    setAllowMigrations(ns, true);
+
+    // Check that permitMigrations has been unset (that implies migrations are allowed) on the
+    // configsvr config.collections.
+    assert.eq(undefined, configDB.collections.findOne({_id: ns}).permitMigrations);
+};
+
+// Test cases that should disable the balancer.
+testBalancer(false /* setAllowMigrations */, {});
+testBalancer(false /* setAllowMigrations */, {noBalance: false});
+testBalancer(false /* setAllowMigrations */, {noBalance: true});
+
+// Test the setAllowMigrations command.
+testSetAllowMigrationsCommand();
+
+st.stop();
+})();