SERVER-73903 Add invariant that yieldable Query plans are created with a yieldable CollectionPtr

9 files changed, 31 insertions, 11 deletions

diff --git a/src/mongo/db/catalog/collection.cpp b/src/mongo/db/catalog/collection.cpp
index 4e4b3b409e5..0107be98308 100644
--- a/src/mongo/db/catalog/collection.cpp
+++ b/src/mongo/db/catalog/collection.cpp
@@ -45,22 +45,22 @@ CollectionPtr::CollectionPtr(CollectionPtr&&) = default;
 CollectionPtr::~CollectionPtr() {}
 CollectionPtr& CollectionPtr::operator=(CollectionPtr&&) = default;
 
-bool CollectionPtr::_canYield() const {
+bool CollectionPtr::yieldable() const {
     // We only set the opCtx when this CollectionPtr is yieldable.
-    return _opCtx;
+    return _opCtx || !_collection;
 }
 
 void CollectionPtr::yield() const {
     // Yield if we are yieldable and have a valid collection
-    invariant(_canYield());
     if (_collection) {
+        invariant(_opCtx);
         _yieldedUUID = _collection->uuid();
         _collection = nullptr;
     }
 }
 void CollectionPtr::restore() const {
     // Restore from yield if we are yieldable and if uuid was set in a previous yield.
-    invariant(_canYield());
+    invariant(_opCtx);
     if (_yieldedUUID) {
         // We may only do yield restore when we were holding locks that was yielded so we need to
         // refresh from the catalog to make sure we have a valid collection pointer.
diff --git a/src/mongo/db/catalog/collection.h b/src/mongo/db/catalog/collection.h
index ca61ea91657..c1e8508d90e 100644
--- a/src/mongo/db/catalog/collection.h
+++ b/src/mongo/db/catalog/collection.h
@@ -766,6 +766,7 @@ public:
         *this = CollectionPtr();
     }
 
+    bool yieldable() const override;
     void yield() const override;
     void restore() const override;
 
@@ -779,8 +780,6 @@ public:
     }
 
 private:
-    bool _canYield() const;
-
     // These members needs to be mutable so the yield/restore interface can be const. We don't want
     // yield/restore to require a non-const instance when it otherwise could be const.
     mutable const Collection* _collection = nullptr;
diff --git a/src/mongo/db/db_raii.cpp b/src/mongo/db/db_raii.cpp
index d85146c06d2..e4259bb33a3 100644
--- a/src/mongo/db/db_raii.cpp
+++ b/src/mongo/db/db_raii.cpp
@@ -1321,6 +1321,13 @@ AutoGetCollectionForReadLockFreePITCatalog::AutoGetCollectionForReadLockFreePITC
             _lockFreeReadsBlock.reset();
         }
         _collectionPtr = std::move(collection);
