3 files changed, 48 insertions, 32 deletions

diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index b180654ffb4..dc5219316c3 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -957,40 +957,67 @@ void Balancer::abortCollectionDefragmentation(OperationContext* opCtx, const Nam
     _defragmentationPolicy->abortCollectionDefragmentation(opCtx, nss);
 }
 
-Balancer::BalancerStatus Balancer::getBalancerStatusForNs(OperationContext* opCtx,
-                                                          const NamespaceString& ns) {
+BalancerCollectionStatusResponse Balancer::getBalancerStatusForNs(OperationContext* opCtx,
+                                                                  const NamespaceString& ns) {
+    CollectionType coll;
     try {
-        auto coll = Grid::get(opCtx)->catalogClient()->getCollection(opCtx, ns, {});
-        bool isDefragmenting = coll.getDefragmentCollection();
-        if (isDefragmenting) {
-            return {false,
-                    kBalancerPolicyStatusDefragmentingChunks.toString(),
-                    _defragmentationPolicy->reportProgressOn(coll.getUuid())};
+        coll = Grid::get(opCtx)->catalogClient()->getCollection(opCtx, ns, {});
+    } catch (const ExceptionFor<ErrorCodes::NamespaceNotFound>&) {
+        uasserted(ErrorCodes::NamespaceNotSharded, "Collection unsharded or undefined");
+    }
+
+    const auto maxChunkSizeMB = [&]() -> int64_t {
+        int64_t value = 0;
+        if (const auto& collOverride = coll.getMaxChunkSizeBytes(); collOverride.is_initialized()) {
+            value = *collOverride;
+        } else {
+            auto balancerConfig = Grid::get(opCtx)->getBalancerConfiguration();
+            uassertStatusOK(balancerConfig->refreshAndCheck(opCtx));
+            value = balancerConfig->getMaxChunkSizeBytes();
         }
-    } catch (DBException&) {
-        // Catch exceptions to keep consistency with errors thrown before defragmentation
+        return value / (1024 * 1024);
+    }();
+    BalancerCollectionStatusResponse response(maxChunkSizeMB, true /*balancerCompliant*/);
+    auto setViolationOnResponse = [&response](const StringData& reason,
+                                              const boost::optional<BSONObj>& details =
+                                                  boost::none) {
+        response.setBalancerCompliant(false);
+        response.setFirstComplianceViolation(reason);
+        response.setDetails(details);
+    };
+
+    bool isDefragmenting = coll.getDefragmentCollection();
+    if (isDefragmenting) {
+        setViolationOnResponse(kBalancerPolicyStatusDefragmentingChunks,
+                               _defragmentationPolicy->reportProgressOn(coll.getUuid()));
+        return response;
     }
 
     auto splitChunks = uassertStatusOK(_chunkSelectionPolicy->selectChunksToSplit(opCtx, ns));
     if (!splitChunks.empty()) {
-        return {false, kBalancerPolicyStatusZoneViolation.toString()};
+        setViolationOnResponse(kBalancerPolicyStatusZoneViolation);
+        return response;
     }
+
     auto chunksToMove = uassertStatusOK(_chunkSelectionPolicy->selectChunksToMove(opCtx, ns));
     if (chunksToMove.empty()) {
-        return {true, boost::none, boost::none};
+        return response;
     }
-    const auto& migrationInfo = chunksToMove.front();
 
+    const auto& migrationInfo = chunksToMove.front();
     switch (migrationInfo.reason) {
         case MigrateInfo::drain:
-            return {false, kBalancerPolicyStatusDraining.toString(), boost::none};
+            setViolationOnResponse(kBalancerPolicyStatusDraining);
+            break;
         case MigrateInfo::zoneViolation:
-            return {false, kBalancerPolicyStatusZoneViolation.toString(), boost::none};
+            setViolationOnResponse(kBalancerPolicyStatusZoneViolation);
+            break;
         case MigrateInfo::chunksImbalance:
-            return {false, kBalancerPolicyStatusChunksImbalance.toString(), boost::none};
+            setViolationOnResponse(kBalancerPolicyStatusChunksImbalance);
+            break;
     }
 
-    return {true, boost::none, boost::none};
+    return response;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/s/balancer/balancer.h b/src/mongo/db/s/balancer/balancer.h
index c3ea9b0578f..f0710e245b6 100644
--- a/src/mongo/db/s/balancer/balancer.h
+++ b/src/mongo/db/s/balancer/balancer.h
@@ -33,6 +33,7 @@
 #include "mongo/db/s/balancer/balancer_chunk_selection_policy.h"
 #include "mongo/db/s/balancer/balancer_random.h"
 #include "mongo/platform/mutex.h"
+#include "mongo/s/request_types/balancer_collection_status_gen.h"
 #include "mongo/stdx/condition_variable.h"
 #include "mongo/stdx/thread.h"
 
@@ -183,16 +184,12 @@ public:
      */
     void abortCollectionDefragmentation(OperationContext* opCtx, const NamespaceString& nss);
 
-    struct BalancerStatus {
-        bool balancerCompliant;
-        boost::optional<std::string> firstComplianceViolation;
-        boost::optional<BSONObj> details;
-    };
     /**
      * Returns if a given collection is draining due to a removed shard, has chunks on an invalid
      * zone or the number of chunks is imbalanced across the cluster
      */
-    BalancerStatus getBalancerStatusForNs(OperationContext* opCtx, const NamespaceString& nss);
+    BalancerCollectionStatusResponse getBalancerStatusForNs(OperationContext* opCtx,
+                                                            const NamespaceString& nss);
 
 private:
     /**
diff --git a/src/mongo/db/s/config/configsvr_balancer_collection_status_command.cpp b/src/mongo/db/s/config/configsvr_balancer_collection_status_command.cpp
index 54e21cd2be8..1791cab3943 100644
--- a/src/mongo/db/s/config/configsvr_balancer_collection_status_command.cpp
+++ b/src/mongo/db/s/config/configsvr_balancer_collection_status_command.cpp
@@ -71,15 +71,7 @@ public:
             uassert(ErrorCodes::InvalidNamespace,
                     str::stream() << "Invalid namespace specified '" << nss.ns() << "'",
                     nss.isValid());
-            const auto& balancerStatus = Balancer::get(opCtx)->getBalancerStatusForNs(opCtx, nss);
-            Response response(balancerStatus.balancerCompliant);
-            response.setDetails(balancerStatus.details);
-            response.setFirstComplianceViolation(
-                balancerStatus.firstComplianceViolation.is_initialized()
-                    ? boost::optional<StringData>(
-                          StringData(*balancerStatus.firstComplianceViolation))
-                    : boost::optional<StringData>(boost::none));
-            return response;
+            return Balancer::get(opCtx)->getBalancerStatusForNs(opCtx, nss);
         }
 
     private: