diff options
Diffstat (limited to 'src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp')
| -rw-r--r-- | src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp | 591 |
1 files changed, 23 insertions, 568 deletions
diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp index 6152927cfe9..a06171ed7da 100644 --- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp +++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_test.cpp @@ -33,11 +33,7 @@ #include "mongo/db/s/balancer/balancer_chunk_selection_policy_impl.h" #include "mongo/db/s/balancer/cluster_statistics_impl.h" #include "mongo/db/s/balancer/migration_test_fixture.h" -#include "mongo/idl/server_parameter_test_util.h" -#include "mongo/logv2/log.h" #include "mongo/platform/random.h" -#include "mongo/s/balancer_configuration.h" -#include "mongo/s/request_types/get_stats_for_balancing_gen.h" #include "mongo/s/type_collection_common_types_gen.h" namespace mongo { @@ -58,24 +54,6 @@ protected: std::make_unique<BalancerChunkSelectionPolicyImpl>(_clusterStats.get(), _random)) {} /** - * Generates a default chunks distribution across shards with the form: - * [MinKey, 0), [0, 1), [1, 2) ... [N - 2, MaxKey) - */ - std::map<ShardId, std::vector<ChunkRange>> generateDefaultChunkRanges( - const std::vector<ShardId>& shards) { - - std::map<ShardId, std::vector<ChunkRange>> chunksPerShard; - for (auto i = 0U; i < shards.size(); ++i) { - const ShardId& shardId = shards[i]; - const auto min = (i == 0 ? kKeyPattern.globalMin() : BSON(kPattern << int(i - 1))); - const auto max = - (i == shards.size() - 1 ? kKeyPattern.globalMax() : BSON(kPattern << int(i))); - chunksPerShard[shardId].push_back(ChunkRange(min, max)); - } - return chunksPerShard; - } - - /** * Sets up mock network to expect a listDatabases command and returns a BSON response with * a dummy sizeOnDisk. */ @@ -97,162 +75,33 @@ protected: } /** - * Sets up mock network for all the shards to expect the commands executed for computing cluster - * stats, which include listDatabase and serverStatus. - */ - void expectGetStatsCommands(int numShards) { - for (int i = 0; i < numShards; i++) { - expectListDatabasesCommand(); - } - } - - /** - * Sets up mock network for all the shards to expect the command `_shardsvrGetStatsForBalancing` - * Given a request sent to a specific shard with below structure ... - * { - * "_shardsvrGetStatsForBalancing" : 1, - * "collections" : [ - * { - * "ns" : "TestDb.TestColl", - * "UUID" : "xxxx" - * }, - * ... - * ] - * } - * - * ... mocks a reply with the following structure: - * { - * "stats" : [ - * { - * "namespace" : "TestDb.TestColl", - * "collSize" : 12345, - * }, - * ... - * ] - * } + * Sets up mock network to expect a serverStatus command and returns a BSON response with + * a dummy version. */ - void expectGetStatsForBalancingCommands(const std::map<ShardId, int64_t>& collSizePerShard) { - const auto& numShards = collSizePerShard.size(); - for (auto i = 0U; i < numShards; ++i) { - BSONObjBuilder resultBuilder; - CommandHelpers::appendCommandStatusNoThrow(resultBuilder, Status::OK()); - - // Build a response for given request - onCommand([&](const RemoteCommandRequest& request) { - ASSERT(request.cmdObj[ShardsvrGetStatsForBalancing::kCommandName]); - - // Get `shardId` - ShardId shardId = getShardIdByHost(request.target); - resultBuilder.append("shardId", shardId); - - // Build `stats` array: [ {"namespace": <nss>, "collSize": <collSize>}, ...] - { - BSONArrayBuilder statsArrayBuilder(resultBuilder.subarrayStart("stats")); - - ASSERT_EQ(1, collSizePerShard.count(shardId)); - const auto& collSize = collSizePerShard.at(shardId); - - for (const auto& reqColl : - request.cmdObj[ShardsvrGetStatsForBalancing::kCollectionsFieldName] - .Array()) { - const auto nss = - NamespaceWithOptionalUUID::parse( - IDLParserErrorContext("BalancerChunkSelectionPolicyTest"), - reqColl.Obj()) - .getNs(); - - statsArrayBuilder.append(CollStatsForBalancing(nss, collSize).toBSON()); - } - } - return resultBuilder.obj(); - }); - } - } + void expectServerStatusCommand() { + BSONObjBuilder resultBuilder; + CommandHelpers::appendCommandStatusNoThrow(resultBuilder, Status::OK()); - /** - * Same as expectGetStatsForBalancingCommands with the difference that this function will expect - * only one migration between the specified shards - */ - void expectGetStatsForBalancingCommandsWithOneMigration(uint32_t numShards, - ShardId donorShardId, - ShardId recipientShardId) { - ASSERT_NE(donorShardId, recipientShardId); - - const auto maxChunkSizeBytes = - Grid::get(operationContext())->getBalancerConfiguration()->getMaxChunkSizeBytes(); - const auto defaultCollSizeOnShard = 2 * maxChunkSizeBytes; - const auto imbalancedCollSizeOnRecipient = maxChunkSizeBytes; - const auto imbalancedCollSizeOnDonor = 5 * maxChunkSizeBytes; - - for (auto i = 0U; i < numShards; ++i) { - BSONObjBuilder resultBuilder; - CommandHelpers::appendCommandStatusNoThrow(resultBuilder, Status::OK()); - - // Build a response for every given request - onCommand([&](const RemoteCommandRequest& request) { - ASSERT(request.cmdObj[ShardsvrGetStatsForBalancing::kCommandName]); - - // Get `shardId` - ShardId shardId = getShardIdByHost(request.target); - resultBuilder.append("shardId", shardId); - - // Build `stats` array: [ {"namespace": <nss>, "collSize": <collSize>}, ...] - { - bool firstColl = true; - BSONArrayBuilder statsArrayBuilder(resultBuilder.subarrayStart("stats")); - for (const auto& reqColl : - request.cmdObj[ShardsvrGetStatsForBalancing::kCollectionsFieldName] - .Array()) { - const auto nss = - NamespaceWithOptionalUUID::parse( - IDLParserErrorContext("BalancerChunkSelectionPolicyTest"), - reqColl.Obj()) - .getNs(); - - const auto collSize = [&]() { - if (firstColl && shardId == donorShardId) { - return imbalancedCollSizeOnDonor; - } else if (firstColl && shardId == recipientShardId) { - return imbalancedCollSizeOnRecipient; - } - return defaultCollSizeOnShard; - }(); - - statsArrayBuilder.append(CollStatsForBalancing(nss, collSize).toBSON()); - firstColl = false; - } - } - return resultBuilder.obj(); - }); - } + onCommand([&resultBuilder](const RemoteCommandRequest& request) { + ASSERT(request.cmdObj["serverStatus"]); + resultBuilder.append("version", "MONGO_VERSION"); + return resultBuilder.obj(); + }); } /** - * Sets up a collection and its chunks according to the given range distribution across - * shards + * Sets up mock network for all the shards to expect the commands executed for computing cluster + * stats, which include listDatabase and serverStatus. */ - UUID setUpCollectionWithChunks( - const NamespaceString& ns, - const std::map<ShardId, std::vector<ChunkRange>>& chunksPerShard) { - const UUID collUuid = UUID::gen(); - ChunkVersion version(2, 0, OID::gen(), Timestamp(42)); - - for (const auto& [shardId, chunkRanges] : chunksPerShard) { - for (const auto& chunkRange : chunkRanges) { - setUpChunk(collUuid, chunkRange.getMin(), chunkRange.getMax(), shardId, version); - version.incMinor(); - } - version.incMajor(); + void expectGetStatsCommands(int numShards) { + for (int i = 0; i < numShards; i++) { + expectListDatabasesCommand(); + expectServerStatusCommand(); } - - setUpCollection(ns, collUuid, version); - - return collUuid; } /** - * Returns a new BSON object with the zone encoded using the legacy field "tags" - * (to mimic the expected schema of config.shards) + * Returns a new BSON object with the tags appended. */ BSONObj appendTags(const BSONObj shardBSON, std::vector<std::string> tags) { BSONObjBuilder appendedShardBSON(shardBSON); @@ -268,21 +117,8 @@ protected: BalancerRandomSource _random; std::unique_ptr<ClusterStatistics> _clusterStats; std::unique_ptr<BalancerChunkSelectionPolicy> _chunkSelectionPolicy; - stdx::unordered_set<NamespaceString> _imbalancedCollectionsCache; }; -stdx::unordered_set<ShardId> getAllShardIds( - const std::vector<ClusterStatistics::ShardStatistics>& shardStats) { - stdx::unordered_set<ShardId> shards; - std::transform(shardStats.begin(), - shardStats.end(), - std::inserter(shards, shards.end()), - [](const ClusterStatistics::ShardStatistics& shardStaticstics) -> ShardId { - return shardStaticstics.shardId; - }); - return shards; -} - TEST_F(BalancerChunkSelectionTest, TagRangesOverlap) { // Set up two shards in the metadata. ASSERT_OK(catalogClient()->insertConfigDocument( @@ -337,8 +173,6 @@ TEST_F(BalancerChunkSelectionTest, TagRangesOverlap) { } TEST_F(BalancerChunkSelectionTest, TagRangeMaxNotAlignedWithChunkMax) { - RAIIServerParameterControllerForTest featureFlagBalanceAccordingToDataSize{ - "featureFlagBalanceAccordingToDataSize", false}; // Set up two shards in the metadata. ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), ShardType::ConfigNS, @@ -375,12 +209,9 @@ TEST_F(BalancerChunkSelectionTest, TagRangeMaxNotAlignedWithChunkMax) { shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - _imbalancedCollectionsCache.clear(); - auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); + stdx::unordered_set<ShardId> usedShards; + auto candidateChunksStatus = + _chunkSelectionPolicy.get()->selectChunksToMove(opCtx.get(), &usedShards); ASSERT_OK(candidateChunksStatus.getStatus()); // The balancer does not bubble up the IllegalOperation error, but it is expected @@ -450,8 +281,6 @@ TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeAutoSplitted } TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeBalanced) { - RAIIServerParameterControllerForTest featureFlagBalanceAccordingToDataSize{ - "featureFlagBalanceAccordingToDataSize", false}; // Set up two shards in the metadata. ASSERT_OK(catalogClient()->insertConfigDocument( operationContext(), ShardType::ConfigNS, kShard0, kMajorityWriteConcern)); @@ -487,13 +316,9 @@ TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeBalanced) { shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - _imbalancedCollectionsCache.clear(); - auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); + stdx::unordered_set<ShardId> usedShards; + auto candidateChunksStatus = + _chunkSelectionPolicy.get()->selectChunksToMove(opCtx.get(), &usedShards); ASSERT_OK(candidateChunksStatus.getStatus()); ASSERT_EQUALS(1, candidateChunksStatus.getValue().size()); @@ -503,375 +328,5 @@ TEST_F(BalancerChunkSelectionTest, ShardedTimeseriesCollectionsCanBeBalanced) { future.default_timed_get(); } -TEST_F(BalancerChunkSelectionTest, AllImbalancedCollectionsShouldEventuallyBeSelectedForBalancing) { - // Set up two shards in the metadata. - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard0, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard1, - kMajorityWriteConcern)); - - // Set up database - setUpDatabase(kDbName, kShardId0); - - // Override collections batch size to 4 for speeding up the test - FailPointEnableBlock overrideBatchSizeGuard("overrideStatsForBalancingBatchSize", - BSON("size" << 4)); - - // Set up 7 imbalanced collections (more than `kStatsForBalancingBatchSize`) - const int numCollections = 7; - const int maxIterations = 1000; - - for (auto i = 0; i < numCollections; ++i) { - const std::string collName = "TestColl" + std::to_string(i); - setUpCollectionWithChunks(NamespaceString(kDbName, collName), - generateDefaultChunkRanges({kShardId0, kShardId1})); - } - - std::set<NamespaceString> collectionsSelected; - _imbalancedCollectionsCache.clear(); - - auto i = 0; - for (; i < maxIterations; ++i) { - - auto future = launchAsync([this, &collectionsSelected]() { - ThreadClient tc(getServiceContext()); - auto opCtx = Client::getCurrent()->makeOperationContext(); - - // Requests chunks to be relocated requires running commands on each shard to - // get shard statistics. Set up dummy hosts for the source shards. - shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); - shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - auto chunksToMoveWithStatus = _chunkSelectionPolicy->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); - ASSERT_OK(chunksToMoveWithStatus.getStatus()); - - for (const auto& chunkToMove : chunksToMoveWithStatus.getValue()) { - collectionsSelected.insert(chunkToMove.nss); - } - }); - - expectGetStatsCommands(2 /*numShards*/); - - // Collection size distribution for each collection: - // Shard0 -> 512 MB - // Shard1 -> 0 MB - expectGetStatsForBalancingCommands( - {{kShardId0, 512 * 1024 * 1024 /*Bytes*/}, {kShardId1, 0 /*Bytes*/}}); - - future.default_timed_get(); - - if (collectionsSelected.size() == numCollections) { - break; - } - } - - LOGV2(6867000, - "AllImbalancedCollectionsShouldEventuallyBeSelectedForBalancing test results", - "numCollectionsSelected"_attr = collectionsSelected.size(), - "iterations"_attr = i); - - // Check that all collections were selected for balancing at least once. - ASSERT_EQ(numCollections, collectionsSelected.size()); -} - -TEST_F(BalancerChunkSelectionTest, SelectedCollectionsShouldBeCached) { - // Set up two shards in the metadata. - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard0, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard1, - kMajorityWriteConcern)); - - // Set up database - setUpDatabase(kDbName, kShardId0); - - // Set up 4 collections - const int numCollections = 4; - for (auto i = 0; i < numCollections; ++i) { - const std::string collName = "TestColl" + std::to_string(i); - setUpCollectionWithChunks(NamespaceString(kDbName, collName), - generateDefaultChunkRanges({kShardId0, kShardId1})); - } - - std::set<NamespaceString> collectionsSelected; - _imbalancedCollectionsCache.clear(); - - for (auto i = 0; i < 5; ++i) { - - auto future = launchAsync([this, &collectionsSelected]() { - ThreadClient tc(getServiceContext()); - auto opCtx = Client::getCurrent()->makeOperationContext(); - - // Requests chunks to be relocated requires running commands on each shard to - // get shard statistics. Set up dummy hosts for the source shards. - shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); - shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - auto chunksToMoveWithStatus = _chunkSelectionPolicy->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); - ASSERT_OK(chunksToMoveWithStatus.getStatus()); - - for (const auto& chunkToMove : chunksToMoveWithStatus.getValue()) { - collectionsSelected.insert(chunkToMove.nss); - } - }); - - expectGetStatsCommands(2 /*numShards*/); - - // Collection size distribution for each collection: - // Shard0 -> 512 MB - // Shard1 -> 0 MB - expectGetStatsForBalancingCommands( - {{kShardId0, 512 * 1024 * 1024 /*Bytes*/}, {kShardId1, 0 /*Bytes*/}}); - - future.default_timed_get(); - } - - // Check that all selected collections are cached - for (const auto& coll : collectionsSelected) { - ASSERT_TRUE(_imbalancedCollectionsCache.count(coll)); - } - ASSERT_EQ(_imbalancedCollectionsCache.size(), collectionsSelected.size()); -} - -TEST_F(BalancerChunkSelectionTest, CachedCollectionsShouldBeSelected) { - // Set up two shards in the metadata. - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard0, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard1, - kMajorityWriteConcern)); - - // Set up database - setUpDatabase(kDbName, kShardId0); - - _imbalancedCollectionsCache.clear(); - std::vector<NamespaceString> allCollections; - - // Set up 4 collections and add all them into the imbalanced collections cache - const int numCollections = 4; - for (auto i = 0; i < numCollections; ++i) { - NamespaceString nss(kDbName, "TestColl" + std::to_string(i)); - allCollections.push_back(nss); - setUpCollectionWithChunks(nss, generateDefaultChunkRanges({kShardId0, kShardId1})); - - _imbalancedCollectionsCache.insert(nss); - } - - std::set<NamespaceString> collectionsSelected; - - for (auto i = 0; i < 1000; ++i) { - - auto future = launchAsync([this, &collectionsSelected]() { - ThreadClient tc(getServiceContext()); - auto opCtx = Client::getCurrent()->makeOperationContext(); - - // Requests chunks to be relocated requires running commands on each shard to - // get shard statistics. Set up dummy hosts for the source shards. - shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); - shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - auto chunksToMoveWithStatus = _chunkSelectionPolicy->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); - ASSERT_OK(chunksToMoveWithStatus.getStatus()); - - for (const auto& chunkToMove : chunksToMoveWithStatus.getValue()) { - collectionsSelected.insert(chunkToMove.nss); - } - }); - - expectGetStatsCommands(2 /*numShards*/); - - // Collection size distribution for each collection: - // Shard0 -> 512 MB - // Shard1 -> 0 MB - expectGetStatsForBalancingCommands( - {{kShardId0, 512 * 1024 * 1024 /*Bytes*/}, {kShardId1, 0 /*Bytes*/}}); - - future.default_timed_get(); - - if (collectionsSelected.size() == allCollections.size()) { - break; - } - } - - // Check that all selected collections are cached - for (const auto& nss : allCollections) { - ASSERT_TRUE(collectionsSelected.count(nss)); - } - ASSERT_EQ(allCollections.size(), collectionsSelected.size()); -} - -TEST_F(BalancerChunkSelectionTest, MaxTimeToScheduleBalancingOperationsExceeded) { - // Set up 4 shards in the metadata. - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard0, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard1, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard2, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard3, - kMajorityWriteConcern)); - - // Set up database - setUpDatabase(kDbName, kShardId0); - - // Override collections batch size to 4 for speeding up the test - FailPointEnableBlock overrideBatchSizeGuard("overrideStatsForBalancingBatchSize", - BSON("size" << 4)); - - // Set up 5 collections to process more than 1 batch - for (auto i = 0U; i < 5; ++i) { - NamespaceString ns{kDbName, "coll" + std::to_string(i)}; - setUpCollectionWithChunks( - ns, generateDefaultChunkRanges({kShardId0, kShardId1, kShardId2, kShardId3})); - } - - auto future = launchAsync([&] { - ThreadClient tc(getServiceContext()); - auto opCtx = Client::getCurrent()->makeOperationContext(); - - // Requesting chunks to be relocated requires running commands on each shard to get - // shard statistics. Set up dummy hosts for the source shards. - shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); - shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - shardTargeterMock(opCtx.get(), kShardId2)->setFindHostReturnValue(kShardHost2); - shardTargeterMock(opCtx.get(), kShardId3)->setFindHostReturnValue(kShardHost3); - - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - _imbalancedCollectionsCache.clear(); - - // Forcing timeout to exceed by setting it to 0 - RAIIServerParameterControllerForTest balancerChunksSelectionTimeoutMsIsZero( - "balancerChunksSelectionTimeoutMs", 0); - - auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); - - ASSERT_OK(candidateChunksStatus.getStatus()); - - // We know that timeout exceeded because we only got 1 migration instead of the 2 migrations - // expected in a normal scenario with 4 shards - ASSERT_EQUALS(1U, candidateChunksStatus.getValue().size()); - }); - - expectGetStatsCommands(4); - - // We need to get at least 1 migration per batch since the timeout only exceeds when balancer - // has found at least one candidate migration On the other side, we must get less than 2 - // migrations per batch since the maximum number of migrations per balancing round is 2 (with 4 - // shards) - expectGetStatsForBalancingCommandsWithOneMigration( - 4 /*numShards*/, kShardId0 /*donor*/, kShardId1 /*recipient*/); - - future.default_timed_get(); -} - -TEST_F(BalancerChunkSelectionTest, MakeSureMoreThanOneBatchIsProcessedIfNeeded) { - // Set up 4 shards in the metadata. - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard0, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard1, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard2, - kMajorityWriteConcern)); - ASSERT_OK(catalogClient()->insertConfigDocument(operationContext(), - NamespaceString::kConfigsvrShardsNamespace, - kShard3, - kMajorityWriteConcern)); - - // Set up database - setUpDatabase(kDbName, kShardId0); - - // Override collections batch size to 4 for speeding up the test - FailPointEnableBlock overrideBatchSizeGuard("overrideStatsForBalancingBatchSize", - BSON("size" << 4)); - - // Set up 5 collections to process 2 batches - for (auto i = 0; i < 5; ++i) { - NamespaceString ns{kDbName, "coll" + std::to_string(i)}; - setUpCollectionWithChunks( - ns, generateDefaultChunkRanges({kShardId0, kShardId1, kShardId2, kShardId3})); - } - - auto future = launchAsync([&] { - ThreadClient tc(getServiceContext()); - auto opCtx = Client::getCurrent()->makeOperationContext(); - - // Requesting chunks to be relocated requires running commands on each shard to get - // shard statistics. Set up dummy hosts for the source shards. - shardTargeterMock(opCtx.get(), kShardId0)->setFindHostReturnValue(kShardHost0); - shardTargeterMock(opCtx.get(), kShardId1)->setFindHostReturnValue(kShardHost1); - shardTargeterMock(opCtx.get(), kShardId2)->setFindHostReturnValue(kShardHost2); - shardTargeterMock(opCtx.get(), kShardId3)->setFindHostReturnValue(kShardHost3); - - std::vector<ClusterStatistics::ShardStatistics> shardStats = - uassertStatusOK(_clusterStats.get()->getStats(opCtx.get())); - auto availableShards = getAllShardIds(shardStats); - - _imbalancedCollectionsCache.clear(); - - // Forcing timeout to exceed - auto candidateChunksStatus = _chunkSelectionPolicy.get()->selectChunksToMove( - opCtx.get(), shardStats, &availableShards, &_imbalancedCollectionsCache); - - ASSERT_OK(candidateChunksStatus.getStatus()); - - // We know that timeout exceeded because we only got 1 migration instead of the 2 migrations - // expected in a normal scenario with 4 shards - ASSERT_EQUALS(2U, candidateChunksStatus.getValue().size()); - }); - - expectGetStatsCommands(4); - - // We are scheduling one migration on the first batch to make sure that the second batch is - // processed - expectGetStatsForBalancingCommandsWithOneMigration( - 4 /*numShards*/, kShardId0 /*donor*/, kShardId1 /*recipient*/); - expectGetStatsForBalancingCommandsWithOneMigration( - 4 /*numShards*/, kShardId2 /*donor*/, kShardId3 /*recipient*/); - - future.default_timed_get(); -} } // namespace } // namespace mongo |