SERVER-60507 ClusterStatistics need to provide collection level size statistics

6 files changed, 124 insertions, 16 deletions

diff --git a/src/mongo/db/s/balancer/cluster_statistics.cpp b/src/mongo/db/s/balancer/cluster_statistics.cpp
index df276f8baf3..b9fffb0af51 100644
--- a/src/mongo/db/s/balancer/cluster_statistics.cpp
+++ b/src/mongo/db/s/balancer/cluster_statistics.cpp
@@ -34,6 +34,15 @@
 #include "mongo/bson/bsonobj.h"
 #include "mongo/bson/bsonobjbuilder.h"
 
+namespace {
+uint64_t convertMBToBytes(uint64_t inMB) {
+    if (inMB > std::numeric_limits<uint64_t>::max() / (1024 * 1024)) {
+        return std::numeric_limits<uint64_t>::max();
+    }
+    return inMB * 1024 * 1024;
+}
+}  // namespace
+
 namespace mongo {
 
 ClusterStatistics::ClusterStatistics() = default;
@@ -41,31 +50,46 @@ ClusterStatistics::ClusterStatistics() = default;
 ClusterStatistics::~ClusterStatistics() = default;
 
 ClusterStatistics::ShardStatistics::ShardStatistics(ShardId inShardId,
-                                                    uint64_t inMaxSizeMB,
-                                                    uint64_t inCurrSizeMB,
+                                                    uint64_t inMaxSizeBytes,
+                                                    uint64_t inCurrSizeBytes,
                                                     bool inIsDraining,
                                                     std::set<std::string> inShardTags,
-                                                    std::string inMongoVersion)
+                                                    std::string inMongoVersion,
+                                                    use_bytes_t t)
     : shardId(std::move(inShardId)),
-      maxSizeMB(inMaxSizeMB),
-      currSizeMB(inCurrSizeMB),
+      maxSizeBytes(inMaxSizeBytes),
+      currSizeBytes(inCurrSizeBytes),
       isDraining(inIsDraining),
       shardTags(std::move(inShardTags)),
       mongoVersion(std::move(inMongoVersion)) {}
 
