mongo/jstests/sharding/libs/defragmentation_util.js

import {findChunksUtil} from "jstests/sharding/libs/find_chunks_util.js";
export var defragmentationUtil = (function () {
    let createFragmentedCollection = function (
        mongos,
        ns,
        numChunks,
        maxChunkFillMB,
        numZones,
        docSizeBytesRange,
        chunkSpacing,
        disableCollectionBalancing,
    ) {
        jsTest.log(
            "Creating fragmented collection " +
                ns +
                " with parameters: numChunks = " +
                numChunks +
                ", numZones = " +
                numZones +
                ", docSizeBytesRange = " +
                docSizeBytesRange +
                ", maxChunkFillMB = " +
                maxChunkFillMB +
                ", chunkSpacing = " +
                chunkSpacing,
        );
        assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: {key: 1}}));
        // Turn off balancer for this collection
        if (disableCollectionBalancing) {
            // Use 'retryWrites' when writing to the configsvr because mongos does not automatically
            // retry those.
            const mongosSession = mongos.startSession({retryWrites: true});
            const sessionConfigDB = mongosSession.getDatabase("config");
            assert.commandWorked(sessionConfigDB.collections.update({_id: ns}, {$set: {"noBalance": true}}));
        }

        createAndDistributeChunks(mongos, ns, numChunks, chunkSpacing);
        // Created zones will line up exactly with existing chunks so as not to trigger zone
        // violations in the balancer.
        createRandomZones(mongos, ns, numZones);
        fillChunksToRandomSize(mongos, ns, docSizeBytesRange, maxChunkFillMB);

        const beginningNumberChunks = findChunksUtil.countChunksForNs(mongos.getDB("config"), ns);
        const beginningNumberZones = mongos.getDB("config").tags.countDocuments({ns: ns});
        jsTest.log(
            "Collection " +
                ns +
                " created with " +
                beginningNumberChunks +
                " chunks and " +
                beginningNumberZones +
                " zones.",
        );
    };

    let createAndDistributeChunks = function (mongos, ns, numChunks, chunkSpacing) {
        const shards = mongos.getCollection("config.shards").find().toArray();
        const existingNumChunks = findChunksUtil.countChunksForNs(mongos.getDB("config"), ns);
        let numChunksToCreate = numChunks - existingNumChunks;
        if (numChunksToCreate <= 0) {
            return;
        }
        for (let i = -Math.floor(numChunksToCreate / 2); i < Math.ceil(numChunksToCreate / 2); i++) {
            assert.commandWorked(mongos.adminCommand({split: ns, middle: {key: i * chunkSpacing}}));
            assert.soon(() => {
                let toShard = Random.randInt(shards.length);
                let res = mongos.adminCommand({moveChunk: ns, find: {key: i * chunkSpacing}, to: shards[toShard]._id});
                return res.ok;
            });
        }
    };

    let createRandomZones = function (mongos, ns, numZones) {
        let existingChunks = findChunksUtil.findChunksByNs(mongos.getDB("config"), ns);
        existingChunks = Array.shuffle(existingChunks.toArray());
        for (let i = 0; i < numZones; i++) {
            let zoneName = "Zone" + i;
            let shardForZone = existingChunks[i].shard;
            assert.commandWorked(mongos.adminCommand({addShardToZone: shardForZone, zone: zoneName}));
            assert.commandWorked(
                mongos.adminCommand({
                    updateZoneKeyRange: ns,
                    min: existingChunks[i].min,
                    max: existingChunks[i].max,
                    zone: zoneName,
                }),
            );
        }
    };

    let fillChunksToRandomSize = function (mongos, ns, docSizeBytesRange, maxChunkFillMB) {
        const chunks = findChunksUtil.findChunksByNs(mongos.getDB("config"), ns).toArray();
        const coll = mongos.getCollection(ns);
        let bulk = coll.initializeUnorderedBulkOp();
        assert.gte(docSizeBytesRange[1], docSizeBytesRange[0]);
        chunks.forEach((chunk) => {
            let chunkSize = Random.randInt(maxChunkFillMB);
            let docSizeBytes = Random.randInt(docSizeBytesRange[1] - docSizeBytesRange[0] + 1) + docSizeBytesRange[0];
            docSizeBytes = docSizeBytes === 0 ? 1 : docSizeBytes;
            let bigString = "X".repeat(docSizeBytes);
            let docsPerChunk = (chunkSize * 1024 * 1024) / docSizeBytes;
            if (docsPerChunk === 0) {
                return;
            }
            let minKey = chunk["min"]["key"];
            if (minKey === MinKey) minKey = Number.MIN_SAFE_INTEGER;
            let maxKey = chunk["max"]["key"];
            if (maxKey === MaxKey) maxKey = Number.MAX_SAFE_INTEGER;
            let currKey = minKey;
            const gap = (maxKey - minKey) / docsPerChunk;
            for (let i = 0; i < docsPerChunk; i++) {
                bulk.insert({key: currKey, longString: bigString});
                currKey += gap;
            }
        });
        assert.commandWorked(bulk.execute());
    };

    let checkPostDefragmentationState = function (configSvr, mongos, ns, maxChunkSizeMB, shardKey) {
        const oversizedChunkThreshold = (maxChunkSizeMB * 1024 * 1024 * 4) / 3;
        const chunks = findChunksUtil
            .findChunksByNs(mongos.getDB("config"), ns)
            .sort({[shardKey]: 1})
            .toArray();
        const coll = mongos.getCollection(ns);
        const pipeline = [
            {"$collStats": {"storageStats": {}}},
            {"$project": {"shard": true, "storageStats": {"avgObjSize": true}}},
        ];
        const storageStats = coll.aggregate(pipeline).toArray();
        let avgObjSizeByShard = {};
        storageStats.forEach((storageStat) => {
            avgObjSizeByShard[storageStat["shard"]] =
                typeof storageStat["storageStats"]["avgObjSize"] === "undefined"
                    ? 0
                    : storageStat["storageStats"]["avgObjSize"];
        });
        for (let i = 1; i < chunks.length; i++) {
            let leftChunk = chunks[i - 1];
            let rightChunk = chunks[i];
            // Check for mergeable chunks with combined size less than maxChunkSize
            if (
                leftChunk["shard"] === rightChunk["shard"] &&
                bsonWoCompare(leftChunk["max"], rightChunk["min"]) === 0
            ) {
                let leftChunkZone = getZoneForRange(mongos, ns, leftChunk.min, leftChunk.max);
                let rightChunkZone = getZoneForRange(mongos, ns, rightChunk.min, rightChunk.max);
                if (bsonWoCompare(leftChunkZone, rightChunkZone) === 0) {
                    assert(false, `Chunks ${tojson(leftChunk)} and ${tojson(rightChunk)} should have been merged`);
                }
            }
        }
    };

    let getZoneForRange = function (mongos, ns, minKey, maxKey) {
        const tags = mongos
            .getDB("config")
            .tags.find({ns: ns, min: {$lte: minKey}, max: {$gte: maxKey}})
            .toArray();
        assert.lte(tags.length, 1);
        if (tags.length === 1) {
            return tags[0].tag;
        }
        return null;
    };

    let waitForEndOfDefragmentation = function (mongos, ns) {
        jsTest.log("Waiting end of defragmentation for " + ns);
        assert.soon(function () {
            let balancerStatus = assert.commandWorked(mongos.adminCommand({balancerCollectionStatus: ns}));
            return (
                balancerStatus.balancerCompliant || balancerStatus.firstComplianceViolation !== "defragmentingChunks"
            );
        });
        jsTest.log("Defragmentation completed for " + ns);
    };

    return {
        createFragmentedCollection,
        checkPostDefragmentationState,
        getZoneForRange,
        waitForEndOfDefragmentation,
    };
})();