SERVER-41583 Refactor the registration and de-registration of collection and catalog entry

(cherry picked from commit 5eda33f9fa40a1a17f9f63f904a8c147700d648c)

5 files changed, 36 insertions, 66 deletions

diff --git a/src/mongo/db/storage/kv/kv_catalog.cpp b/src/mongo/db/storage/kv/kv_catalog.cpp
index a102a8707d4..bcdec978c0e 100644
--- a/src/mongo/db/storage/kv/kv_catalog.cpp
+++ b/src/mongo/db/storage/kv/kv_catalog.cpp
@@ -35,7 +35,6 @@
 
 #include "mongo/bson/util/bson_extract.h"
 #include "mongo/bson/util/builder.h"
-#include "mongo/db/catalog/collection_catalog.h"
 #include "mongo/db/concurrency/d_concurrency.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/operation_context.h"
@@ -731,10 +730,11 @@ std::unique_ptr<CollectionCatalogEntry> KVCatalog::makeCollectionCatalogEntry(
         _engine, this, nss.ns(), ident, std::move(rs));
 }
 
-Status KVCatalog::createCollection(OperationContext* opCtx,
-                                   const NamespaceString& nss,
-                                   const CollectionOptions& options,
-                                   bool allocateDefaultSpace) {
+StatusWith<std::unique_ptr<CollectionCatalogEntry>> KVCatalog::createCollection(
+    OperationContext* opCtx,
+    const NamespaceString& nss,
+    const CollectionOptions& options,
+    bool allocateDefaultSpace) {
     invariant(opCtx->lockState()->isDbLockedForMode(nss.db(), MODE_IX));
     invariant(nss.coll().size() > 0);
 
@@ -769,19 +769,13 @@ Status KVCatalog::createCollection(OperationContext* opCtx,
     opCtx->recoveryUnit()->onRollback([ opCtx, catalog = this, nss, ident, uuid ]() {
         // Intentionally ignoring failure
         catalog->_engine->getEngine()->dropIdent(opCtx, ident).ignore();
-
-        CollectionCatalog::get(opCtx).deregisterCatalogEntry(uuid);
     });
 
     auto rs = _engine->getEngine()->getGroupedRecordStore(opCtx, nss.ns(), ident, options, prefix);
     invariant(rs);
 
-    CollectionCatalog::get(getGlobalServiceContext())
-        .registerCatalogEntry(uuid,
-                              std::make_unique<KVCollectionCatalogEntry>(
-                                  _engine, this, nss.ns(), ident, std::move(rs)));
-
-    return Status::OK();
+    return std::make_unique<KVCollectionCatalogEntry>(
+        _engine, this, nss.ns(), ident, std::move(rs));
 }
 
 Status KVCatalog::renameCollection(OperationContext* opCtx,
@@ -808,29 +802,6 @@ Status KVCatalog::renameCollection(OperationContext* opCtx,
     return Status::OK();
 }
 
-class KVCatalog::FinishDropCatalogEntryChange : public RecoveryUnit::Change {
-public:
-    FinishDropCatalogEntryChange(CollectionCatalog& catalog,
-                                 std::unique_ptr<CollectionCatalogEntry> collectionCatalogEntry,
-                                 CollectionUUID uuid)
-        : _catalog(catalog),
-          _collectionCatalogEntry(std::move(collectionCatalogEntry)),
-          _uuid(uuid) {}
-
-    void commit(boost::optional<Timestamp>) override {
-        _collectionCatalogEntry.reset();
-    }
-
-    void rollback() override {
-        _catalog.registerCatalogEntry(_uuid, std::move(_collectionCatalogEntry));
-    }
-
-private:
-    CollectionCatalog& _catalog;
-    std::unique_ptr<CollectionCatalogEntry> _collectionCatalogEntry;
-    CollectionUUID _uuid;
-};
-
 Status KVCatalog::dropCollection(OperationContext* opCtx, const NamespaceString& nss) {
     invariant(opCtx->lockState()->isCollectionLockedForMode(nss, MODE_X));
 
@@ -863,13 +834,6 @@ Status KVCatalog::dropCollection(OperationContext* opCtx, const NamespaceString&
         return status;
     }
 
-    // Remove catalog entry
-    std::unique_ptr<CollectionCatalogEntry> removedCatalogEntry =
-        CollectionCatalog::get(opCtx).deregisterCatalogEntry(uuid.get());
-
-    opCtx->recoveryUnit()->registerChange(new FinishDropCatalogEntryChange(
-        CollectionCatalog::get(opCtx), std::move(removedCatalogEntry), uuid.get()));
-
     // This will lazily delete the KVCollectionCatalogEntry and notify the storageEngine to
     // drop the collection only on WUOW::commit().
     opCtx->recoveryUnit()->onCommit(
@@ -888,7 +852,6 @@ Status KVCatalog::dropCollection(OperationContext* opCtx, const NamespaceString&
             }
         });
 
-
     return Status::OK();
 }
 }
diff --git a/src/mongo/db/storage/kv/kv_catalog.h b/src/mongo/db/storage/kv/kv_catalog.h
index b22b6cf68ff..a173209b53d 100644
--- a/src/mongo/db/storage/kv/kv_catalog.h
+++ b/src/mongo/db/storage/kv/kv_catalog.h
@@ -34,6 +34,7 @@
 #include <string>
 
 #include "mongo/base/string_data.h"
+#include "mongo/db/catalog/collection_catalog.h"
 #include "mongo/db/catalog/collection_options.h"
 #include "mongo/db/record_id.h"
 #include "mongo/db/storage/bson_collection_catalog_entry.h"
@@ -114,10 +115,11 @@ public:
                                                                        const NamespaceString& nss,
                                                                        bool forRepair);
 
