SERVER-46253 Only register rollback handlers for registered collections

3 files changed, 15 insertions, 17 deletions

diff --git a/src/mongo/db/catalog/collection_catalog.h b/src/mongo/db/catalog/collection_catalog.h
index 85ff4a8daad..1f2c98de80d 100644
--- a/src/mongo/db/catalog/collection_catalog.h
+++ b/src/mongo/db/catalog/collection_catalog.h
@@ -254,7 +254,6 @@ private:
 
     const std::vector<CollectionUUID>& _getOrdering_inlock(const StringData& db,
                                                            const stdx::lock_guard<Latch>&);
-
     mutable mongo::Mutex _catalogLock;
 
     /**
diff --git a/src/mongo/db/catalog/uncommitted_collections.cpp b/src/mongo/db/catalog/uncommitted_collections.cpp
index 34b367c1912..db9a6e3c974 100644
--- a/src/mongo/db/catalog/uncommitted_collections.cpp
+++ b/src/mongo/db/catalog/uncommitted_collections.cpp
@@ -67,7 +67,6 @@ void UncommittedCollections::addToTxn(OperationContext* opCtx,
         UncommittedCollections::erase(uuid, nss, collListUnowned.lock().get());
     });
 
-    auto svcCtx = opCtx->getServiceContext();
 
     opCtx->recoveryUnit()->registerPreCommitHook(
         [collListUnowned, uuid, createTime](OperationContext* opCtx) {
@@ -80,9 +79,6 @@ void UncommittedCollections::addToTxn(OperationContext* opCtx,
             invariant(collPtr->getMinimumVisibleSnapshot() == createTime);
             UncommittedCollections::clear(collListUnowned.lock().get());
         });
-    opCtx->recoveryUnit()->onRollback([svcCtx, collListUnowned, uuid, nss]() {
-        UncommittedCollections::rollback(svcCtx, uuid, collListUnowned.lock().get());
-    });
 }
 
 Collection* UncommittedCollections::getForTxn(OperationContext* opCtx,
@@ -121,13 +117,10 @@ void UncommittedCollections::erase(UUID uuid, NamespaceString nss, UncommittedCo
 void UncommittedCollections::rollback(ServiceContext* svcCtx,
                                       CollectionUUID uuid,
                                       UncommittedCollectionsMap* map) {
-    auto it = std::find(map->_registeredUUIDs.begin(), map->_registeredUUIDs.end(), uuid);
-    if (it != map->_registeredUUIDs.end()) {
-        auto collPtr = CollectionCatalog::get(svcCtx).deregisterCollection(uuid);
-        auto nss = collPtr.get()->ns();
-        map->_collections[uuid] = std::move(collPtr);
-        map->_nssIndex.insert({nss, uuid});
-    }
+    auto collPtr = CollectionCatalog::get(svcCtx).deregisterCollection(uuid);
+    auto nss = collPtr.get()->ns();
+    map->_collections[uuid] = std::move(collPtr);
+    map->_nssIndex.insert({nss, uuid});
 }
 
 void UncommittedCollections::commit(OperationContext* opCtx,
@@ -147,7 +140,12 @@ void UncommittedCollections::commit(OperationContext* opCtx,
     CollectionCatalog::get(opCtx).registerCollection(uuid, &(it->second));
     map->_collections.erase(it);
     map->_nssIndex.erase(nss);
-    map->_registeredUUIDs.push_back(uuid);
+    auto svcCtx = opCtx->getServiceContext();
+    auto collListUnowned = getUncommittedCollections(opCtx).getResources();
+
+    opCtx->recoveryUnit()->onRollback([svcCtx, collListUnowned, uuid]() {
+        UncommittedCollections::rollback(svcCtx, uuid, collListUnowned.lock().get());
+    });
 }
 
 bool UncommittedCollections::isUncommittedCollection(OperationContext* opCtx,
diff --git a/src/mongo/db/catalog/uncommitted_collections.h b/src/mongo/db/catalog/uncommitted_collections.h
index fa2025354aa..2d13b855ab8 100644
--- a/src/mongo/db/catalog/uncommitted_collections.h
+++ b/src/mongo/db/catalog/uncommitted_collections.h
@@ -58,7 +58,6 @@ public:
 
         std::map<UUID, std::unique_ptr<Collection>> _collections;
         std::map<NamespaceString, UUID> _nssIndex;
-        std::vector<UUID> _registeredUUIDs;
     };
 
     UncommittedCollections() {
@@ -92,6 +91,8 @@ public:
     /**
      * Registers any uncommitted collections with the CollectionCatalog. If registering a collection
      * name conflicts with an existing entry, this method will throw a `WriteConflictException`.
+     * This method also clears the entries for the collection identified by `uuid` from
+     * UncommittedCollections.
      */
     static void commit(OperationContext* opCtx,
                        UUID uuid,
@@ -99,8 +100,9 @@ public:
                        UncommittedCollectionsMap* map);
 
     /**
-     * If the collection with uuid `uuid` was previously registered with the CollectionCatalog as
-     * part of the current multi-document transaction, this handler deregisters it.
+     * Deregisters the collection with uuid `uuid` from the CollectionCatalog, and re-adds the
+     * entries for the collection identified by `uuid` to UncommittedCollections. This function
+     * assumes `commit` has previously been called for `uuid`.
      */
     static void rollback(ServiceContext* svcCtx,
                          CollectionUUID uuid,
@@ -120,7 +122,6 @@ public:
     static void clear(UncommittedCollectionsMap* map) {
         map->_collections.clear();
         map->_nssIndex.clear();
-        map->_registeredUUIDs.clear();
     }
 
 private: