SERVER-37988 Add an optional predicate argument in forEachCollectionFromDb() to allow skipping unsatisfied collections without grabbing them in the desired lock mode

10 files changed, 215 insertions, 35 deletions

diff --git a/src/mongo/db/catalog/SConscript b/src/mongo/db/catalog/SConscript
index bf49127cdab..11b6f55ebf0 100644
--- a/src/mongo/db/catalog/SConscript
+++ b/src/mongo/db/catalog/SConscript
@@ -339,7 +339,9 @@ env.CppUnitTest(
         'collection_catalog_test.cpp',
     ],
     LIBDEPS=[
+        'catalog_test_fixture',
         'collection_catalog',
+        'collection_catalog_helper',
         '$BUILD_DIR/mongo/db/service_context',
         '$BUILD_DIR/mongo/db/storage/kv/kv_prefix',
     ],
diff --git a/src/mongo/db/catalog/collection_catalog.cpp b/src/mongo/db/catalog/collection_catalog.cpp
index 1bfb90dba4e..3d491c268b3 100644
--- a/src/mongo/db/catalog/collection_catalog.cpp
+++ b/src/mongo/db/catalog/collection_catalog.cpp
@@ -281,6 +281,10 @@ void CollectionCatalog::onOpenCatalog(OperationContext* opCtx) {
 
 Collection* CollectionCatalog::lookupCollectionByUUID(CollectionUUID uuid) const {
     stdx::lock_guard<stdx::mutex> lock(_catalogLock);
+    return _lookupCollectionByUUID(lock, uuid);
+}
+
+Collection* CollectionCatalog::_lookupCollectionByUUID(WithLock, CollectionUUID uuid) const {
     auto foundIt = _catalog.find(uuid);
     return foundIt == _catalog.end() || foundIt->second.collectionPtr == nullptr
         ? nullptr
@@ -298,6 +302,11 @@ Collection* CollectionCatalog::lookupCollectionByNamespace(const NamespaceString
 CollectionCatalogEntry* CollectionCatalog::lookupCollectionCatalogEntryByUUID(
     CollectionUUID uuid) const {
     stdx::lock_guard<stdx::mutex> lock(_catalogLock);
+    return _lookupCollectionCatalogEntryByUUID(lock, uuid);
+}
+
+CollectionCatalogEntry* CollectionCatalog::_lookupCollectionCatalogEntryByUUID(
+    WithLock, CollectionUUID uuid) const {
     auto foundIt = _catalog.find(uuid);
     return foundIt == _catalog.end() ? nullptr : foundIt->second.collectionCatalogEntry.get();
 }
@@ -346,6 +355,21 @@ boost::optional<CollectionUUID> CollectionCatalog::lookupUUIDByNSS(
     return boost::none;
 }
 
+bool CollectionCatalog::checkIfCollectionSatisfiable(CollectionUUID uuid,
+                                                     CollectionInfoFn predicate) const {
+    invariant(predicate);
+
+    stdx::lock_guard<stdx::mutex> lock(_catalogLock);
+    auto collection = _lookupCollectionByUUID(lock, uuid);
+    auto catalogEntry = _lookupCollectionCatalogEntryByUUID(lock, uuid);
+
+    if (!collection || !catalogEntry) {
+        return false;
+    }
+
+    return predicate(collection, catalogEntry);
+}
+
 std::vector<CollectionUUID> CollectionCatalog::getAllCollectionUUIDsFromDb(
     StringData dbName) const {
     stdx::lock_guard<stdx::mutex> lock(_catalogLock);
diff --git a/src/mongo/db/catalog/collection_catalog.h b/src/mongo/db/catalog/collection_catalog.h
index 8a7dcaae650..475d12241fe 100644
--- a/src/mongo/db/catalog/collection_catalog.h
+++ b/src/mongo/db/catalog/collection_catalog.h
@@ -56,6 +56,9 @@ class CollectionCatalog {
     struct CollectionInfo;
 
 public:
+    using CollectionInfoFn = std::function<bool(const Collection* collection,
+                                                const CollectionCatalogEntry* catalogEntry)>;
+
     class iterator {
     public:
         using value_type = Collection*;
@@ -217,6 +220,12 @@ public:
     boost::optional<CollectionUUID> lookupUUIDByNSS(const NamespaceString& nss) const;
 
     /**
+     * Returns whether the collection with 'uuid' satisfies the provided 'predicate'. If the
+     * collection with 'uuid' is not found, false is returned.
+     */
+    bool checkIfCollectionSatisfiable(CollectionUUID uuid, CollectionInfoFn predicate) const;
+
+    /**
      * This function gets the UUIDs of all collections from `dbName`.
      *
      * If the caller does not take a strong database lock, some of UUIDs might no longer exist (due
@@ -284,6 +293,10 @@ private:
     class FinishDropChange;
     friend class CollectionCatalog::iterator;
 
+    Collection* _lookupCollectionByUUID(WithLock, CollectionUUID uuid) const;
+    CollectionCatalogEntry* _lookupCollectionCatalogEntryByUUID(WithLock,
+                                                                CollectionUUID uuid) const;
+
     const std::vector<CollectionUUID>& _getOrdering_inlock(const StringData& db,
                                                            const stdx::lock_guard<stdx::mutex>&);
     mutable mongo::stdx::mutex _catalogLock;
diff --git a/src/mongo/db/catalog/collection_catalog_helper.cpp b/src/mongo/db/catalog/collection_catalog_helper.cpp
index 34b6ce992e1..7719ef531d1 100644
--- a/src/mongo/db/catalog/collection_catalog_helper.cpp
+++ b/src/mongo/db/catalog/collection_catalog_helper.cpp
@@ -36,15 +36,19 @@
 namespace mongo {
 namespace catalog {
 
-void forEachCollectionFromDb(
-    OperationContext* opCtx,
-    StringData dbName,
-    LockMode collLockMode,
-    std::function<bool(Collection* collection, CollectionCatalogEntry* catalogEntry)> callback) {
+void forEachCollectionFromDb(OperationContext* opCtx,
+                             StringData dbName,
+                             LockMode collLockMode,
+                             CollectionCatalog::CollectionInfoFn callback,
+                             CollectionCatalog::CollectionInfoFn predicate) {
 
     CollectionCatalog& catalog = CollectionCatalog::get(opCtx);
     for (auto collectionIt = catalog.begin(dbName); collectionIt != catalog.end(); ++collectionIt) {
         auto uuid = collectionIt.uuid().get();
+        if (predicate && !catalog.checkIfCollectionSatisfiable(uuid, predicate)) {
+            continue;
+        }
+
         auto nss = catalog.lookupNSSByUUID(uuid);
 
         // If the NamespaceString can't be resolved from the uuid, then the collection was dropped.
diff --git a/src/mongo/db/catalog/collection_catalog_helper.h b/src/mongo/db/catalog/collection_catalog_helper.h
index 50df494842e..db9010db85a 100644
--- a/src/mongo/db/catalog/collection_catalog_helper.h
+++ b/src/mongo/db/catalog/collection_catalog_helper.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/db/catalog/collection_catalog.h"
 #include "mongo/db/concurrency/lock_manager_defs.h"
 #include "mongo/db/operation_context.h"
 
@@ -40,14 +41,21 @@ class CollectionCatalogEntry;
 namespace catalog {
 
 /**
- * Looping through all the collections in the database and run callback function on each one of
- * them. The return value of the callback decides whether we should continue the loop.
+ * Iterates through all the collections in the given database and runs the callback function on each
+ * collection. If a predicate is provided, then the callback will only be executed against the
+ * collections that satisfy the predicate.
+ *
+ * Additionally, no collection lock is held while checking the outcome of the predicate. The
+ * predicate must not block, as an internal collection catalog mutex is held during its evaluation.
+ * The collection lock is acquired when executing the callback only on the satisfying collections.
+ *
+ * Iterating through the remaining collections stops when the callback returns false.
  */
-void forEachCollectionFromDb(
-    OperationContext* opCtx,
-    StringData dbName,
-    LockMode collLockMode,
-    std::function<bool(Collection* collection, CollectionCatalogEntry* catalogEntry)> callback);
+void forEachCollectionFromDb(OperationContext* opCtx,
+                             StringData dbName,
+                             LockMode collLockMode,
+                             CollectionCatalog::CollectionInfoFn callback,
+                             CollectionCatalog::CollectionInfoFn predicate = nullptr);
 
 }  // namespace catalog
 }  // namespace mongo
diff --git a/src/mongo/db/catalog/collection_catalog_test.cpp b/src/mongo/db/catalog/collection_catalog_test.cpp
index 7a4f96fd2f5..ea931c9acc9 100644
--- a/src/mongo/db/catalog/collection_catalog_test.cpp
+++ b/src/mongo/db/catalog/collection_catalog_test.cpp
@@ -31,7 +31,9 @@
 #include <algorithm>
 #include <boost/optional/optional_io.hpp>
 
+#include "mongo/db/catalog/catalog_test_fixture.h"
 #include "mongo/db/catalog/collection_catalog_entry_mock.h"
+#include "mongo/db/catalog/collection_catalog_helper.h"
 #include "mongo/db/catalog/collection_mock.h"
 #include "mongo/db/concurrency/lock_manager_defs.h"
 #include "mongo/db/operation_context_noop.h"
@@ -666,4 +668,127 @@ TEST_F(CollectionCatalogTest, GetAllCollectionNamesAndGetAllDbNames) {
 
     catalog.deregisterAllCatalogEntriesAndCollectionObjects();
 }
+
+class ForEachCollectionFromDbTest : public CatalogTestFixture {
+public:
+    void createTestData() {
+        CollectionOptions emptyCollOptions;
+
+        CollectionOptions tempCollOptions;
+        tempCollOptions.temp = true;
+
+        ASSERT_OK(storageInterface()->createCollection(
+            operationContext(), NamespaceString("db", "coll1"), emptyCollOptions));
+        ASSERT_OK(storageInterface()->createCollection(
+            operationContext(), NamespaceString("db", "coll2"), tempCollOptions));
+        ASSERT_OK(storageInterface()->createCollection(
+            operationContext(), NamespaceString("db", "coll3"), tempCollOptions));
+        ASSERT_OK(storageInterface()->createCollection(
+            operationContext(), NamespaceString("db2", "coll4"), emptyCollOptions));
+    }
+};
+
+TEST_F(ForEachCollectionFromDbTest, ForEachCollectionFromDb) {
+    createTestData();
+    auto opCtx = operationContext();
+
+    {
+        auto dbLock = std::make_unique<Lock::DBLock>(opCtx, "db", MODE_IX);
+        int numCollectionsTraversed = 0;
+        catalog::forEachCollectionFromDb(
+            opCtx,
+            "db",
+            MODE_X,
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_X));
+                numCollectionsTraversed++;
+                return true;
+            });
+
+        ASSERT_EQUALS(numCollectionsTraversed, 3);
+    }
+
+    {
+        auto dbLock = std::make_unique<Lock::DBLock>(opCtx, "db2", MODE_IX);
+        int numCollectionsTraversed = 0;
+        catalog::forEachCollectionFromDb(
+            opCtx,
+            "db2",
+            MODE_IS,
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_IS));
+                numCollectionsTraversed++;
+                return true;
+            });
+
+        ASSERT_EQUALS(numCollectionsTraversed, 1);
+    }
+
+    {
+        auto dbLock = std::make_unique<Lock::DBLock>(opCtx, "db3", MODE_IX);
+        int numCollectionsTraversed = 0;
+        catalog::forEachCollectionFromDb(
+            opCtx,
+            "db3",
+            MODE_S,
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                numCollectionsTraversed++;
+                return true;
+            });
+
+        ASSERT_EQUALS(numCollectionsTraversed, 0);
+    }
+}
+
+TEST_F(ForEachCollectionFromDbTest, ForEachCollectionFromDbWithPredicate) {
+    createTestData();
+    auto opCtx = operationContext();
+
+    {
+        auto dbLock = std::make_unique<Lock::DBLock>(opCtx, "db", MODE_IX);
+        int numCollectionsTraversed = 0;
+        catalog::forEachCollectionFromDb(
+            opCtx,
+            "db",
+            MODE_X,
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_X));
+                numCollectionsTraversed++;
+                return true;
+            },
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_NONE));
+                return catalogEntry->getCollectionOptions(opCtx).temp;
+            });
+
+        ASSERT_EQUALS(numCollectionsTraversed, 2);
+    }
+
+    {
+        auto dbLock = std::make_unique<Lock::DBLock>(opCtx, "db", MODE_IX);
+        int numCollectionsTraversed = 0;
+        catalog::forEachCollectionFromDb(
+            opCtx,
+            "db",
+            MODE_IX,
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_IX));
+                numCollectionsTraversed++;
+                return true;
+            },
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+                ASSERT_TRUE(
+                    opCtx->lockState()->isCollectionLockedForMode(collection->ns(), MODE_NONE));
+                return !catalogEntry->getCollectionOptions(opCtx).temp;
+            });
+
+        ASSERT_EQUALS(numCollectionsTraversed, 1);
+    }
+}
+
 }  // namespace
