Update upstream source from tag 'upstream/6.0.0'master

Update to upstream version '6.0.0'
with Debian dir 5604a80ec1c96ca76f25f40d78e6ef855abec322

1 files changed, 69 insertions, 270 deletions

diff --git a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
index 8deef770d27..c9a41e33c09 100644
--- a/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_chunk_selection_policy_impl.cpp
@@ -40,7 +40,6 @@
 #include "mongo/base/status_with.h"
 #include "mongo/bson/bsonobj_comparator_interface.h"
 #include "mongo/db/s/sharding_config_server_parameters_gen.h"
-#include "mongo/db/s/sharding_util.h"
 #include "mongo/logv2/log.h"
 #include "mongo/platform/bits.h"
 #include "mongo/s/balancer_configuration.h"
@@ -48,14 +47,8 @@
 #include "mongo/s/catalog/type_collection.h"
 #include "mongo/s/catalog/type_tags.h"
 #include "mongo/s/catalog_cache.h"
-#include "mongo/s/chunk_manager.h"
 #include "mongo/s/grid.h"
-#include "mongo/s/request_types/get_stats_for_balancing_gen.h"
-#include "mongo/s/sharding_feature_flags_gen.h"
 #include "mongo/util/str.h"
-#include "mongo/util/timer.h"
-
-MONGO_FAIL_POINT_DEFINE(overrideStatsForBalancingBatchSize);
 
 namespace mongo {
 
@@ -70,92 +63,39 @@ StatusWith<DistributionStatus> createCollectionDistributionStatus(
     const NamespaceString& nss,
     const ShardStatisticsVector& allShards,
     const ChunkManager& chunkMgr) {
+    ShardToChunksMap shardToChunksMap;
 
-    auto swZoneInfo =
-        createCollectionZoneInfo(opCtx, nss, chunkMgr.getShardKeyPattern().getKeyPattern());
-    if (!swZoneInfo.isOK()) {
-        return swZoneInfo.getStatus();
+    // Makes sure there is an entry in shardToChunksMap for every shard, so empty shards will also
+    // be accounted for
+    for (const auto& stat : allShards) {
+        shardToChunksMap[stat.shardId];
     }
 
-    return {DistributionStatus{nss, std::move(swZoneInfo.getValue()), chunkMgr}};
-}
+    chunkMgr.forEachChunk([&](const auto& chunkEntry) {
+        ChunkType chunk;
+        chunk.setCollectionUUID(chunkMgr.getUUID());
+        chunk.setMin(chunkEntry.getMin());
+        chunk.setMax(chunkEntry.getMax());
+        chunk.setJumbo(chunkEntry.isJumbo());
+        chunk.setShard(chunkEntry.getShardId());
+        chunk.setVersion(chunkEntry.getLastmod());
 
-stdx::unordered_map<NamespaceString, CollectionDataSizeInfoForBalancing>
-getDataSizeInfoForCollections(OperationContext* opCtx,
-                              const std::vector<CollectionType>& collections) {
-    const auto balancerConfig = Grid::get(opCtx)->getBalancerConfiguration();
-    uassertStatusOK(balancerConfig->refreshAndCheck(opCtx));
+        shardToChunksMap[chunkEntry.getShardId()].push_back(chunk);
 
-    const auto shardRegistry = Grid::get(opCtx)->shardRegistry();
-    const auto shardIds = shardRegistry->getAllShardIds(opCtx);
+        return true;
+    });
 
-    // Map to be returned, incrementally populated with the collected statistics
-    stdx::unordered_map<NamespaceString, CollectionDataSizeInfoForBalancing> dataSizeInfoMap;
+    DistributionStatus distribution(nss, std::move(shardToChunksMap));
 
-    std::vector<NamespaceWithOptionalUUID> namespacesWithUUIDsForStatsRequest;
-    for (const auto& coll : collections) {
-        const auto& nss = coll.getNss();
-        const auto maxChunkSizeBytes =
-            coll.getMaxChunkSizeBytes().value_or(balancerConfig->getMaxChunkSizeBytes());
-
-        dataSizeInfoMap.emplace(
-            nss,
-            CollectionDataSizeInfoForBalancing(std::map<ShardId, int64_t>(), maxChunkSizeBytes));
-
-        NamespaceWithOptionalUUID nssWithUUID(nss);
-        nssWithUUID.setUUID(coll.getUuid());
-        namespacesWithUUIDsForStatsRequest.push_back(nssWithUUID);
-    }
-
-    ShardsvrGetStatsForBalancing req{namespacesWithUUIDsForStatsRequest};
-    req.setScaleFactor(1);
-    const auto reqObj = req.toBSON({});
-
-    const auto executor = Grid::get(opCtx)->getExecutorPool()->getFixedExecutor();
-    const auto responsesFromShards =
-        sharding_util::sendCommandToShards(opCtx,
-                                           NamespaceString::kAdminDb.toString(),
-                                           reqObj,
-                                           shardIds,
-                                           executor,
-                                           false /* throwOnError */);
-
-    for (auto&& response : responsesFromShards) {
-        try {
-            const auto& shardId = response.shardId;
-            const auto errorContext =
-                "Failed to get stats for balancing from shard '{}'"_format(shardId.toString());
-            const auto responseValue =
-                uassertStatusOKWithContext(std::move(response.swResponse), errorContext);
-
-            const ShardsvrGetStatsForBalancingReply reply =
-                ShardsvrGetStatsForBalancingReply::parse(
-                    IDLParserErrorContext("ShardsvrGetStatsForBalancingReply"),
-                    std::move(responseValue.data));
-            const auto collStatsFromShard = reply.getStats();
-
-            invariant(collStatsFromShard.size() == collections.size());
-            for (const auto& stats : collStatsFromShard) {
-                invariant(dataSizeInfoMap.contains(stats.getNs()));
-                dataSizeInfoMap.at(stats.getNs()).shardToDataSizeMap[shardId] = stats.getCollSize();
-            }
-        } catch (const ExceptionFor<ErrorCodes::ShardNotFound>& ex) {
-            // Handle `removeShard`: skip shards removed during a balancing round
-            LOGV2_DEBUG(6581603,
-                        1,
-                        "Skipping shard for the current balancing round",
-                        "error"_attr = redact(ex));
-        }
-    }
+    const auto& keyPattern = chunkMgr.getShardKeyPattern().getKeyPattern();
 
-    return dataSizeInfoMap;
-}
+    // Cache the collection tags
+    auto status = ZoneInfo::addTagsFromCatalog(opCtx, nss, keyPattern, distribution.zoneInfo());
+    if (!status.isOK()) {
+        return status;
+    }
 
-const CollectionDataSizeInfoForBalancing getDataSizeInfoForCollection(OperationContext* opCtx,
-                                                                      const NamespaceString& nss) {
-    const auto coll = Grid::get(opCtx)->catalogClient()->getCollection(opCtx, nss);
-    std::vector<CollectionType> vec{coll};
-    return std::move(getDataSizeInfoForCollections(opCtx, vec).at(nss));
+    return {std::move(distribution)};
 }
 
 /**
@@ -225,12 +165,12 @@ private:
  * range boundaries.
  */
 void getSplitCandidatesToEnforceTagRanges(const ChunkManager& cm,
-                                          const ZoneInfo& zoneInfo,
+                                          const DistributionStatus& distribution,
                                           SplitCandidatesBuffer* splitCandidates) {
     const auto& globalMax = cm.getShardKeyPattern().getKeyPattern().globalMax();
 
     // For each tag range, find chunks that need to be split.
-    for (const auto& tagRangeEntry : zoneInfo.zoneRanges()) {
+    for (const auto& tagRangeEntry : distribution.tagRanges()) {
         const auto& tagRange = tagRangeEntry.second;
 
         const auto chunkAtZoneMin = cm.findIntersectingChunkWithSimpleCollation(tagRange.min);
@@ -382,185 +322,60 @@ StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::selectChunksToSpli
 }
 
 StatusWith<MigrateInfoVector> BalancerChunkSelectionPolicyImpl::selectChunksToMove(
-    OperationContext* opCtx,
-    const std::vector<ClusterStatistics::ShardStatistics>& shardStats,
-    stdx::unordered_set<ShardId>* availableShards,
-    stdx::unordered_set<NamespaceString>* imbalancedCollectionsCachePtr) {
-    invariant(availableShards);
-    invariant(imbalancedCollectionsCachePtr);
-
-    if (availableShards->size() < 2) {
-        return MigrateInfoVector{};
+    OperationContext* opCtx, stdx::unordered_set<ShardId>* usedShards) {
+    auto shardStatsStatus = _clusterStats->getStats(opCtx);
+    if (!shardStatsStatus.isOK()) {
+        return shardStatsStatus.getStatus();
     }
 
-    Timer chunksSelectionTimer;
+    const auto& shardStats = shardStatsStatus.getValue();
 
-    auto collections = Grid::get(opCtx)->catalogClient()->getCollections(
-        opCtx,
-        {},
-        repl::ReadConcernLevel::kMajorityReadConcern,
-        BSON(CollectionType::kNssFieldName << 1));
+    if (shardStats.size() < 2) {
+        return MigrateInfoVector{};
+    }
 
+    auto collections = Grid::get(opCtx)->catalogClient()->getCollections(opCtx, {});
     if (collections.empty()) {
         return MigrateInfoVector{};
     }
 
     MigrateInfoVector candidateChunks;
 
-    const uint32_t kStatsForBalancingBatchSize = [&]() {
-        auto batchSize = 100U;
-        overrideStatsForBalancingBatchSize.execute([&batchSize](const BSONObj& data) {
-            batchSize = data["size"].numberInt();
-            LOGV2(7617200, "Overriding collections batch size", "size"_attr = batchSize);
-        });
-        return batchSize;
-    }();
-
-    const uint32_t kMaxCachedCollectionsSize = 0.75 * kStatsForBalancingBatchSize;
-
-    // Lambda function used to get a CollectionType leveraging the `collections` vector
-    // The `collections` vector must be sorted by nss when it is called
-    auto getCollectionTypeByNss = [&collections](const NamespaceString& nss)
-        -> std::pair<boost::optional<CollectionType>, std::vector<CollectionType>::iterator> {
-        // Using a lower_bound to perform a binary search on the `collections` vector
-        const auto collIt =
-            std::lower_bound(collections.begin(),
-                             collections.end(),
-                             nss,
-                             [](const CollectionType& coll, const NamespaceString& ns) {
-                                 return coll.getNss() < ns;
-                             });
-
-        if (collIt == collections.end() || collIt->getNss() != nss) {
-            return std::make_pair(boost::none, collections.end());
-        }
-        return std::make_pair(*collIt, collIt);
-    };
+    std::shuffle(collections.begin(), collections.end(), _random);
+
+    for (const auto& coll : collections) {
+        const NamespaceString& nss(coll.getNss());
 
-    // Lambda function to check if a collection is explicitly disabled for balancing
-    const auto canBalanceCollection = [](const CollectionType& coll) -> bool {
         if (!coll.getAllowBalance() || !coll.getAllowMigrations() || !coll.getPermitMigrations() ||
             coll.getDefragmentCollection()) {
             LOGV2_DEBUG(5966401,
                         1,
                         "Not balancing explicitly disabled collection",
-                        "namespace"_attr = coll.getNss(),
+                        "namespace"_attr = nss,
                         "allowBalance"_attr = coll.getAllowBalance(),
                         "allowMigrations"_attr = coll.getAllowMigrations(),
                         "permitMigrations"_attr = coll.getPermitMigrations(),
                         "defragmentCollection"_attr = coll.getDefragmentCollection());
-            return false;
-        }
-        return true;
-    };
-
-    // Lambda function to select migrate candidates from a batch of collections
-    const auto processBatch = [&](std::vector<CollectionType>& collBatch) {
-        boost::optional<stdx::unordered_map<NamespaceString, CollectionDataSizeInfoForBalancing>>
-            collsDataSizeInfo;
-        if (feature_flags::gBalanceAccordingToDataSize.isEnabled(
-                serverGlobalParams.featureCompatibility)) {
-            collsDataSizeInfo.emplace(getDataSizeInfoForCollections(opCtx, collBatch));
-        }
-
-        std::shuffle(collBatch.begin(), collBatch.end(), _random);
-        for (const auto& coll : collBatch) {
-
-            if (availableShards->size() < 2) {
-                break;
-            }
-
-            const auto& nss = coll.getNss();
-
-            boost::optional<CollectionDataSizeInfoForBalancing> optDataSizeInfo;
-            if (collsDataSizeInfo.has_value()) {
-                optDataSizeInfo.emplace(std::move(collsDataSizeInfo->at(nss)));
-            }
-
-            auto swMigrateCandidates = _getMigrateCandidatesForCollection(
-                opCtx, nss, shardStats, optDataSizeInfo, availableShards);
-            if (swMigrateCandidates == ErrorCodes::NamespaceNotFound) {
-                // Namespace got dropped before we managed to get to it, so just skip it
-                imbalancedCollectionsCachePtr->erase(nss);
-                continue;
-            } else if (!swMigrateCandidates.isOK()) {
-                LOGV2_WARNING(21853,
-                              "Unable to balance collection",
-                              "namespace"_attr = nss.ns(),
-                              "error"_attr = swMigrateCandidates.getStatus());
-                continue;
-            }
-
-            candidateChunks.insert(
-                candidateChunks.end(),
-                std::make_move_iterator(swMigrateCandidates.getValue().first.begin()),
-                std::make_move_iterator(swMigrateCandidates.getValue().first.end()));
-
-            const auto& migrateCandidates = swMigrateCandidates.getValue().first;
-            if (migrateCandidates.empty()) {
-                imbalancedCollectionsCachePtr->erase(nss);
-            } else if (imbalancedCollectionsCachePtr->size() < kMaxCachedCollectionsSize) {
-                imbalancedCollectionsCachePtr->insert(nss);
-            }
-        }
-    };
-
-    // To assess if a collection has chunks to migrate, we need to ask shards the size of that
-    // collection. For efficiency, we ask for a batch of collections per every shard request instead
-    // of a single request per collection
-    std::vector<CollectionType> collBatch;
-
-    // The first batch is partially filled by the imbalanced cached collections
-    for (auto imbalancedNssIt = imbalancedCollectionsCachePtr->begin();
-         imbalancedNssIt != imbalancedCollectionsCachePtr->end();) {
-
-        const auto& [imbalancedColl, collIt] = getCollectionTypeByNss(*imbalancedNssIt);
-
-        if (!imbalancedColl.has_value() || !canBalanceCollection(imbalancedColl.value())) {
-            // The collection was dropped or is no longer enabled for balancing.
-            imbalancedCollectionsCachePtr->erase(imbalancedNssIt++);
             continue;
         }
 
-        collBatch.push_back(imbalancedColl.value());
-        ++imbalancedNssIt;
-
-        // Remove the collection from the whole list of collections to avoid processing it twice
-        collections.erase(collIt);
-    }
-
-    // Iterate all the remaining collections randomly
-    std::shuffle(collections.begin(), collections.end(), _random);
-    for (const auto& coll : collections) {
-
-        if (canBalanceCollection(coll)) {
-            collBatch.push_back(coll);
-        }
-
-        if (collBatch.size() == kStatsForBalancingBatchSize) {
-            processBatch(collBatch);
-            if (availableShards->size() < 2) {
-                return candidateChunks;
-            }
-            collBatch.clear();
-        }
-
-        const auto maxTimeMs = balancerChunksSelectionTimeoutMs.load();
-        if (candidateChunks.size() > 0 && chunksSelectionTimer.millis() > maxTimeMs) {
-            LOGV2_DEBUG(
-                7100900,
-                1,
-                "Exceeded max time while searching for candidate chunks to migrate in this round.",
-                "maxTime"_attr = Milliseconds(maxTimeMs),
-                "chunksSelectionTime"_attr = chunksSelectionTimer.elapsed(),
-                "numCandidateChunks"_attr = candidateChunks.size());
-
-            return candidateChunks;
+        auto candidatesStatus =
+            _getMigrateCandidatesForCollection(opCtx, nss, shardStats, usedShards);
+        if (candidatesStatus == ErrorCodes::NamespaceNotFound) {
+            // Namespace got dropped before we managed to get to it, so just skip it
+            continue;
+        } else if (!candidatesStatus.isOK()) {
+            LOGV2_WARNING(21853,
+                          "Unable to balance collection {namespace}: {error}",
+                          "Unable to balance collection",
+                          "namespace"_attr = nss.ns(),
+                          "error"_attr = candidatesStatus.getStatus());
+            continue;
         }
-    }
 
-    if (collBatch.size() > 0) {
-        processBatch(collBatch);
+        candidateChunks.insert(candidateChunks.end(),
+                               std::make_move_iterator(candidatesStatus.getValue().first.begin()),
+                               std::make_move_iterator(candidatesStatus.getValue().first.end()));
     }
 
     return candidateChunks;
@@ -579,23 +394,9 @@ StatusWith<MigrateInfosWithReason> BalancerChunkSelectionPolicyImpl::selectChunk
     // doesn't.
     Grid::get(opCtx)->catalogClient()->getCollection(opCtx, nss);
 
-    stdx::unordered_set<ShardId> availableShards;
-    std::transform(shardStats.begin(),
-                   shardStats.end(),
-                   std::inserter(availableShards, availableShards.end()),
-                   [](const ClusterStatistics::ShardStatistics& shardStatistics) -> ShardId {
-                       return shardStatistics.shardId;
-                   });
+    stdx::unordered_set<ShardId> usedShards;
 
-
-    boost::optional<CollectionDataSizeInfoForBalancing> optCollDataSizeInfo;
-    if (feature_flags::gBalanceAccordingToDataSize.isEnabled(
-            serverGlobalParams.featureCompatibility)) {
-        optCollDataSizeInfo.emplace(getDataSizeInfoForCollection(opCtx, nss));
-    }
-
-    auto candidatesStatus = _getMigrateCandidatesForCollection(
-        opCtx, nss, shardStats, optCollDataSizeInfo, &availableShards);
+    auto candidatesStatus = _getMigrateCandidatesForCollection(opCtx, nss, shardStats, &usedShards);
     if (!candidatesStatus.isOK()) {
         return candidatesStatus.getStatus();
     }
@@ -675,7 +476,7 @@ Status BalancerChunkSelectionPolicyImpl::checkMoveAllowed(OperationContext* opCt
     }
 
     return BalancerPolicy::isShardSuitableReceiver(*newShardIterator,
-                                                   distribution.getTagForRange(chunk.getRange()));
+                                                   distribution.getTagForChunk(chunk));
 }
 
 StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::_getSplitCandidatesForCollection(
@@ -688,26 +489,26 @@ StatusWith<SplitInfoVector> BalancerChunkSelectionPolicyImpl::_getSplitCandidate
 
     const auto& cm = routingInfoStatus.getValue();
 
-    const auto swCollZoneInfo =
-        createCollectionZoneInfo(opCtx, nss, cm.getShardKeyPattern().getKeyPattern());
-    if (!swCollZoneInfo.isOK()) {
-        return swCollZoneInfo.getStatus();
+    const auto collInfoStatus = createCollectionDistributionStatus(opCtx, nss, shardStats, cm);
+    if (!collInfoStatus.isOK()) {
+        return collInfoStatus.getStatus();
     }
-    const auto& collZoneInfo = swCollZoneInfo.getValue();
+
+    const DistributionStatus& distribution = collInfoStatus.getValue();
 
     // Accumulate split points for the same chunk together
     SplitCandidatesBuffer splitCandidates(nss, cm.getVersion());
 
     if (nss == NamespaceString::kLogicalSessionsNamespace) {
-        if (!collZoneInfo.allZones().empty()) {
+        if (!distribution.tags().empty()) {
             LOGV2_WARNING(4562401,
                           "Ignoring zones for the sessions collection",
-                          "tags"_attr = collZoneInfo.allZones());
+                          "tags"_attr = distribution.tags());
         }
 
         getSplitCandidatesForSessionsCollection(opCtx, cm, &splitCandidates);
     } else {
-        getSplitCandidatesToEnforceTagRanges(cm, collZoneInfo, &splitCandidates);
+        getSplitCandidatesToEnforceTagRanges(cm, distribution, &splitCandidates);
     }
 
     return splitCandidates.done();
@@ -718,8 +519,7 @@ BalancerChunkSelectionPolicyImpl::_getMigrateCandidatesForCollection(
     OperationContext* opCtx,
     const NamespaceString& nss,
     const ShardStatisticsVector& shardStats,
-    const boost::optional<CollectionDataSizeInfoForBalancing>& collDataSizeInfo,
-    stdx::unordered_set<ShardId>* availableShards) {
+    stdx::unordered_set<ShardId>* usedShards) {
     auto routingInfoStatus =
         Grid::get(opCtx)->catalogCache()->getShardedCollectionRoutingInfoWithRefresh(opCtx, nss);
     if (!routingInfoStatus.isOK()) {
@@ -737,7 +537,7 @@ BalancerChunkSelectionPolicyImpl::_getMigrateCandidatesForCollection(
 
     const DistributionStatus& distribution = collInfoStatus.getValue();
 
-    for (const auto& tagRangeEntry : distribution.getZoneInfo().zoneRanges()) {
+    for (const auto& tagRangeEntry : distribution.tagRanges()) {
         const auto& tagRange = tagRangeEntry.second;
 
         const auto chunkAtZoneMin = cm.findIntersectingChunkWithSimpleCollation(tagRange.min);
@@ -775,8 +575,7 @@ BalancerChunkSelectionPolicyImpl::_getMigrateCandidatesForCollection(
     return BalancerPolicy::balance(
         shardStats,
         distribution,
-        collDataSizeInfo,
-        availableShards,
+        usedShards,
         Grid::get(opCtx)->getBalancerConfiguration()->attemptToBalanceJumboChunks());
 }