+ClusterStatistics::ShardStatistics::ShardStatistics(ShardId inShardId,
+                                                    uint64_t inMaxSizeMB,
+                                                    uint64_t inCurrSizeMB,
+                                                    bool inIsDraining,
+                                                    std::set<std::string> inShardTags,
+                                                    std::string inMongoVersion)
+    : ShardStatistics(inShardId,
+                      convertMBToBytes(inMaxSizeMB),
+                      convertMBToBytes(inCurrSizeMB),
+                      inIsDraining,
+                      std::move(inShardTags),
+                      std::move(inMongoVersion),
+                      use_bytes_t{}) {}
+
 bool ClusterStatistics::ShardStatistics::isSizeMaxed() const {
-    if (!maxSizeMB || !currSizeMB) {
+    if (!maxSizeBytes || !currSizeBytes) {
         return false;
     }
 
-    return currSizeMB >= maxSizeMB;
+    return currSizeBytes >= maxSizeBytes;
 }
 
 BSONObj ClusterStatistics::ShardStatistics::toBSON() const {
     BSONObjBuilder builder;
     builder.append("id", shardId.toString());
-    builder.append("maxSizeMB", static_cast<long long>(maxSizeMB));
-    builder.append("currSizeMB", static_cast<long long>(currSizeMB));
+    builder.append("maxSizeMB", static_cast<long long>(maxSizeBytes / 1024 / 1024));
+    builder.append("currSizeMB", static_cast<long long>(currSizeBytes / 1024 / 1024));
     builder.append("draining", isDraining);
     if (!shardTags.empty()) {
         BSONArrayBuilder arrayBuilder(builder.subarrayStart("tags"));
diff --git a/src/mongo/db/s/balancer/cluster_statistics.h b/src/mongo/db/s/balancer/cluster_statistics.h
index 32f6cb9f7c5..71c7ce9484f 100644
--- a/src/mongo/db/s/balancer/cluster_statistics.h
+++ b/src/mongo/db/s/balancer/cluster_statistics.h
@@ -58,6 +58,18 @@ public:
      */
     struct ShardStatistics {
     public:
+        // hack so we do not accidentally miss code using MB
+        struct use_bytes_t {
+            explicit use_bytes_t() = default;
+        };
+        ShardStatistics(ShardId shardId,
+                        uint64_t maxSizeBytes,
+                        uint64_t currSizeBytes,
+                        bool isDraining,
+                        std::set<std::string> shardTags,
+                        std::string mongoVersion,
+                        use_bytes_t t);
+
         ShardStatistics(ShardId shardId,
                         uint64_t maxSizeMB,
                         uint64_t currSizeMB,
@@ -80,10 +92,10 @@ public:
         ShardId shardId;
 
         // The maximum storage size allowed for the shard. Zero means no maximum specified.
-        uint64_t maxSizeMB{0};
+        uint64_t maxSizeBytes{0};
 
         // The current storage size of the shard.
-        uint64_t currSizeMB{0};
+        uint64_t currSizeBytes{0};
 
         // Whether the shard is in draining mode
         bool isDraining{false};
@@ -103,6 +115,9 @@ public:
      */
     virtual StatusWith<std::vector<ShardStatistics>> getStats(OperationContext* opCtx) = 0;
 
+    virtual StatusWith<std::vector<ShardStatistics>> getCollStats(OperationContext* opCtx,
+                                                                  NamespaceString const& ns) = 0;
+
 protected:
     ClusterStatistics();
 };
diff --git a/src/mongo/db/s/balancer/cluster_statistics_impl.cpp b/src/mongo/db/s/balancer/cluster_statistics_impl.cpp
index 289f6517de6..495e8a77774 100644
--- a/src/mongo/db/s/balancer/cluster_statistics_impl.cpp
+++ b/src/mongo/db/s/balancer/cluster_statistics_impl.cpp
@@ -89,7 +89,6 @@ StatusWith<std::string> retrieveShardMongoDVersion(OperationContext* opCtx, Shar
 
     return version;
 }
-
 }  // namespace
 
 using ShardStatistics = ClusterStatistics::ShardStatistics;
@@ -99,6 +98,16 @@ ClusterStatisticsImpl::ClusterStatisticsImpl(BalancerRandomSource& random) : _ra
 ClusterStatisticsImpl::~ClusterStatisticsImpl() = default;
 
 StatusWith<std::vector<ShardStatistics>> ClusterStatisticsImpl::getStats(OperationContext* opCtx) {
+    return _getStats(opCtx, boost::none);
+}
+
+StatusWith<std::vector<ShardStatistics>> ClusterStatisticsImpl::getCollStats(
+    OperationContext* opCtx, NamespaceString const& ns) {
+    return _getStats(opCtx, ns);
+}
+
+StatusWith<std::vector<ShardStatistics>> ClusterStatisticsImpl::_getStats(
+    OperationContext* opCtx, boost::optional<NamespaceString> ns) {
     // Get a list of all the shards that are participating in this balance round along with any
     // maximum allowed quotas and current utilization. We get the latter by issuing
     // db.serverStatus() (mem.mapped) to all shards.
@@ -118,10 +127,13 @@ StatusWith<std::vector<ShardStatistics>> ClusterStatisticsImpl::getStats(Operati
 
     for (const auto& shard : shards) {
         const auto shardSizeStatus = [&]() -> StatusWith<long long> {
+            if (ns) {
+                return shardutil::retrieveCollectionShardSize(opCtx, shard.getName(), *ns);
+            }
+            // optimization for the case where the balancer does not care about the dataSize
             if (!shard.getMaxSizeMB()) {
                 return 0;
             }
-
             return shardutil::retrieveTotalShardSize(opCtx, shard.getName());
         }();
 
@@ -155,11 +167,12 @@ StatusWith<std::vector<ShardStatistics>> ClusterStatisticsImpl::getStats(Operati
         }
 
         stats.emplace_back(shard.getName(),
-                           shard.getMaxSizeMB(),
-                           shardSizeStatus.getValue() / 1024 / 1024,
+                           shard.getMaxSizeMB() * 1024 * 1024,
+                           shardSizeStatus.getValue(),
                            shard.getDraining(),
                            std::move(shardTags),
-                           std::move(mongoDVersion));
+                           std::move(mongoDVersion),
+                           ShardStatistics::use_bytes_t{});
     }
 
     return stats;
diff --git a/src/mongo/db/s/balancer/cluster_statistics_impl.h b/src/mongo/db/s/balancer/cluster_statistics_impl.h
index b9d5a1e2830..56037628a47 100644
--- a/src/mongo/db/s/balancer/cluster_statistics_impl.h
+++ b/src/mongo/db/s/balancer/cluster_statistics_impl.h
@@ -46,7 +46,13 @@ public:
 
     StatusWith<std::vector<ShardStatistics>> getStats(OperationContext* opCtx) override;
 
+    StatusWith<std::vector<ShardStatistics>> getCollStats(OperationContext* opCtx,
+                                                          NamespaceString const& ns) override;
+
 private:
+    StatusWith<std::vector<ShardStatistics>> _getStats(OperationContext* opCtx,
+                                                       boost::optional<NamespaceString> ns);
+
     // Source of randomness when metadata needs to be randomized.
     BalancerRandomSource& _random;
 };
diff --git a/src/mongo/s/shard_util.cpp b/src/mongo/s/shard_util.cpp
index c53c1b8af2f..84e2d5cfcc0 100644
--- a/src/mongo/s/shard_util.cpp
+++ b/src/mongo/s/shard_util.cpp
@@ -90,6 +90,43 @@ StatusWith<long long> retrieveTotalShardSize(OperationContext* opCtx, const Shar
     return totalSizeElem.numberLong();
 }
 
+StatusWith<long long> retrieveCollectionShardSize(OperationContext* opCtx,
+                                                  const ShardId& shardId,
+                                                  NamespaceString const& ns,
+                                                  bool estimate) {
+    auto shardStatus = Grid::get(opCtx)->shardRegistry()->getShard(opCtx, shardId);
+    if (!shardStatus.isOK()) {
+        return shardStatus.getStatus();
+    }
+
+    const Minutes maxTimeMSOverride{10};
+    const auto cmdObj = BSON("dataSize" << ns.ns() << "estimate" << estimate);
+    auto statStatus = shardStatus.getValue()->runCommandWithFixedRetryAttempts(
+        opCtx,
+        ReadPreferenceSetting{ReadPreference::PrimaryPreferred},
+        ns.db().toString(),
+        cmdObj,
+        maxTimeMSOverride,
+        Shard::RetryPolicy::kIdempotent);
+
+    auto stat = Shard::CommandResponse::getEffectiveStatus(statStatus);
+    if (!stat.isOK()) {
+        if (stat == ErrorCodes::NamespaceNotFound) {
+            return 0;
+        }
+        return stat;
+    }
+
+
+    BSONElement sizeElem = statStatus.getValue().response["size"];
+    if (!sizeElem.isNumber()) {
+        return {ErrorCodes::NoSuchKey, "size field not found in dataSize"};
+    }
+
+    return sizeElem.safeNumberLong();
+}
+
+
 StatusWith<std::vector<BSONObj>> selectChunkSplitPoints(OperationContext* opCtx,
                                                         const ShardId& shardId,
                                                         const NamespaceString& nss,
diff --git a/src/mongo/s/shard_util.h b/src/mongo/s/shard_util.h
index 7c0ce16fa8b..3152558df4d 100644
--- a/src/mongo/s/shard_util.h
+++ b/src/mongo/s/shard_util.h
@@ -62,6 +62,19 @@ namespace shardutil {
 StatusWith<long long> retrieveTotalShardSize(OperationContext* opCtx, const ShardId& shardId);
 
 /**
+ * Executes the dataSize command against the specified shard and obtains the total data
+ * size for the collection in bytes (essentially, the dataSize field).
+ *
+ * Returns OK with the total size in bytes or an error. Known errors are:
+ *  ShardNotFound if shard by that id is not available on the registry
+ *  NoSuchKey if the total shard size could not be retrieved
+ */
+StatusWith<long long> retrieveCollectionShardSize(OperationContext* opCtx,
+                                                  const ShardId& shardId,
+                                                  NamespaceString const& ns,
+                                                  bool estimate = true);
+
+/**
  * Ask the specified shard to figure out the split points for a given chunk.
  *
  * shardId The shard id to query.