SERVER-67383 Use `CollectionNamespaceOrUUIDLock` to lock collections via UUID

7 files changed, 67 insertions, 57 deletions

diff --git a/src/mongo/db/catalog/drop_collection.cpp b/src/mongo/db/catalog/drop_collection.cpp
index 5e9089fce0c..645362a327d 100644
--- a/src/mongo/db/catalog/drop_collection.cpp
+++ b/src/mongo/db/catalog/drop_collection.cpp
@@ -198,7 +198,7 @@ Status _abortIndexBuildsAndDrop(OperationContext* opCtx,
     // We only need to hold an intent lock to send abort signals to the active index builder on this
     // collection.
     boost::optional<AutoGetDb> optionalAutoDb(std::move(autoDb));
-    boost::optional<Lock::CollectionLock> collLock;
+    boost::optional<CollectionNamespaceOrUUIDLock> collLock;
     collLock.emplace(opCtx, startingNss, MODE_IX);
 
     // Abandon the snapshot as the index catalog will compare the in-memory state to the disk state,
diff --git a/src/mongo/db/catalog/validate_state.h b/src/mongo/db/catalog/validate_state.h
index a186839e839..c508a6b5850 100644
--- a/src/mongo/db/catalog/validate_state.h
+++ b/src/mongo/db/catalog/validate_state.h
@@ -237,7 +237,7 @@ private:
     boost::optional<ShouldNotConflictWithSecondaryBatchApplicationBlock> _noPBWM;
     boost::optional<Lock::GlobalLock> _globalLock;
     boost::optional<AutoGetDb> _databaseLock;
-    boost::optional<Lock::CollectionLock> _collectionLock;
+    boost::optional<CollectionNamespaceOrUUIDLock> _collectionLock;
 
     Database* _database;
     CollectionPtr _collection;
diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index 3197c370cea..d05e5e2e094 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -122,7 +122,7 @@ void acquireCollectionLocksInResourceIdOrder(
     LockMode modeColl,
     Date_t deadline,
     const std::vector<NamespaceStringOrUUID>& secondaryNssOrUUIDs,
-    std::vector<Lock::CollectionLock>* collLocks) {
+    std::vector<CollectionNamespaceOrUUIDLock>* collLocks) {
     invariant(collLocks->empty());
     auto catalog = CollectionCatalog::get(opCtx);
 
@@ -195,6 +195,32 @@ Database* AutoGetDb::ensureDbExists(OperationContext* opCtx) {
     return _db;
 }
 
+CollectionNamespaceOrUUIDLock::CollectionNamespaceOrUUIDLock(OperationContext* opCtx,
+                                                             const NamespaceStringOrUUID& nsOrUUID,
+                                                             LockMode mode,
+                                                             Date_t deadline)
+    : _lock([opCtx, &nsOrUUID, mode, deadline] {
+          if (auto& ns = nsOrUUID.nss()) {
+              return Lock::CollectionLock{opCtx, *ns, mode, deadline};
+          }
+
+          auto resolveNs = [opCtx, &nsOrUUID] {
+              return CollectionCatalog::get(opCtx)->resolveNamespaceStringOrUUID(opCtx, nsOrUUID);
+          };
+
+          // We cannot be sure that the namespace we lock matches the UUID given because we resolve
+          // the namespace from the UUID without the safety of a lock. Therefore, we will continue
+          // to re-lock until the namespace we resolve from the UUID before and after taking the
+          // lock is the same.
+          while (true) {
+              auto ns = resolveNs();
+              Lock::CollectionLock lock{opCtx, ns, mode, deadline};
+              if (ns == resolveNs()) {
+                  return lock;
+              }
+          }
+      }()) {}
+
 AutoGetCollection::AutoGetCollection(OperationContext* opCtx,
                                      const NamespaceStringOrUUID& nsOrUUID,
                                      LockMode modeColl,
diff --git a/src/mongo/db/catalog_raii.h b/src/mongo/db/catalog_raii.h
index 6eed6afd845..7b0a0af30af 100644
--- a/src/mongo/db/catalog_raii.h
+++ b/src/mongo/db/catalog_raii.h
@@ -90,6 +90,27 @@ private:
     std::vector<Lock::DBLock> _secondaryDbLocks;
 };
 
+/**
+ * Light wrapper around Lock::CollectionLock which allows acquiring the lock based on UUID rather
+ * than namespace.
+ *
+ * The lock manager manages resources based on namespace and does not have a concept of UUIDs, so
+ * there must be some additional concurrency checks around resolving the UUID to a namespace and
+ * then subsequently acquiring the lock.
+ */
+class CollectionNamespaceOrUUIDLock {
+public:
+    CollectionNamespaceOrUUIDLock(OperationContext* opCtx,
+                                  const NamespaceStringOrUUID& nsOrUUID,
+                                  LockMode mode,
+                                  Date_t deadline = Date_t::max());
+
+    CollectionNamespaceOrUUIDLock(CollectionNamespaceOrUUIDLock&& other) = default;
+
+private:
+    Lock::CollectionLock _lock;
+};
+
 namespace auto_get_collection {
 enum class ViewMode { kViewsPermitted, kViewsForbidden };
 
@@ -224,7 +245,7 @@ protected:
     // Ordering matters, the _collLocks should destruct before the _autoGetDb releases the
     // rstl/global/database locks.
     boost::optional<AutoGetDb> _autoDb;
-    std::vector<Lock::CollectionLock> _collLocks;
+    std::vector<CollectionNamespaceOrUUIDLock> _collLocks;
 
     CollectionPtr _coll = nullptr;
     std::shared_ptr<const ViewDefinition> _view;
diff --git a/src/mongo/db/concurrency/d_concurrency.cpp b/src/mongo/db/concurrency/d_concurrency.cpp
index a678e49bbb0..f46c09d636e 100644
--- a/src/mongo/db/concurrency/d_concurrency.cpp
+++ b/src/mongo/db/concurrency/d_concurrency.cpp
@@ -33,7 +33,6 @@
 #include <string>
 #include <vector>
 
-#include "mongo/db/catalog/collection_catalog.h"
 #include "mongo/db/concurrency/flow_control_ticketholder.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/service_context.h"
@@ -246,49 +245,15 @@ Lock::DBLock::~DBLock() {
 }
 
 Lock::CollectionLock::CollectionLock(OperationContext* opCtx,
-                                     const NamespaceStringOrUUID& nssOrUUID,
+                                     const NamespaceString& ns,
                                      LockMode mode,
                                      Date_t deadline)
-    : _opCtx(opCtx) {
-    if (nssOrUUID.nss()) {
-        auto& nss = *nssOrUUID.nss();
-        _id = {RESOURCE_COLLECTION, nss};
-
-        invariant(nss.coll().size(), str::stream() << "expected non-empty collection name:" << nss);
-        dassert(_opCtx->lockState()->isDbLockedForMode(nss.dbName(),
-                                                       isSharedLockMode(mode) ? MODE_IS : MODE_IX));
-
-        _opCtx->lockState()->lock(_opCtx, _id, mode, deadline);
-        return;
-    }
-
-    // 'nsOrUUID' must be a UUID and dbName.
-
-    auto nss = CollectionCatalog::get(opCtx)->resolveNamespaceStringOrUUID(opCtx, nssOrUUID);
-
-    // The UUID cannot move between databases so this one dassert is sufficient.
-    dassert(_opCtx->lockState()->isDbLockedForMode(nss.dbName(),
+    : _id(RESOURCE_COLLECTION, ns), _opCtx(opCtx) {
+    invariant(!ns.coll().empty());
+    dassert(_opCtx->lockState()->isDbLockedForMode(ns.dbName(),
                                                    isSharedLockMode(mode) ? MODE_IS : MODE_IX));
 
-    // We cannot be sure that the namespace we lock matches the UUID given because we resolve the
-    // namespace from the UUID without the safety of a lock. Therefore, we will continue to re-lock
-    // until the namespace we resolve from the UUID before and after taking the lock is the same.
-    bool locked = false;
-    NamespaceString prevResolvedNss;
-    do {
-        if (locked) {
-            _opCtx->lockState()->unlock(_id);
-        }
-
-        _id = ResourceId(RESOURCE_COLLECTION, nss);
-        _opCtx->lockState()->lock(_opCtx, _id, mode, deadline);
-        locked = true;
-
-        // We looked up UUID without a collection lock so it's possible that the
-        // collection name changed now. Look it up again.
-        prevResolvedNss = nss;
-        nss = CollectionCatalog::get(opCtx)->resolveNamespaceStringOrUUID(opCtx, nssOrUUID);
-    } while (nss != prevResolvedNss);
+    _opCtx->lockState()->lock(_opCtx, _id, mode, deadline);
 }
 
 Lock::CollectionLock::CollectionLock(CollectionLock&& otherLock)
diff --git a/src/mongo/db/concurrency/d_concurrency.h b/src/mongo/db/concurrency/d_concurrency.h
index b999c6103da..2d4b5b6d6a8 100644
--- a/src/mongo/db/concurrency/d_concurrency.h
+++ b/src/mongo/db/concurrency/d_concurrency.h
@@ -37,8 +37,6 @@
 
 namespace mongo {
 
-class NamespaceStringOrUUID;
-
 class Lock {
 public:
     /**
@@ -350,7 +348,7 @@ public:
 
     public:
         CollectionLock(OperationContext* opCtx,
-                       const NamespaceStringOrUUID& nssOrUUID,
+                       const NamespaceString& ns,
                        LockMode mode,
                        Date_t deadline = Date_t::max());
 
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index 5de2a889dc6..7708d48886d 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -668,7 +668,7 @@ Status IndexBuildsCoordinator::_setUpResumeIndexBuild(OperationContext* opCtx,
     }
 
     Lock::DBLock dbLock(opCtx, dbName, MODE_IX);
-    Lock::CollectionLock collLock(opCtx, nssOrUuid, MODE_X);
+    CollectionNamespaceOrUUIDLock collLock(opCtx, nssOrUuid, MODE_X);
 
     CollectionWriter collection(opCtx, resumeInfo.getCollectionUUID());
     invariant(collection);
@@ -1251,7 +1251,7 @@ bool IndexBuildsCoordinator::abortIndexBuildByBuildUUID(OperationContext* opCtx,
             // taking a strong collection lock. See SERVER-42621.
             unlockRSTL(opCtx);
         }
-        Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_X);
+        CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_X);
         AutoGetCollection indexBuildEntryColl(
             opCtx, NamespaceString::kIndexBuildEntryNamespace, MODE_IX);
 
@@ -1892,7 +1892,7 @@ Status IndexBuildsCoordinator::_setUpIndexBuildForTwoPhaseRecovery(
     // case when an index builds is restarted during recovery.
 
     Lock::DBLock dbLock(opCtx, dbName, MODE_IX);
-    Lock::CollectionLock collLock(opCtx, nssOrUuid, MODE_X);
+    CollectionNamespaceOrUUIDLock collLock(opCtx, nssOrUuid, MODE_X);
     auto collection = CollectionCatalog::get(opCtx)->lookupCollectionByUUID(opCtx, collectionUUID);
     invariant(collection);
     const auto& nss = collection->ns();
@@ -2249,7 +2249,7 @@ void IndexBuildsCoordinator::_cleanUpSinglePhaseAfterFailure(
             unlockRSTL(abortCtx);
 
             const NamespaceStringOrUUID dbAndUUID(replState->dbName, replState->collectionUUID);
-            Lock::CollectionLock collLock(abortCtx, dbAndUUID, MODE_X);
+            CollectionNamespaceOrUUIDLock collLock(abortCtx, dbAndUUID, MODE_X);
             AutoGetCollection indexBuildEntryColl(
                 abortCtx, NamespaceString::kIndexBuildEntryNamespace, MODE_IX);
             _completeSelfAbort(abortCtx, replState, *indexBuildEntryColl, status);
@@ -2289,7 +2289,7 @@ void IndexBuildsCoordinator::_cleanUpTwoPhaseAfterFailure(
                                                          << "; Database: " << replState->dbName));
             }
 
-            Lock::CollectionLock collLock(abortCtx, dbAndUUID, MODE_X);
+            CollectionNamespaceOrUUIDLock collLock(abortCtx, dbAndUUID, MODE_X);
             AutoGetCollection indexBuildEntryColl(
                 abortCtx, NamespaceString::kIndexBuildEntryNamespace, MODE_IX);
             _completeSelfAbort(abortCtx, replState, *indexBuildEntryColl, status);
@@ -2590,7 +2590,7 @@ void IndexBuildsCoordinator::_scanCollectionAndInsertSortedKeysIntoIndex(
 
         Lock::DBLock autoDb(opCtx, replState->dbName, MODE_IX);
         const NamespaceStringOrUUID dbAndUUID(replState->dbName, replState->collectionUUID);
-        Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_IX);
+        CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_IX);
 
         auto collection = _setUpForScanCollectionAndInsertSortedKeysIntoIndex(opCtx, replState);
 
@@ -2609,7 +2609,7 @@ void IndexBuildsCoordinator::_insertSortedKeysIntoIndexForResume(
     {
         Lock::DBLock autoDb(opCtx, replState->dbName, MODE_IX);
         const NamespaceStringOrUUID dbAndUUID(replState->dbName, replState->collectionUUID);
-        Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_IX);
+        CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_IX);
 
         auto collection = _setUpForScanCollectionAndInsertSortedKeysIntoIndex(opCtx, replState);
         uassertStatusOK(_indexBuildsManager.resumeBuildingIndexFromBulkLoadPhase(
@@ -2649,7 +2649,7 @@ void IndexBuildsCoordinator::_insertKeysFromSideTablesWithoutBlockingWrites(
     const NamespaceStringOrUUID dbAndUUID(replState->dbName, replState->collectionUUID);
     {
         Lock::DBLock autoDb(opCtx, replState->dbName, MODE_IX);
-        Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_IX);
+        CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_IX);
 
         uassertStatusOK(_indexBuildsManager.drainBackgroundWrites(
             opCtx,
@@ -2678,7 +2678,7 @@ void IndexBuildsCoordinator::_insertKeysFromSideTablesBlockingWrites(
         // Unlock RSTL to avoid deadlocks with prepare conflicts and state transitions. See
         // SERVER-42621.
         unlockRSTL(opCtx);
-        Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_S);
+        CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_S);
 
         uassertStatusOK(_indexBuildsManager.drainBackgroundWrites(
             opCtx,
@@ -2717,7 +2717,7 @@ IndexBuildsCoordinator::CommitResult IndexBuildsCoordinator::_insertKeysFromSide
 
     // Need to return the collection lock back to exclusive mode to complete the index build.
     const NamespaceStringOrUUID dbAndUUID(replState->dbName, replState->collectionUUID);
-    Lock::CollectionLock collLock(opCtx, dbAndUUID, MODE_X);
+    CollectionNamespaceOrUUIDLock collLock(opCtx, dbAndUUID, MODE_X);
     AutoGetCollection indexBuildEntryColl(
         opCtx, NamespaceString::kIndexBuildEntryNamespace, MODE_IX);