-    Status createCollection(OperationContext* opCtx,
-                            const NamespaceString& nss,
-                            const CollectionOptions& options,
-                            bool allocateDefaultSpace);
+    StatusWith<std::unique_ptr<CollectionCatalogEntry>> createCollection(
+        OperationContext* opCtx,
+        const NamespaceString& nss,
+        const CollectionOptions& options,
+        bool allocateDefaultSpace);
 
     Status renameCollection(OperationContext* opCtx,
                             const NamespaceString& fromNss,
@@ -129,7 +131,6 @@ public:
 private:
     class AddIdentChange;
     class RemoveIdentChange;
-    class FinishDropCatalogEntryChange;
 
     friend class KVStorageEngine;
     friend class KVCatalogTest;
diff --git a/src/mongo/db/storage/kv/kv_collection_catalog_entry_test.cpp b/src/mongo/db/storage/kv/kv_collection_catalog_entry_test.cpp
index 7d4e035c222..4beb5d860fc 100644
--- a/src/mongo/db/storage/kv/kv_collection_catalog_entry_test.cpp
+++ b/src/mongo/db/storage/kv/kv_collection_catalog_entry_test.cpp
@@ -84,8 +84,14 @@ public:
             const bool allocateDefaultSpace = true;
             CollectionOptions options;
             options.uuid = UUID::gen();
-            ASSERT_OK(_storageEngine.getCatalog()->createCollection(
-                opCtx.get(), _nss, options, allocateDefaultSpace));
+            auto statusWithCatalogEntry = _storageEngine.getCatalog()->createCollection(
+                opCtx.get(), _nss, options, allocateDefaultSpace);
+            ASSERT_OK(statusWithCatalogEntry.getStatus());
+            auto collection = std::make_unique<CollectionMock>(_nss);
+            CollectionCatalog::get(opCtx.get())
+                .registerCollection(options.uuid.get(),
+                                    std::move(statusWithCatalogEntry.getValue()),
+                                    std::move(collection));
             wuow.commit();
         }
     }
