SERVER-65429 Clean up defragmentation policy memory on stepdown

(cherry picked from commit b4384a6a6774ae494e57adc0c625769fec107828)

4 files changed, 16 insertions, 0 deletions

diff --git a/src/mongo/db/s/balancer/balancer.cpp b/src/mongo/db/s/balancer/balancer.cpp
index 52bacd4d743..7d9b40d5ac5 100644
--- a/src/mongo/db/s/balancer/balancer.cpp
+++ b/src/mongo/db/s/balancer/balancer.cpp
@@ -389,6 +389,8 @@ void Balancer::report(OperationContext* opCtx, BSONObjBuilder* builder) {
 }
 
 void Balancer::_consumeActionStreamLoop() {
+    ScopeGuard onExitCleanup([this] { _defragmentationPolicy->interruptAllDefragmentations(); });
+
     Client::initThread("BalancerSecondary");
     auto applyThrottling = [lastActionTime = Date_t::fromMillisSinceEpoch(0)]() mutable {
         const Milliseconds throttle{chunkDefragmentationThrottlingMS.load()};
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy.h b/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
index 3cada882af9..26fd3c6212b 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy.h
@@ -59,6 +59,13 @@ public:
                                                 const NamespaceString& nss) = 0;
 
     /**
+     * Requests to stop the emission of any new defragmentation action request. Does not alter the
+     * persisted state of the affected collections. startCollectionDefragmentation() can be invoked
+     * on a later stage to resume the defragmentation on each item.
+     */
+    virtual void interruptAllDefragmentations() = 0;
+
+    /**
      * Returns true if the specified collection is currently being defragmented.
      */
     virtual bool isDefragmentingCollection(const UUID& uuid) = 0;
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
index fbbb0f53236..42c415003f9 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.cpp
@@ -1431,6 +1431,11 @@ void BalancerDefragmentationPolicyImpl::abortCollectionDefragmentation(Operation
     }
 }
 
+void BalancerDefragmentationPolicyImpl::interruptAllDefragmentations() {
+    stdx::lock_guard<Latch> lk(_stateMutex);
+    _defragmentationStates.clear();
+}
+
 bool BalancerDefragmentationPolicyImpl::isDefragmentingCollection(const UUID& uuid) {
     stdx::lock_guard<Latch> lk(_stateMutex);
     return _defragmentationStates.contains(uuid);
diff --git a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
index 68b30cdfc29..80a59a89a08 100644
--- a/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
+++ b/src/mongo/db/s/balancer/balancer_defragmentation_policy_impl.h
@@ -76,6 +76,8 @@ public:
 
     ~BalancerDefragmentationPolicyImpl() {}
 
+    void interruptAllDefragmentations() override;
+
     bool isDefragmentingCollection(const UUID& uuid) override;
 
     virtual BSONObj reportProgressOn(const UUID& uuid) override;