diff --git a/src/mongo/db/catalog/database_impl.cpp b/src/mongo/db/catalog/database_impl.cpp
index ff36de7092e..3ab8a5f536b 100644
--- a/src/mongo/db/catalog/database_impl.cpp
+++ b/src/mongo/db/catalog/database_impl.cpp
@@ -194,32 +194,32 @@ void DatabaseImpl::init(OperationContext* const opCtx) const {
 }
 
 void DatabaseImpl::clearTmpCollections(OperationContext* opCtx) const {
-    invariant(opCtx->lockState()->isDbLockedForMode(name(), MODE_X));
-
-    for (const auto& nss :
-         CollectionCatalog::get(opCtx).getAllCollectionNamesFromDb(opCtx, _name)) {
-        CollectionCatalogEntry* coll =
-            CollectionCatalog::get(opCtx).lookupCollectionCatalogEntryByNamespace(nss);
-        CollectionOptions options = coll->getCollectionOptions(opCtx);
+    invariant(opCtx->lockState()->isDbLockedForMode(name(), MODE_IX));
 
-        if (!options.temp)
-            continue;
+    CollectionCatalog::CollectionInfoFn callback = [&](const Collection* collection,
+                                                       const CollectionCatalogEntry* catalogEntry) {
         try {
-            WriteUnitOfWork wunit(opCtx);
-            Status status = dropCollection(opCtx, nss, {});
-
+            WriteUnitOfWork wuow(opCtx);
+            Status status = dropCollection(opCtx, collection->ns(), {});
             if (!status.isOK()) {
-                warning() << "could not drop temp collection '" << nss << "': " << redact(status);
-                continue;
+                warning() << "could not drop temp collection '" << collection->ns()
+                          << "': " << redact(status);
             }
-
-            wunit.commit();
+            wuow.commit();
         } catch (const WriteConflictException&) {
-            warning() << "could not drop temp collection '" << nss << "' due to "
-                                                                      "WriteConflictException";
+            warning() << "could not drop temp collection '" << collection->ns()
+                      << "' due to WriteConflictException";
             opCtx->recoveryUnit()->abandonSnapshot();
         }
-    }
+        return true;
+    };
+
+    CollectionCatalog::CollectionInfoFn predicate =
+        [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
+            return catalogEntry->getCollectionOptions(opCtx).temp;
+        };
+
+    catalog::forEachCollectionFromDb(opCtx, name(), MODE_X, callback, predicate);
 }
 
 Status DatabaseImpl::setProfilingLevel(OperationContext* opCtx, int newLevel) {
@@ -283,7 +283,7 @@ void DatabaseImpl::getStats(OperationContext* opCtx, BSONObjBuilder* output, dou
         opCtx,
         name(),
         MODE_IS,
-        [&](Collection* collection, CollectionCatalogEntry* catalogEntry) -> bool {
+        [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) -> bool {
             nCollections += 1;
             objects += collection->numRecords(opCtx);
             size += collection->dataSize(opCtx);
diff --git a/src/mongo/db/commands/dbhash.cpp b/src/mongo/db/commands/dbhash.cpp
index 9cb4a5adbd5..d93488bbacf 100644
--- a/src/mongo/db/commands/dbhash.cpp
+++ b/src/mongo/db/commands/dbhash.cpp
@@ -223,7 +223,7 @@ public:
             opCtx,
             dbname,
             MODE_IS,
-            [&](Collection* collection, CollectionCatalogEntry* catalogEntry) {
+            [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
                 auto collNss = collection->ns();
 
                 if (collNss.size() - 1 <= dbname.size()) {
diff --git a/src/mongo/db/commands/list_collections.cpp b/src/mongo/db/commands/list_collections.cpp
index 2e1c8fece35..35a2e0d8008 100644
--- a/src/mongo/db/commands/list_collections.cpp
+++ b/src/mongo/db/commands/list_collections.cpp
@@ -324,7 +324,8 @@ public:
                         opCtx,
                         dbname,
                         MODE_IS,
-                        [&](Collection* collection, CollectionCatalogEntry* catalogEntry) {
+                        [&](const Collection* collection,
+                            const CollectionCatalogEntry* catalogEntry) {
                             if (authorizedCollections &&
                                 (!as->isAuthorizedForAnyActionOnResource(
                                     ResourcePattern::forExactNamespace(collection->ns())))) {
diff --git a/src/mongo/db/storage/kv/kv_storage_engine.cpp b/src/mongo/db/storage/kv/kv_storage_engine.cpp
index d922f598d8e..9931bcf93ee 100644
--- a/src/mongo/db/storage/kv/kv_storage_engine.cpp
+++ b/src/mongo/db/storage/kv/kv_storage_engine.cpp
@@ -915,7 +915,10 @@ int64_t KVStorageEngine::sizeOnDiskForDb(OperationContext* opCtx, StringData dbN
     int64_t size = 0;
 
     catalog::forEachCollectionFromDb(
-        opCtx, dbName, MODE_IS, [&](Collection* collection, CollectionCatalogEntry* catalogEntry) {
+        opCtx,
+        dbName,
+        MODE_IS,
+        [&](const Collection* collection, const CollectionCatalogEntry* catalogEntry) {
             size += catalogEntry->getRecordStore()->storageSize(opCtx);
 
             std::vector<std::string> indexNames;