SERVER-63848 Extend AllIndicesRequiredChecker to allow multiple collections

3 files changed, 39 insertions, 28 deletions

diff --git a/src/mongo/db/query/all_indices_required_checker.cpp b/src/mongo/db/query/all_indices_required_checker.cpp
index 26619abe8f4..182d0159447 100644
--- a/src/mongo/db/query/all_indices_required_checker.cpp
+++ b/src/mongo/db/query/all_indices_required_checker.cpp
@@ -31,24 +31,36 @@
 
 namespace mongo {
 
-AllIndicesRequiredChecker::AllIndicesRequiredChecker(const CollectionPtr& collection) {
-    auto allEntriesShared = collection->getIndexCatalog()->getAllReadyEntriesShared();
-    _indexCatalogEntries.reserve(allEntriesShared.size());
-    _indexNames.reserve(allEntriesShared.size());
-    for (auto&& index : allEntriesShared) {
-        _indexCatalogEntries.emplace_back(index);
-        _indexNames.push_back(index->descriptor()->indexName());
+AllIndicesRequiredChecker::AllIndicesRequiredChecker(
+    const MultipleCollectionAccessor& collections) {
+    saveIndicesForCollection(collections.getMainCollection());
+    for (auto& [_, collection] : collections.getSecondaryCollections()) {
+        saveIndicesForCollection(collection);
+    }
+}
+
+void AllIndicesRequiredChecker::saveIndicesForCollection(const CollectionPtr& collection) {
+    if (collection) {
+        auto allEntriesShared = collection->getIndexCatalog()->getAllReadyEntriesShared();
+        auto& indexMap = _indexCatalogEntries[collection->ns()];
+        for (auto&& index : allEntriesShared) {
+            indexMap[index->descriptor()->indexName()] = index;
+        }
     }
 }
 
 void AllIndicesRequiredChecker::check() const {
-    size_t i = 0;
-    for (auto&& index : _indexCatalogEntries) {
-        auto indexCatalogEntry = index.lock();
-        uassert(ErrorCodes::QueryPlanKilled,
-                str::stream() << "query plan killed :: index '" << _indexNames[i] << "' dropped",
-                indexCatalogEntry && !indexCatalogEntry->isDropped());
-        ++i;
+    for (auto& [ns, indexMap] : _indexCatalogEntries) {
+        auto& nsCopy = ns;
+        for (auto& [name, index] : indexMap) {
+            // Create copies for 'ns' and 'name' to make them accessible to the uassert below.
+            auto& nameCopy = name;
+            auto indexCatalogEntry = index.lock();
+            uassert(ErrorCodes::QueryPlanKilled,
+                    str::stream() << "query plan killed :: index '" << nameCopy
+                                  << "' for collection '" << nsCopy << "' dropped",
+                    indexCatalogEntry && !indexCatalogEntry->isDropped());
+        }
     }
 }
 
diff --git a/src/mongo/db/query/all_indices_required_checker.h b/src/mongo/db/query/all_indices_required_checker.h
index 49e9066ee39..d4217dda643 100644
--- a/src/mongo/db/query/all_indices_required_checker.h
+++ b/src/mongo/db/query/all_indices_required_checker.h
@@ -34,6 +34,7 @@
 
 #include "mongo/db/catalog/collection.h"
 #include "mongo/db/catalog/index_catalog_entry.h"
+#include "mongo/db/query/multiple_collection_accessor.h"
 
 namespace mongo {
 
@@ -45,10 +46,10 @@ class AllIndicesRequiredChecker {
 public:
     /**
      * Constructs an 'AllIndicesRequiredChecker' which can be used later to ensure that none of the
-     * indices for the given 'collection' have been dropped. The caller must hold the appropriate
-     * db_raii object in order to read the collection's index catalog.
+     * indices from 'collections' have been dropped. The caller must hold the appropriate db_raii
+     * object in order to read the collection's index catalog.
      */
-    explicit AllIndicesRequiredChecker(const CollectionPtr& collection);
+    explicit AllIndicesRequiredChecker(const MultipleCollectionAccessor& collections);
 
     /**
      * Throws a 'QueryPlanKilled' error if any of the indices which existed at the time of
@@ -57,14 +58,15 @@ public:
     void check() const;
 
 private:
-    // This class holds weak pointers to all of the index catalog entries known at the time of
-    // construction. Later, we can attempt to lock each weak pointer in order to determine whether
-    // an index in the list has been destroyed. If we can lock the weak pointer, we need to check
-    // the 'isDropped()' flag on the index catalog entry.
-    std::vector<std::weak_ptr<const IndexCatalogEntry>> _indexCatalogEntries;
+    void saveIndicesForCollection(const CollectionPtr& collection);
 
-    // The names of the indices above. Used for error reporting.
-    std::vector<std::string> _indexNames;
+    // This map of map holds weak pointers to all of the index catalog entries known at the time of
+    // construction, grouped first by collection namespace then by index name. Later, we can attempt
+    // to lock each weak pointer in order to determine whether an index in the list has been
+    // destroyed. If we can lock the weak pointer, we need to check the 'isDropped()' flag on the
+    // index catalog entry.
+    std::map<NamespaceString, StringMap<std::weak_ptr<const IndexCatalogEntry>>>
+        _indexCatalogEntries;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h
index c360105c01c..6dc49449165 100644
--- a/src/mongo/db/query/sbe_runtime_planner.h
+++ b/src/mongo/db/query/sbe_runtime_planner.h
@@ -87,7 +87,7 @@ public:
           _cq(cq),
           _queryParams(queryParams),
           _yieldPolicy(yieldPolicy),
-          _indexExistenceChecker(collections.getMainCollection()) {
+          _indexExistenceChecker(collections) {
         invariant(_opCtx);
     }
 
@@ -138,9 +138,6 @@ protected:
     const CanonicalQuery& _cq;
     const QueryPlannerParams _queryParams;
     PlanYieldPolicySBE* const _yieldPolicy;
-
-    // TODO SERVER-62913: When support for indexed nested loop join is added, this member needs
-    //  to be extended to support checking for index existence on multiple collections.
     const AllIndicesRequiredChecker _indexExistenceChecker;
 };
 }  // namespace mongo::sbe