From ede5f0da8f7bdf020843559ae137a763a6400cf0 Mon Sep 17 00:00:00 2001
From: Paolo Polato <paolo.polato@mongodb.com>
Date: Mon, 16 May 2022 14:09:29 +0000
Subject: SERVER-66377 make ClusterChunksResizePolicy resilient to
 renameCollection() with dropTarget

(cherry picked from commit 46c54afca7f15c78a578a4941b3a765f8f4e90c7)
---
 .../balancer/cluster_chunks_resize_policy_impl.cpp | 27 ++++++++++++++++++----
 .../s/balancer/cluster_chunks_resize_policy_impl.h |  4 +++-
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.cpp b/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.cpp
index c05d998fe86..e886b8b9e48 100644
--- a/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.cpp
+++ b/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.cpp
@@ -90,7 +90,7 @@ void applyResultToCollState(OperationContext* opCtx,
                       "Marking collection chunks resize state as pending to be restarded",
                       "namespace"_attr = collectionState.getNss().ns(),
                       "err"_attr = result);
-        collectionState.errorDetectedOnActionCompleted();
+        collectionState.errorDetectedOnActionCompleted(opCtx);
     }
 }
 
@@ -142,7 +142,7 @@ boost::optional<DefragmentationAction> CollectionState::popNextAction(OperationC
                   "pending to be restarted",
                   "namespace"_attr = _nss.ns(),
                   "error"_attr = redact(e));
-            _restartRequested = true;
+            _requestRestart(opCtx);
         }
     }
     return boost::none;
@@ -168,8 +168,8 @@ void CollectionState::actionCompleted(std::vector<ActionRequestInfo>&& followUpR
     }
 }
 
-void CollectionState::errorDetectedOnActionCompleted() {
-    _restartRequested = true;
+void CollectionState::errorDetectedOnActionCompleted(OperationContext* opCtx) {
+    _requestRestart(opCtx);
     --_numOutstandingActions;
     _pendingRequests.clear();
 }
@@ -190,6 +190,12 @@ const UUID& CollectionState::getUuid() const {
     return _uuid;
 }
 
+void CollectionState::_requestRestart(OperationContext* opCtx) {
+    // Invalidate the CollectionState object and the related entry in the cache.
+    _restartRequested = true;
+    Grid::get(opCtx)->catalogCache()->invalidateCollectionEntry_LINEARIZABLE(_nss);
+}
+
 ClusterChunksResizePolicyImpl::ClusterChunksResizePolicyImpl(
     const std::function<void()>& onStateUpdated)
     : _onStateUpdated(onStateUpdated) {}
@@ -266,11 +272,24 @@ boost::optional<DefragmentationAction> ClusterChunksResizePolicyImpl::getNextStr
             if (collState.restartNeeded()) {
                 auto refreshedCollState = _buildInitialStateFor(opCtx, collState.getNss());
                 if (!refreshedCollState) {
+                    // The collection does not longer exist - discard it
                     auto entryToErase = it;
                     it = std::next(it);
                     _collectionsBeingProcessed.erase(entryToErase);
                     continue;
                 }
+                if (refreshedCollState->getUuid() != collUuid) {
+                    // the collection has been dropped and re-created since the creation of the
+                    // element; re-insert the entry with an updated UUID (this operation will
+                    // invalidate the iterator).
+                    _collectionsBeingProcessed.erase(it);
+                    _collectionsBeingProcessed.emplace(refreshedCollState->getUuid(),
+                                                       std::move(*refreshedCollState));
+                    it = _collectionsBeingProcessed.begin();
+                    continue;
+                }
+
+                // update the current element
                 collState = std::move(*refreshedCollState);
             }
 
diff --git a/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.h b/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.h
index 7bdcdd79596..b53b7afbeda 100644
--- a/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.h
+++ b/src/mongo/db/s/balancer/cluster_chunks_resize_policy_impl.h
@@ -57,7 +57,7 @@ public:
 
     void actionCompleted(std::vector<ActionRequestInfo>&& followUpRequests);
 
-    void errorDetectedOnActionCompleted();
+    void errorDetectedOnActionCompleted(OperationContext* opCtx);
 
     bool restartNeeded() const;
 
@@ -86,6 +86,8 @@ private:
     std::vector<ActionRequestInfo> _pendingRequests;
     int _numOutstandingActions;
     bool _restartRequested;
+
+    void _requestRestart(OperationContext* opCtx);
 };
 
 class ClusterChunksResizePolicyImpl : public ClusterChunksResizePolicy {
-- 
cgit v1.2.1