+        // Nested operations should never yield as we don't yield when the global lock is held
+        // recursively. But this is not known when we create the Query plan for this sub operation.
+        // Pretend that we are yieldable but don't allow yield to actually be called.
+        _collectionPtr.makeYieldable(opCtx, [](OperationContext*, UUID) {
+            MONGO_UNREACHABLE;
+            return nullptr;
+        });
     } else {
         auto catalogStateForNamespace =
             acquireCatalogStateForNamespace(opCtx,
diff --git a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
index 8751e488bca..92db73ed2f9 100644
--- a/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
+++ b/src/mongo/db/exec/sbe/sbe_plan_stage_test.h
@@ -268,6 +268,9 @@ protected:
 
 private:
     class MockYieldable : public Yieldable {
+        bool yieldable() const override {
+            return true;
+        }
         void yield() const override {}
         void restore() const override {}
     };
diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp
index dd660ef657e..8243f760328 100644
--- a/src/mongo/db/query/plan_yield_policy.cpp
+++ b/src/mongo/db/query/plan_yield_policy.cpp
@@ -49,7 +49,12 @@ PlanYieldPolicy::PlanYieldPolicy(YieldPolicy policy,
     : _policy(policy),
       _yieldable(yieldable),
       _callbacks(std::move(callbacks)),
-      _elapsedTracker(cs, yieldIterations, yieldPeriod) {}
+      _elapsedTracker(cs, yieldIterations, yieldPeriod) {
+    invariant(!_yieldable || _yieldable->yieldable() ||
+              policy == YieldPolicy::WRITE_CONFLICT_RETRY_ONLY || policy == YieldPolicy::NO_YIELD ||
+              policy == YieldPolicy::INTERRUPT_ONLY || policy == YieldPolicy::ALWAYS_TIME_OUT ||
+              policy == YieldPolicy::ALWAYS_MARK_KILLED);
+}
 
 bool PlanYieldPolicy::shouldYieldOrInterrupt(OperationContext* opCtx) {
     if (_policy == YieldPolicy::INTERRUPT_ONLY) {
diff --git a/src/mongo/db/yieldable.h b/src/mongo/db/yieldable.h
index d420789933d..d9f21518801 100644
--- a/src/mongo/db/yieldable.h
+++ b/src/mongo/db/yieldable.h
@@ -33,6 +33,7 @@ namespace mongo {
 class Yieldable {
 public:
     virtual ~Yieldable() {}
+    virtual bool yieldable() const = 0;
     virtual void yield() const = 0;
     virtual void restore() const = 0;
 };
diff --git a/src/mongo/dbtests/plan_executor_invalidation_test.cpp b/src/mongo/dbtests/plan_executor_invalidation_test.cpp
index 55087fb5e8f..28b2a72e944 100644
--- a/src/mongo/dbtests/plan_executor_invalidation_test.cpp
+++ b/src/mongo/dbtests/plan_executor_invalidation_test.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/client/dbclient_cursor.h"
 #include "mongo/db/catalog/collection.h"
+#include "mongo/db/catalog/collection_yield_restore.h"
 #include "mongo/db/catalog/database.h"
 #include "mongo/db/catalog/index_catalog.h"
 #include "mongo/db/client.h"
@@ -96,7 +97,7 @@ public:
                                         std::move(ws),
                                         std::move(scan),
                                         &collection(),
-                                        PlanYieldPolicy::YieldPolicy::YIELD_MANUAL,
+                                        PlanYieldPolicy::YieldPolicy::NO_YIELD,
                                         QueryPlannerParams::DEFAULT);
 
         ASSERT_OK(statusWithPlanExecutor.getStatus());
@@ -115,7 +116,7 @@ public:
                                           startKey,
                                           endKey,
                                           BoundInclusion::kIncludeBothStartAndEndKeys,
-                                          PlanYieldPolicy::YieldPolicy::YIELD_MANUAL);
+                                          PlanYieldPolicy::YieldPolicy::NO_YIELD);
     }
 
     int N() {
diff --git a/src/mongo/dbtests/query_plan_executor.cpp b/src/mongo/dbtests/query_plan_executor.cpp
index 45872141182..33aebde3cc6 100644
--- a/src/mongo/dbtests/query_plan_executor.cpp
+++ b/src/mongo/dbtests/query_plan_executor.cpp
@@ -101,7 +101,7 @@ public:
     unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeCollScanExec(
         const CollectionPtr* coll,
         BSONObj& filterObj,
-        PlanYieldPolicy::YieldPolicy yieldPolicy = PlanYieldPolicy::YieldPolicy::YIELD_MANUAL,
+        PlanYieldPolicy::YieldPolicy yieldPolicy = PlanYieldPolicy::YieldPolicy::NO_YIELD,
         TailableModeEnum tailableMode = TailableModeEnum::kNormal) {
         CollectionScanParams csparams;
         csparams.direction = CollectionScanParams::FORWARD;
@@ -170,7 +170,7 @@ public:
                                         std::move(ws),
                                         std::move(root),
                                         coll,
-                                        PlanYieldPolicy::YieldPolicy::YIELD_MANUAL,
+                                        PlanYieldPolicy::YieldPolicy::NO_YIELD,
                                         QueryPlannerParams::DEFAULT);
         ASSERT_OK(statusWithPlanExecutor.getStatus());
         return std::move(statusWithPlanExecutor.getValue());
diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp
index b8403c7ca8c..eb17f6b951c 100644
--- a/src/mongo/dbtests/query_stage_sort.cpp
+++ b/src/mongo/dbtests/query_stage_sort.cpp
@@ -35,6 +35,7 @@
 #include "mongo/db/bson/dotted_path_support.h"
 #include "mongo/db/catalog/collection.h"
 #include "mongo/db/catalog/collection_write_path.h"
+#include "mongo/db/catalog/collection_yield_restore.h"
 #include "mongo/db/catalog/database.h"
 #include "mongo/db/client.h"
 #include "mongo/db/db_raii.h"
@@ -355,6 +356,7 @@ public:
             coll = CollectionPtr(db->createCollection(&_opCtx, nss()));
             wuow.commit();
         }
+        coll.makeYieldable(&_opCtx, LockedCollectionYieldRestore(&_opCtx, coll));
 
         fillData();
 
@@ -484,6 +486,7 @@ public:
             coll = CollectionPtr(db->createCollection(&_opCtx, nss()));
             wuow.commit();
         }
+        coll.makeYieldable(&_opCtx, LockedCollectionYieldRestore(&_opCtx, coll));
 
         fillData();
 
@@ -590,6 +593,7 @@ public:
             coll = CollectionPtr(db->createCollection(&_opCtx, nss()));
             wuow.commit();
         }
+        coll.makeYieldable(&_opCtx, LockedCollectionYieldRestore(&_opCtx, coll));
 
         auto ws = std::make_unique<WorkingSet>();
         auto queuedDataStage = std::make_unique<QueuedDataStage>(_expCtx.get(), ws.get());