SERVER-67301 Balancer may perform one unnecessary migration for a completely balanced collection

2 files changed, 114 insertions, 14 deletions

diff --git a/src/mongo/shell/shardingtest.js b/src/mongo/shell/shardingtest.js
index 7117f7575ac..79b20bc08cc 100644
--- a/src/mongo/shell/shardingtest.js
+++ b/src/mongo/shell/shardingtest.js
@@ -600,21 +600,12 @@ var ShardingTest = function(params) {
     };
 
     /**
-     * Waits up to one minute for the difference in chunks between the most loaded shard and
-     * least loaded shard to be 0 or 1, indicating that the collection is well balanced. This should
-     * only be called after creating a big enough chunk difference to trigger balancing.
-     */
+     * Waits up to the specified timeout (with a default of 60s) for the collection to be
+     * considered well balanced.
+     **/
     this.awaitBalance = function(collName, dbName, timeToWait, interval) {
-        timeToWait = timeToWait || 60000;
-        interval = interval || 200;
-
-        const mongos = this.s;
-        assert.soon(function() {
-            return assert
-                .commandWorked(
-                    mongos.adminCommand({balancerCollectionStatus: dbName + '.' + collName}))
-                .balancerCompliant;
-        }, 'Timed out waiting for the collection to be balanced', timeToWait, interval);
+        const coll = this.s.getCollection(dbName + "." + collName);
+        this.awaitCollectionBalance(coll, timeToWait, interval);
     };
 
     this.getShard = function(coll, query, includeEmpty) {
diff --git a/src/mongo/shell/utils_sh.js b/src/mongo/shell/utils_sh.js
index 3bb6db08d49..fc7574c5cb1 100644
--- a/src/mongo/shell/utils_sh.js
+++ b/src/mongo/shell/utils_sh.js
@@ -96,6 +96,9 @@ sh.help = function() {
         "returns wheter the specified collection is balanced or the balancer needs to take more actions on it");
     print("\tsh.configureCollectionBalancing(fullName, params)       " +
           "configure balancing settings for a specific collection");
+    print("\tsh.awaitCollectionBalance(coll)         waits for a collection to be balanced");
+    print(
+        "\tsh.verifyCollectionIsBalanced(coll)       verifies that a collection is well balanced by checking the actual data size on each shard");
 };
 
 sh.status = function(verbose, configDB) {
@@ -314,6 +317,112 @@ sh.enableBalancing = function(coll) {
         {writeConcern: {w: 'majority', wtimeout: 60000}}));
 };
 
+sh.awaitCollectionBalance = function(coll, timeout, interval) {
+    if (coll === undefined) {
+        throw Error("Must specify collection");
+    }
+    timeout = timeout || 60000;
+    interval = interval || 200;
+
+    const ns = coll.getFullName();
+    const orphanDocsPipeline = [
+        {'$collStats': {'storageStats': {}}},
+        {'$project': {'shard': true, 'storageStats': {'numOrphanDocs': true}}},
+        {'$group': {'_id': null, 'totalNumOrphanDocs': {'$sum': '$storageStats.numOrphanDocs'}}}
+    ];
+
+    var oldDb = db;
+    try {
+        db = coll.getDB();
+
+        assert.soon(
+            function() {
+                assert.soon(function() {
+                    return assert
+                        .commandWorked(sh._adminCommand({balancerCollectionStatus: ns}, true))
+                        .balancerCompliant;
+                }, 'Timed out waiting for the collection to be balanced', timeout, interval);
+
+                // (SERVER-67301) Wait for orphans counter to be 0 to account for potential stale
+                // orphans count
+                sh.disableBalancing(coll);
+                assert.soon(function() {
+                    return coll.aggregate(orphanDocsPipeline).toArray()[0].totalNumOrphanDocs === 0;
+                }, 'Timed out waiting for orphans counter to be 0', timeout, interval);
+                sh.enableBalancing(coll);
+
+                return assert.commandWorked(sh._adminCommand({balancerCollectionStatus: ns}, true))
+                    .balancerCompliant;
+            },
+            'Timed out waiting for collection to be balanced and orphans counter to be 0',
+            timeout,
+            interval);
+    } finally {
+        db = oldDb;
+    }
+};
+
+/**
+ * Verifies if given collection is properly balanced according to the data size aware balancing
+ * policy
+ */
+sh.verifyCollectionIsBalanced = function(coll) {
+    if (coll === undefined) {
+        throw Error("Must specify collection");
+    }
+
+    var oldDb = db;
+    try {
+        db = coll.getDB();
+
+        const configDB = sh._getConfigDB();
+        const ns = coll.getFullName();
+        const collection = configDB.collections.findOne({_id: ns});
+
+        let collSizeOnShards = [];
+        let shards = [];
+        const collStatsPipeline = [
+            {'$collStats': {'storageStats': {}}},
+            {
+                '$project': {
+                    'shard': true,
+                    'storageStats':
+                        {'count': true, 'size': true, 'avgObjSize': true, 'numOrphanDocs': true}
+                }
+            },
+            {'$sort': {'shard': 1}}
+        ];
+
+        let kChunkSize = 1024 * 1024 *
+            assert.commandWorked(sh._adminCommand({balancerCollectionStatus: ns})).chunkSize;
+        // TODO SERVER-67898 delete kChunkSize overwrite after completing the ticket
+        if (kChunkSize == 0) {
+            kChunkSize = collection.maxChunkSizeBytes;
+        }
+
+        // Get coll size per shard
+        const storageStats = coll.aggregate(collStatsPipeline).toArray();
+        coll.aggregate(collStatsPipeline).forEach((shardStats) => {
+            shards.push(shardStats['shard']);
+            const collSize = (shardStats['storageStats']['count'] -
+                              shardStats['storageStats']['numOrphanDocs']) *
+                shardStats['storageStats']['avgObjSize'];
+            collSizeOnShards.push(collSize);
+        });
+
+        let errorMsg = "Collection not balanced. collection= " + tojson(collection) +
+            ", shards= " + tojson(shards) + ", collSizeOnShards=" + tojson(collSizeOnShards) +
+            ", storageStats=" + tojson(storageStats) + ", kChunkSize=" + tojson(kChunkSize);
+
+        assert.lte((Math.max(...collSizeOnShards) - Math.min(...collSizeOnShards)),
+                   3 * kChunkSize,
+                   errorMsg);
+
+    } finally {
+        db = oldDb;
+    }
+};
+
 /*
  * Can call _lastMigration( coll ), _lastMigration( db ), _lastMigration( st ), _lastMigration(
  * mongos )