diff --git a/src/mongo/db/storage/kv/kv_storage_engine.cpp b/src/mongo/db/storage/kv/kv_storage_engine.cpp
index d295631076b..328fa7cd1c5 100644
--- a/src/mongo/db/storage/kv/kv_storage_engine.cpp
+++ b/src/mongo/db/storage/kv/kv_storage_engine.cpp
@@ -245,8 +245,7 @@ void KVStorageEngine::_initCollection(OperationContext* opCtx,
     auto collection = collectionFactory->make(opCtx, catalogEntry.get());
 
     auto& collectionCatalog = CollectionCatalog::get(getGlobalServiceContext());
-    collectionCatalog.registerCatalogEntry(uuid, std::move(catalogEntry));
-    collectionCatalog.registerCollectionObject(uuid, std::move(collection));
+    collectionCatalog.registerCollection(uuid, std::move(catalogEntry), std::move(collection));
 }
 
 void KVStorageEngine::closeCatalog(OperationContext* opCtx) {
@@ -594,10 +593,18 @@ Status KVStorageEngine::_dropCollectionsNoTimestamp(OperationContext* opCtx,
     WriteUnitOfWork untimestampedDropWuow(opCtx);
     for (auto& nss : toDrop) {
         invariant(getCatalog());
+        auto uuid = CollectionCatalog::get(opCtx).lookupUUIDByNSS(nss).get();
         Status result = getCatalog()->dropCollection(opCtx, nss);
+
         if (!result.isOK() && firstError.isOK()) {
             firstError = result;
         }
+
+        auto[removedColl, removedCatalogEntry] =
+            CollectionCatalog::get(opCtx).deregisterCollection(uuid);
+        opCtx->recoveryUnit()->registerChange(
+            CollectionCatalog::get(opCtx).makeFinishDropCollectionChange(
+                std::move(removedColl), std::move(removedCatalogEntry), uuid));
     }
 
     untimestampedDropWuow.commit();
@@ -665,17 +672,10 @@ Status KVStorageEngine::repairRecordStore(OperationContext* opCtx, const Namespa
                                                      << status.reason());
     }
 
+    // After repairing, re-initialize the collection with a valid RecordStore.
     auto& collectionCatalog = CollectionCatalog::get(getGlobalServiceContext());
     auto uuid = collectionCatalog.lookupUUIDByNSS(nss).get();
-
-    // It's possible the Collection may not already have been removed if the no DatabaseHolder was
-    // opened for a database.
-    if (collectionCatalog.lookupCollectionByUUID(uuid)) {
-        collectionCatalog.deregisterCollectionObject(uuid);
-    }
-    collectionCatalog.deregisterCatalogEntry(uuid);
-
-    // After repairing, initialize the collection with a valid RecordStore.
+    collectionCatalog.deregisterCollection(uuid);
     _initCollection(opCtx, nss, false);
     return Status::OK();
 }
diff --git a/src/mongo/db/storage/kv/kv_storage_engine_test_fixture.h b/src/mongo/db/storage/kv/kv_storage_engine_test_fixture.h
index 29cf88e7a84..c9cfbafe911 100644
--- a/src/mongo/db/storage/kv/kv_storage_engine_test_fixture.h
+++ b/src/mongo/db/storage/kv/kv_storage_engine_test_fixture.h
@@ -55,10 +55,10 @@ public:
         AutoGetDb db(opCtx, ns.db(), LockMode::MODE_X);
         CollectionOptions options;
         options.uuid = UUID::gen();
-        auto ret = _storageEngine->getCatalog()->createCollection(opCtx, ns, options, true);
-        if (!ret.isOK()) {
-            return ret;
-        }
+        auto catalogEntry = unittest::assertGet(
+            _storageEngine->getCatalog()->createCollection(opCtx, ns, options, true));
+        CollectionCatalog::get(opCtx).registerCollection(
+            options.uuid.get(), std::move(catalogEntry), std::make_unique<CollectionMock>(ns));
 
         return _storageEngine->getCatalog()->getCollectionIdent(ns);
     }