SERVER-50984 Add CollectionPtr to replace usage of const Collection*

It implements a yieldable interface that is used to re-load the
Collection pointer from the catalog after a yield that released locks.
With lock-free reads and copy-on-write on Collection instances releasing
locks without notifying an AutoGetCollection at a higher level may cause
its pointers to dangle if a MODE_X writer installs a new Collection
instance in the catalog.

CollectionPtr should be passed by const reference so a yield can notify
all the way up.

37 files changed, 199 insertions, 145 deletions

diff --git a/src/mongo/db/query/classic_stage_builder.h b/src/mongo/db/query/classic_stage_builder.h
index bdf68d4677b..1c63e6714b5 100644
--- a/src/mongo/db/query/classic_stage_builder.h
+++ b/src/mongo/db/query/classic_stage_builder.h
@@ -39,7 +39,7 @@ namespace mongo::stage_builder {
 class ClassicStageBuilder : public StageBuilder<PlanStage> {
 public:
     ClassicStageBuilder(OperationContext* opCtx,
-                        const Collection* collection,
+                        const CollectionPtr& collection,
                         const CanonicalQuery& cq,
                         const QuerySolution& solution,
                         WorkingSet* ws)
diff --git a/src/mongo/db/query/collection_query_info.cpp b/src/mongo/db/query/collection_query_info.cpp
index 46f9422e759..38df12525f7 100644
--- a/src/mongo/db/query/collection_query_info.cpp
+++ b/src/mongo/db/query/collection_query_info.cpp
@@ -87,7 +87,7 @@ const UpdateIndexData& CollectionQueryInfo::getIndexKeys(OperationContext* opCtx
     return _indexedPaths;
 }
 
-void CollectionQueryInfo::computeIndexKeys(OperationContext* opCtx, const Collection* coll) {
+void CollectionQueryInfo::computeIndexKeys(OperationContext* opCtx, const CollectionPtr& coll) {
     _indexedPaths.clear();
 
     std::unique_ptr<IndexCatalog::IndexIterator> it =
@@ -160,7 +160,7 @@ void CollectionQueryInfo::computeIndexKeys(OperationContext* opCtx, const Collec
 }
 
 void CollectionQueryInfo::notifyOfQuery(OperationContext* opCtx,
-                                        const Collection* coll,
+                                        const CollectionPtr& coll,
                                         const PlanSummaryStats& summaryStats) const {
     auto& collectionIndexUsageTracker =
         CollectionIndexUsageTrackerDecoration::get(coll->getSharedDecorations());
@@ -181,7 +181,7 @@ void CollectionQueryInfo::notifyOfQuery(OperationContext* opCtx,
     }
 }
 
-void CollectionQueryInfo::clearQueryCache(const Collection* coll) const {
+void CollectionQueryInfo::clearQueryCache(const CollectionPtr& coll) const {
     LOGV2_DEBUG(20907,
                 1,
                 "Clearing plan cache - collection info cache reset",
@@ -196,7 +196,7 @@ PlanCache* CollectionQueryInfo::getPlanCache() const {
 }
 
 void CollectionQueryInfo::updatePlanCacheIndexEntries(OperationContext* opCtx,
-                                                      const Collection* coll) {
+                                                      const CollectionPtr& coll) {
     std::vector<CoreIndexInfo> indexCores;
 
     // TODO We shouldn't need to include unfinished indexes, but we must here because the index
@@ -212,7 +212,7 @@ void CollectionQueryInfo::updatePlanCacheIndexEntries(OperationContext* opCtx,
     _planCache->notifyOfIndexUpdates(indexCores);
 }
 
-void CollectionQueryInfo::init(OperationContext* opCtx, const Collection* coll) {
+void CollectionQueryInfo::init(OperationContext* opCtx, const CollectionPtr& coll) {
     const bool includeUnfinishedIndexes = false;
     std::unique_ptr<IndexCatalog::IndexIterator> ii =
         coll->getIndexCatalog()->getIndexIterator(opCtx, includeUnfinishedIndexes);
@@ -226,7 +226,7 @@ void CollectionQueryInfo::init(OperationContext* opCtx, const Collection* coll)
 }
 
 void CollectionQueryInfo::addedIndex(OperationContext* opCtx,
-                                     const Collection* coll,
+                                     const CollectionPtr& coll,
                                      const IndexDescriptor* desc) {
     invariant(desc);
 
@@ -236,14 +236,14 @@ void CollectionQueryInfo::addedIndex(OperationContext* opCtx,
 }
 
 void CollectionQueryInfo::droppedIndex(OperationContext* opCtx,
-                                       const Collection* coll,
+                                       const CollectionPtr& coll,
                                        StringData indexName) {
     rebuildIndexData(opCtx, coll);
     CollectionIndexUsageTrackerDecoration::get(coll->getSharedDecorations())
         .unregisterIndex(indexName);
 }
 
-void CollectionQueryInfo::rebuildIndexData(OperationContext* opCtx, const Collection* coll) {
+void CollectionQueryInfo::rebuildIndexData(OperationContext* opCtx, const CollectionPtr& coll) {
     clearQueryCache(coll);
 
     _keysComputed = false;
diff --git a/src/mongo/db/query/collection_query_info.h b/src/mongo/db/query/collection_query_info.h
index a9fda7742ed..1175dab7837 100644
--- a/src/mongo/db/query/collection_query_info.h
+++ b/src/mongo/db/query/collection_query_info.h
@@ -48,7 +48,14 @@ class CollectionQueryInfo {
 public:
     CollectionQueryInfo();
 
-    inline static const auto get = Collection::declareDecoration<CollectionQueryInfo>();
+    inline static const auto getCollectionQueryInfo =
+        Collection::declareDecoration<CollectionQueryInfo>();
+    static const CollectionQueryInfo& get(const CollectionPtr& collection) {
+        return CollectionQueryInfo::getCollectionQueryInfo(collection.get());
+    }
+    static CollectionQueryInfo& get(Collection* collection) {
+        return CollectionQueryInfo::getCollectionQueryInfo(collection);
+    }
 
     /**
      * Get the PlanCache for this collection.
@@ -63,7 +70,7 @@ public:
     /**
      * Builds internal cache state based on the current state of the Collection's IndexCatalog.
      */
-    void init(OperationContext* opCtx, const Collection* coll);
+    void init(OperationContext* opCtx, const CollectionPtr& coll);
 
     /**
      * Register a newly-created index with the cache.  Must be called whenever an index is
@@ -71,7 +78,9 @@ public:
      *
      * Must be called under exclusive collection lock.
      */
-    void addedIndex(OperationContext* opCtx, const Collection* coll, const IndexDescriptor* desc);
+    void addedIndex(OperationContext* opCtx,
+                    const CollectionPtr& coll,
+                    const IndexDescriptor* desc);
 
     /**
      * Deregister a newly-dropped index with the cache.  Must be called whenever an index is
@@ -79,26 +88,26 @@ public:
      *
      * Must be called under exclusive collection lock.
      */
-    void droppedIndex(OperationContext* opCtx, const Collection* coll, StringData indexName);
+    void droppedIndex(OperationContext* opCtx, const CollectionPtr& coll, StringData indexName);
 
     /**
      * Removes all cached query plans.
      */
-    void clearQueryCache(const Collection* coll) const;
+    void clearQueryCache(const CollectionPtr& coll) const;
 
     void notifyOfQuery(OperationContext* opCtx,
-                       const Collection* coll,
+                       const CollectionPtr& coll,
                        const PlanSummaryStats& summaryStats) const;
 
 private:
-    void computeIndexKeys(OperationContext* opCtx, const Collection* coll);
-    void updatePlanCacheIndexEntries(OperationContext* opCtx, const Collection* coll);
+    void computeIndexKeys(OperationContext* opCtx, const CollectionPtr& coll);
+    void updatePlanCacheIndexEntries(OperationContext* opCtx, const CollectionPtr& coll);
 
     /**
      * Rebuilds cached information that is dependent on index composition. Must be called
      * when index composition changes.
      */
-    void rebuildIndexData(OperationContext* opCtx, const Collection* coll);
+    void rebuildIndexData(OperationContext* opCtx, const CollectionPtr& coll);
 
     // ---  index keys cache
     bool _keysComputed;
diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp
index c4be694ad3d..003c5fe202c 100644
--- a/src/mongo/db/query/explain.cpp
+++ b/src/mongo/db/query/explain.cpp
@@ -591,7 +591,7 @@ void Explain::statsToBSON(const PlanStageStats& stats,
 }
 
 void Explain::generatePlannerInfo(PlanExecutor* exec,
-                                  const Collection* collection,
+                                  const CollectionPtr& collection,
                                   BSONObj extraInfo,
                                   BSONObjBuilder* out) {
     auto planExecImpl = dynamic_cast<PlanExecutorImpl*>(exec);
@@ -795,7 +795,7 @@ void Explain::generateExecutionInfo(PlanExecutor* exec,
 }
 
 void Explain::explainStages(PlanExecutor* exec,
-                            const Collection* collection,
+                            const CollectionPtr& collection,
                             ExplainOptions::Verbosity verbosity,
                             Status executePlanStatus,
                             PlanStageStats* winningPlanTrialStats,
@@ -833,7 +833,7 @@ void Explain::explainPipelineExecutor(PlanExecutorPipeline* exec,
 }
 
 void Explain::explainStages(PlanExecutor* exec,
-                            const Collection* collection,
+                            const CollectionPtr& collection,
                             ExplainOptions::Verbosity verbosity,
                             BSONObj extraInfo,
                             BSONObjBuilder* out) {
@@ -844,6 +844,7 @@ void Explain::explainStages(PlanExecutor* exec,
     auto winningPlanTrialStats = Explain::getWinningPlanTrialStats(exec);
 
     Status executePlanStatus = Status::OK();
+    const CollectionPtr* collectionPtr = &collection;
 
     // If we need execution stats, then run the plan in order to gather the stats.
     if (verbosity >= ExplainOptions::Verbosity::kExecStats) {
@@ -857,12 +858,12 @@ void Explain::explainStages(PlanExecutor* exec,
         // then the collection may no longer be valid. We conservatively set our collection pointer
         // to null in case it is invalid.
         if (executePlanStatus != ErrorCodes::NoQueryExecutionPlans) {
-            collection = nullptr;
+            collectionPtr = &CollectionPtr::null;
         }
     }
 
     explainStages(exec,
-                  collection,
+                  *collectionPtr,
                   verbosity,
                   executePlanStatus,
                   winningPlanTrialStats.get(),
diff --git a/src/mongo/db/query/explain.h b/src/mongo/db/query/explain.h
index 666959883b2..a8848decfc1 100644
--- a/src/mongo/db/query/explain.h
+++ b/src/mongo/db/query/explain.h
@@ -42,6 +42,7 @@
 namespace mongo {
 
 class Collection;
+class CollectionPtr;
 class OperationContext;
 class PlanExecutorPipeline;
 struct PlanSummaryStats;
@@ -72,7 +73,7 @@ public:
      * added to the "executionStats" section of the explain.
      */
     static void explainStages(PlanExecutor* exec,
-                              const Collection* collection,
+                              const CollectionPtr& collection,
                               ExplainOptions::Verbosity verbosity,
                               BSONObj extraInfo,
                               BSONObjBuilder* out);
@@ -92,7 +93,7 @@ public:
      * - 'out' is the builder for the explain output.
      */
     static void explainStages(PlanExecutor* exec,
-                              const Collection* collection,
+                              const CollectionPtr& collection,
                               ExplainOptions::Verbosity verbosity,
                               Status executePlanStatus,
                               PlanStageStats* winningPlanTrialStats,
@@ -205,7 +206,7 @@ private:
      * - 'out' is a builder for the explain output.
      */
     static void generatePlannerInfo(PlanExecutor* exec,
-                                    const Collection* collection,
+                                    const CollectionPtr& collection,
                                     BSONObj extraInfo,
                                     BSONObjBuilder* out);
 
diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp
index fe77d388ff2..27e10256727 100644
--- a/src/mongo/db/query/find.cpp
+++ b/src/mongo/db/query/find.cpp
@@ -82,7 +82,7 @@ MONGO_FAIL_POINT_DEFINE(failReceivedGetmore);
 MONGO_FAIL_POINT_DEFINE(legacyGetMoreWaitWithCursor)
 
 bool shouldSaveCursor(OperationContext* opCtx,
-                      const Collection* collection,
+                      const CollectionPtr& collection,
                       PlanExecutor::ExecState finalState,
                       PlanExecutor* exec) {
     const QueryRequest& qr = exec->getCanonicalQuery()->getQueryRequest();
@@ -121,7 +121,7 @@ void beginQueryOp(OperationContext* opCtx,
 }
 
 void endQueryOp(OperationContext* opCtx,
-                const Collection* collection,
+                const CollectionPtr& collection,
                 const PlanExecutor& exec,
                 long long numResults,
                 CursorId cursorId) {
@@ -410,7 +410,7 @@ Message getMore(OperationContext* opCtx,
 
     PlanExecutor* exec = cursorPin->getExecutor();
     exec->reattachToOperationContext(opCtx);
-    exec->restoreState();
+    exec->restoreState(readLock ? &readLock->getCollection() : nullptr);
 
     auto planSummary = exec->getPlanSummary();
     {
@@ -476,7 +476,7 @@ Message getMore(OperationContext* opCtx,
 
         // Reacquiring locks.
         readLock.emplace(opCtx, nss);
-        exec->restoreState();
+        exec->restoreState(&readLock->getCollection());
 
         // We woke up because either the timed_wait expired, or there was more data. Either way,
         // attempt to generate another batch of results.
@@ -605,8 +605,8 @@ bool runQuery(OperationContext* opCtx,
     LOGV2_DEBUG(20914, 2, "Running query", "query"_attr = redact(cq->toStringShort()));
 
     // Parse, canonicalize, plan, transcribe, and get a plan executor.
-    AutoGetCollectionForReadCommand ctx(opCtx, nss, AutoGetCollectionViewMode::kViewsForbidden);
-    const Collection* const collection = ctx.getCollection();
+    AutoGetCollectionForReadCommand collection(
+        opCtx, nss, AutoGetCollectionViewMode::kViewsForbidden);
     const QueryRequest& qr = cq->getQueryRequest();
 
     opCtx->setExhaust(qr.isExhaust());
@@ -625,7 +625,8 @@ bool runQuery(OperationContext* opCtx,
     // Get the execution plan for the query.
     constexpr auto verbosity = ExplainOptions::Verbosity::kExecAllPlans;
     expCtx->explain = qr.isExplain() ? boost::make_optional(verbosity) : boost::none;
-    auto exec = uassertStatusOK(getExecutorLegacyFind(opCtx, collection, std::move(cq)));
+    auto exec =
+        uassertStatusOK(getExecutorLegacyFind(opCtx, collection.getCollection(), std::move(cq)));
 
     // If it's actually an explain, do the explain and return rather than falling through
     // to the normal query execution loop.
@@ -634,7 +635,8 @@ bool runQuery(OperationContext* opCtx,
         bb.skip(sizeof(QueryResult::Value));
 
         BSONObjBuilder explainBob;
-        Explain::explainStages(exec.get(), collection, verbosity, BSONObj(), &explainBob);
+        Explain::explainStages(
+            exec.get(), collection.getCollection(), verbosity, BSONObj(), &explainBob);
 
         // Add the resulting object to the return buffer.
         BSONObj explainObj = explainBob.obj();
@@ -721,7 +723,7 @@ bool runQuery(OperationContext* opCtx,
     // this cursorid later.
     long long ccId = 0;
 
-    if (shouldSaveCursor(opCtx, collection, state, exec.get())) {
+    if (shouldSaveCursor(opCtx, collection.getCollection(), state, exec.get())) {
         // We won't use the executor until it's getMore'd.
         exec->saveState();
         exec->detachFromOperationContext();
@@ -763,11 +765,15 @@ bool runQuery(OperationContext* opCtx,
             pinnedCursor.getCursor()->setLeftoverMaxTimeMicros(opCtx->getRemainingMaxTimeMicros());
         }
 
-        endQueryOp(opCtx, collection, *pinnedCursor.getCursor()->getExecutor(), numResults, ccId);
+        endQueryOp(opCtx,
+                   collection.getCollection(),
+                   *pinnedCursor.getCursor()->getExecutor(),
+                   numResults,
+                   ccId);
     } else {
         LOGV2_DEBUG(
             20917, 5, "Not caching executor but returning results", "numResults"_attr = numResults);
-        endQueryOp(opCtx, collection, *exec, numResults, ccId);
+        endQueryOp(opCtx, collection.getCollection(), *exec, numResults, ccId);
     }
 
     // Fill out the output buffer's header.
diff --git a/src/mongo/db/query/find.h b/src/mongo/db/query/find.h
index 805fd18884d..6913404cebb 100644
--- a/src/mongo/db/query/find.h
+++ b/src/mongo/db/query/find.h
@@ -50,7 +50,7 @@ class OperationContext;
  * a cursor ID of 0.
  */
 bool shouldSaveCursor(OperationContext* opCtx,
-                      const Collection* collection,
+                      const CollectionPtr& collection,
                       PlanExecutor::ExecState finalState,
                       PlanExecutor* exec);
 
@@ -79,7 +79,7 @@ void beginQueryOp(OperationContext* opCtx,
  * Uses explain functionality to extract stats from 'exec'.
  */
 void endQueryOp(OperationContext* opCtx,
-                const Collection* collection,
+                const CollectionPtr& collection,
                 const PlanExecutor& exec,
                 long long numResults,
                 CursorId cursorId);
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index b13605fb2aa..daecec36c15 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -145,7 +145,7 @@ bool turnIxscanIntoCount(QuerySolution* soln);
 /**
  * Returns 'true' if 'query' on the given 'collection' can be answered using a special IDHACK plan.
  */
-bool isIdHackEligibleQuery(const Collection* collection, const CanonicalQuery& query) {
+bool isIdHackEligibleQuery(const CollectionPtr& collection, const CanonicalQuery& query) {
     return !query.getQueryRequest().showRecordId() && query.getQueryRequest().getHint().isEmpty() &&
         query.getQueryRequest().getMin().isEmpty() && query.getQueryRequest().getMax().isEmpty() &&
         !query.getQueryRequest().getSkip() &&
@@ -247,7 +247,7 @@ IndexEntry indexEntryFromIndexCatalogEntry(OperationContext* opCtx,
  * If query supports index filters, filter params.indices according to any index filters that have
  * been configured. In addition, sets that there were indeed index filters applied.
  */
-void applyIndexFilters(const Collection* collection,
+void applyIndexFilters(const CollectionPtr& collection,
                        const CanonicalQuery& canonicalQuery,
                        QueryPlannerParams* plannerParams) {
     if (!isIdHackEligibleQuery(collection, canonicalQuery)) {
@@ -266,7 +266,7 @@ void applyIndexFilters(const Collection* collection,
 }
 
 void fillOutPlannerParams(OperationContext* opCtx,
-                          const Collection* collection,
+                          const CollectionPtr& collection,
                           CanonicalQuery* canonicalQuery,
                           QueryPlannerParams* plannerParams) {
     invariant(canonicalQuery);
@@ -334,7 +334,7 @@ void fillOutPlannerParams(OperationContext* opCtx,
 }
 
 bool shouldWaitForOplogVisibility(OperationContext* opCtx,
-                                  const Collection* collection,
+                                  const CollectionPtr& collection,
                                   bool tailable) {
 
     // Only non-tailable cursors on the oplog are affected. Only forward cursors, not reverse
@@ -514,7 +514,7 @@ template <typename PlanStageType, typename ResultType>
 class PrepareExecutionHelper {
 public:
     PrepareExecutionHelper(OperationContext* opCtx,
-                           const Collection* collection,
+                           const CollectionPtr& collection,
                            CanonicalQuery* cq,
                            PlanYieldPolicy* yieldPolicy,
                            size_t plannerOptions)
@@ -527,7 +527,7 @@ public:
     }
 
     StatusWith<std::unique_ptr<ResultType>> prepare() {
-        if (nullptr == _collection) {
+        if (!_collection) {
             LOGV2_DEBUG(20921,
                         2,
                         "Collection does not exist. Using EOF plan",
@@ -720,7 +720,7 @@ protected:
         const QueryPlannerParams& plannerParams) = 0;
 
     OperationContext* _opCtx;
-    const Collection* _collection;
+    const CollectionPtr& _collection;
     CanonicalQuery* _cq;
     PlanYieldPolicy* _yieldPolicy;
     const size_t _plannerOptions;
@@ -733,7 +733,7 @@ class ClassicPrepareExecutionHelper final
     : public PrepareExecutionHelper<std::unique_ptr<PlanStage>, ClassicPrepareExecutionResult> {
 public:
     ClassicPrepareExecutionHelper(OperationContext* opCtx,
-                                  const Collection* collection,
+                                  const CollectionPtr& collection,
                                   WorkingSet* ws,
                                   CanonicalQuery* cq,
                                   PlanYieldPolicy* yieldPolicy,
@@ -954,7 +954,7 @@ private:
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getClassicExecutor(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     size_t plannerOptions) {
@@ -986,7 +986,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getClassicExecu
  */
 std::unique_ptr<sbe::RuntimePlanner> makeRuntimePlannerIfNeeded(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     CanonicalQuery* canonicalQuery,
     size_t numSolutions,
     boost::optional<size_t> decisionWorks,
@@ -1051,7 +1051,7 @@ std::unique_ptr<PlanYieldPolicySBE> makeSbeYieldPolicy(
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getSlotBasedExecutor(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> cq,
     PlanYieldPolicy::YieldPolicy requestedYieldPolicy,
     size_t plannerOptions) {
@@ -1100,7 +1100,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getSlotBasedExe
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     size_t plannerOptions) {
@@ -1119,7 +1119,7 @@ namespace {
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> _getExecutorFind(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     size_t plannerOptions) {
@@ -1134,7 +1134,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> _getExecutorFin
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     bool permitYield,
     size_t plannerOptions) {
@@ -1147,7 +1147,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorLegacyFind(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery) {
     return _getExecutorFind(opCtx,
                             collection,
@@ -1206,7 +1206,7 @@ StatusWith<std::unique_ptr<projection_ast::Projection>> makeProjection(const BSO
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete(
     OpDebug* opDebug,
-    const Collection* collection,
+    const CollectionPtr& collection,
     ParsedDelete* parsedDelete,
     boost::optional<ExplainOptions::Verbosity> verbosity) {
     auto expCtx = parsedDelete->expCtx();
@@ -1364,7 +1364,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDele
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate(
     OpDebug* opDebug,
-    const Collection* collection,
+    const CollectionPtr& collection,
     ParsedUpdate* parsedUpdate,
     boost::optional<ExplainOptions::Verbosity> verbosity) {
     auto expCtx = parsedUpdate->expCtx();
@@ -1662,7 +1662,7 @@ bool getDistinctNodeIndex(const std::vector<IndexEntry>& indices,
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const CountCommand& request,
     bool explain,
     const NamespaceString& nss) {
@@ -1929,7 +1929,7 @@ namespace {
 
 // Get the list of indexes that include the "distinct" field.
 QueryPlannerParams fillOutPlannerParamsForDistinct(OperationContext* opCtx,
-                                                   const Collection* collection,
+                                                   const CollectionPtr& collection,
                                                    size_t plannerOptions,
                                                    const ParsedDistinct& parsedDistinct) {
     QueryPlannerParams plannerParams;
@@ -2013,7 +2013,7 @@ QueryPlannerParams fillOutPlannerParamsForDistinct(OperationContext* opCtx,
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorForSimpleDistinct(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const QueryPlannerParams& plannerParams,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     ParsedDistinct* parsedDistinct) {
@@ -2086,7 +2086,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorForS
 // 'strictDistinctOnly' parameter.
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>>
 getExecutorDistinctFromIndexSolutions(OperationContext* opCtx,
-                                      const Collection* collection,
+                                      const CollectionPtr& collection,
                                       std::vector<std::unique_ptr<QuerySolution>> solutions,
                                       PlanYieldPolicy::YieldPolicy yieldPolicy,
                                       ParsedDistinct* parsedDistinct,
@@ -2126,7 +2126,7 @@ getExecutorDistinctFromIndexSolutions(OperationContext* opCtx,
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorWithoutProjection(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const CanonicalQuery* cq,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     size_t plannerOptions) {
@@ -2148,7 +2148,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorWith
 }  // namespace
 
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinct(
-    const Collection* collection, size_t plannerOptions, ParsedDistinct* parsedDistinct) {
+    const CollectionPtr& collection, size_t plannerOptions, ParsedDistinct* parsedDistinct) {
     auto expCtx = parsedDistinct->getQuery()->getExpCtx();
     OperationContext* opCtx = expCtx->opCtx;
     const auto yieldPolicy = opCtx->inMultiDocumentTransaction()
diff --git a/src/mongo/db/query/get_executor.h b/src/mongo/db/query/get_executor.h
index d14e07f9da3..1cdb332dcdc 100644
--- a/src/mongo/db/query/get_executor.h
+++ b/src/mongo/db/query/get_executor.h
@@ -46,6 +46,7 @@
 namespace mongo {
 
 class Collection;
+class CollectionPtr;
 class CountRequest;
 
 /**
@@ -73,7 +74,7 @@ void filterAllowedIndexEntries(const AllowedIndicesFilter& allowedIndicesFilter,
  * 'collection'.  Exposed for testing.
  */
 void fillOutPlannerParams(OperationContext* opCtx,
-                          const Collection* collection,
+                          const CollectionPtr& collection,
                           CanonicalQuery* canonicalQuery,
                           QueryPlannerParams* plannerParams);
 
@@ -106,7 +107,7 @@ IndexEntry indexEntryFromIndexCatalogEntry(OperationContext* opCtx,
  * collection scans on the oplog.
  */
 bool shouldWaitForOplogVisibility(OperationContext* opCtx,
-                                  const Collection* collection,
+                                  const CollectionPtr& collection,
                                   bool tailable);
 
 /**
@@ -119,7 +120,7 @@ bool shouldWaitForOplogVisibility(OperationContext* opCtx,
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     size_t plannerOptions = 0);
@@ -136,7 +137,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor(
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery,
     bool permitYield = false,
     size_t plannerOptions = QueryPlannerParams::DEFAULT);
@@ -146,7 +147,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorLegacyFind(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<CanonicalQuery> canonicalQuery);
 
 /**
@@ -203,7 +204,7 @@ bool turnIxscanIntoDistinctIxscan(QuerySolution* soln,
  * distinct.
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDistinct(
-    const Collection* collection, size_t plannerOptions, ParsedDistinct* parsedDistinct);
+    const CollectionPtr& collection, size_t plannerOptions, ParsedDistinct* parsedDistinct);
 
 /*
  * Get a PlanExecutor for a query executing as part of a count command.
@@ -214,7 +215,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDist
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCount(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const CountCommand& request,
     bool explain,
     const NamespaceString& nss);
@@ -240,7 +241,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCoun
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDelete(
     OpDebug* opDebug,
-    const Collection* collection,
+    const CollectionPtr& collection,
     ParsedDelete* parsedDelete,
     boost::optional<ExplainOptions::Verbosity> verbosity);
 
@@ -266,7 +267,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDele
  */
 StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpdate(
     OpDebug* opDebug,
-    const Collection* collection,
+    const CollectionPtr& collection,
     ParsedUpdate* parsedUpdate,
     boost::optional<ExplainOptions::Verbosity> verbosity);
 }  // namespace mongo
diff --git a/src/mongo/db/query/internal_plans.cpp b/src/mongo/db/query/internal_plans.cpp
index 8d4f5d8ca47..f7c6beeefac 100644
--- a/src/mongo/db/query/internal_plans.cpp
+++ b/src/mongo/db/query/internal_plans.cpp
@@ -51,7 +51,7 @@ namespace mongo {
 std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collectionScan(
     OperationContext* opCtx,
     StringData ns,
-    const Collection* collection,
+    const CollectionPtr& collection,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     const Direction direction,
     boost::optional<RecordId> resumeAfterRecordId) {
@@ -60,7 +60,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collection
     auto expCtx = make_intrusive<ExpressionContext>(
         opCtx, std::unique_ptr<CollatorInterface>(nullptr), NamespaceString(ns));
 
-    if (nullptr == collection) {
+    if (!collection) {
         auto eof = std::make_unique<EOFStage>(expCtx.get());
         // Takes ownership of 'ws' and 'eof'.
         auto statusWithPlanExecutor = plan_executor_factory::make(
@@ -82,7 +82,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::collection
 
 std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWithCollectionScan(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<DeleteStageParams> params,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     Direction direction) {
@@ -106,7 +106,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith
 
 std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::indexScan(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const IndexDescriptor* descriptor,
     const BSONObj& startKey,
     const BSONObj& endKey,
@@ -137,7 +137,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::indexScan(
 
 std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWithIndexScan(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     std::unique_ptr<DeleteStageParams> params,
     const IndexDescriptor* descriptor,
     const BSONObj& startKey,
@@ -172,7 +172,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::deleteWith
 
 std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::updateWithIdHack(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const UpdateStageParams& params,
     const IndexDescriptor* descriptor,
     const BSONObj& key,
@@ -201,7 +201,7 @@ std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> InternalPlanner::updateWith
 std::unique_ptr<PlanStage> InternalPlanner::_collectionScan(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
     WorkingSet* ws,
-    const Collection* collection,
+    const CollectionPtr& collection,
     Direction direction,
     boost::optional<RecordId> resumeAfterRecordId) {
     invariant(collection);
@@ -223,7 +223,7 @@ std::unique_ptr<PlanStage> InternalPlanner::_collectionScan(
 std::unique_ptr<PlanStage> InternalPlanner::_indexScan(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
     WorkingSet* ws,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const IndexDescriptor* descriptor,
     const BSONObj& startKey,
     const BSONObj& endKey,
diff --git a/src/mongo/db/query/internal_plans.h b/src/mongo/db/query/internal_plans.h
index 3846dca76bc..c7507243369 100644
--- a/src/mongo/db/query/internal_plans.h
+++ b/src/mongo/db/query/internal_plans.h
@@ -39,6 +39,7 @@ namespace mongo {
 
 class BSONObj;
 class Collection;
+class CollectionPtr;
 class IndexDescriptor;
 class OperationContext;
 class PlanStage;
@@ -72,7 +73,7 @@ public:
     static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> collectionScan(
         OperationContext* opCtx,
         StringData ns,
-        const Collection* collection,
+        const CollectionPtr& collection,
         PlanYieldPolicy::YieldPolicy yieldPolicy,
         const Direction direction = FORWARD,
         boost::optional<RecordId> resumeAfterRecordId = boost::none);
@@ -82,7 +83,7 @@ public:
      */
     static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> deleteWithCollectionScan(
         OperationContext* opCtx,
-        const Collection* collection,
+        const CollectionPtr& collection,
         std::unique_ptr<DeleteStageParams> params,
         PlanYieldPolicy::YieldPolicy yieldPolicy,
         Direction direction = FORWARD);
@@ -92,7 +93,7 @@ public:
      */
     static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> indexScan(
         OperationContext* opCtx,
-        const Collection* collection,
+        const CollectionPtr& collection,
         const IndexDescriptor* descriptor,
         const BSONObj& startKey,
         const BSONObj& endKey,
@@ -106,7 +107,7 @@ public:
      */
     static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> deleteWithIndexScan(
         OperationContext* opCtx,
-        const Collection* collection,
+        const CollectionPtr& collection,
         std::unique_ptr<DeleteStageParams> params,
         const IndexDescriptor* descriptor,
         const BSONObj& startKey,
@@ -120,7 +121,7 @@ public:
      */
     static std::unique_ptr<PlanExecutor, PlanExecutor::Deleter> updateWithIdHack(
         OperationContext* opCtx,
-        const Collection* collection,
+        const CollectionPtr& collection,
         const UpdateStageParams& params,
         const IndexDescriptor* descriptor,
         const BSONObj& key,
@@ -135,7 +136,7 @@ private:
     static std::unique_ptr<PlanStage> _collectionScan(
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         WorkingSet* ws,
-        const Collection* collection,
+        const CollectionPtr& collection,
         Direction direction,
         boost::optional<RecordId> resumeAfterRecordId = boost::none);
 
@@ -147,7 +148,7 @@ private:
     static std::unique_ptr<PlanStage> _indexScan(
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         WorkingSet* ws,
-        const Collection* collection,
+        const CollectionPtr& collection,
         const IndexDescriptor* descriptor,
         const BSONObj& startKey,
         const BSONObj& endKey,
diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h
index 3fe99fcad2d..60a832e0491 100644
--- a/src/mongo/db/query/plan_executor.h
+++ b/src/mongo/db/query/plan_executor.h
@@ -187,7 +187,7 @@ public:
      * WriteConflictException is encountered. If the time limit is exceeded during this retry
      * process, throws ErrorCodes::MaxTimeMSExpired.
      */
-    virtual void restoreState() = 0;
+    virtual void restoreState(const Yieldable* yieldable) = 0;
 
     /**
      * Detaches from the OperationContext and releases any storage-engine state.
diff --git a/src/mongo/db/query/plan_executor_factory.cpp b/src/mongo/db/query/plan_executor_factory.cpp
index 3630e975646..93ee590d1dc 100644
--- a/src/mongo/db/query/plan_executor_factory.cpp
+++ b/src/mongo/db/query/plan_executor_factory.cpp
@@ -45,7 +45,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     std::unique_ptr<CanonicalQuery> cq,
     std::unique_ptr<WorkingSet> ws,
     std::unique_ptr<PlanStage> rt,
-    const Collection* collection,
+    const CollectionPtr& collection,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     NamespaceString nss,
     std::unique_ptr<QuerySolution> qs) {
@@ -65,7 +65,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
     std::unique_ptr<WorkingSet> ws,
     std::unique_ptr<PlanStage> rt,
-    const Collection* collection,
+    const CollectionPtr& collection,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     NamespaceString nss,
     std::unique_ptr<QuerySolution> qs) {
@@ -87,7 +87,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     std::unique_ptr<QuerySolution> qs,
     std::unique_ptr<CanonicalQuery> cq,
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     PlanYieldPolicy::YieldPolicy yieldPolicy) {
 
@@ -113,7 +113,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     OperationContext* opCtx,
     std::unique_ptr<CanonicalQuery> cq,
     std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) {
 
@@ -142,7 +142,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     OperationContext* opCtx,
     std::unique_ptr<CanonicalQuery> cq,
     std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     std::queue<std::pair<BSONObj, boost::optional<RecordId>>> stash,
     std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) {
diff --git a/src/mongo/db/query/plan_executor_factory.h b/src/mongo/db/query/plan_executor_factory.h
index 56f6fe87cf4..207e3065e20 100644
--- a/src/mongo/db/query/plan_executor_factory.h
+++ b/src/mongo/db/query/plan_executor_factory.h
@@ -65,7 +65,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     std::unique_ptr<CanonicalQuery> cq,
     std::unique_ptr<WorkingSet> ws,
     std::unique_ptr<PlanStage> rt,
-    const Collection* collection,
+    const CollectionPtr& collection,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     NamespaceString nss = NamespaceString(),
     std::unique_ptr<QuerySolution> qs = nullptr);
@@ -81,7 +81,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
     std::unique_ptr<WorkingSet> ws,
     std::unique_ptr<PlanStage> rt,
-    const Collection* collection,
+    const CollectionPtr& collection,
     PlanYieldPolicy::YieldPolicy yieldPolicy,
     NamespaceString nss = NamespaceString(),
     std::unique_ptr<QuerySolution> qs = nullptr);
@@ -93,7 +93,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     std::unique_ptr<QuerySolution> qs,
     std::unique_ptr<CanonicalQuery> cq,
     const boost::intrusive_ptr<ExpressionContext>& expCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     PlanYieldPolicy::YieldPolicy yieldPolicy);
 
@@ -105,7 +105,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     OperationContext* opCtx,
     std::unique_ptr<CanonicalQuery> cq,
     std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     std::unique_ptr<PlanYieldPolicySBE> yieldPolicy);
 
@@ -118,7 +118,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make(
     OperationContext* opCtx,
     std::unique_ptr<CanonicalQuery> cq,
     std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     std::queue<std::pair<BSONObj, boost::optional<RecordId>>> stash,
     std::unique_ptr<PlanYieldPolicySBE> yieldPolicy);
diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index 7be8aa4ef87..915bc2ad34b 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -83,14 +83,15 @@ MONGO_FAIL_POINT_DEFINE(planExecutorHangBeforeShouldWaitForInserts);
  * Constructs a PlanYieldPolicy based on 'policy'.
  */
 std::unique_ptr<PlanYieldPolicy> makeYieldPolicy(PlanExecutorImpl* exec,
-                                                 PlanYieldPolicy::YieldPolicy policy) {
+                                                 PlanYieldPolicy::YieldPolicy policy,
+                                                 const Yieldable* yieldable) {
     switch (policy) {
         case PlanYieldPolicy::YieldPolicy::YIELD_AUTO:
         case PlanYieldPolicy::YieldPolicy::YIELD_MANUAL:
         case PlanYieldPolicy::YieldPolicy::NO_YIELD:
         case PlanYieldPolicy::YieldPolicy::WRITE_CONFLICT_RETRY_ONLY:
         case PlanYieldPolicy::YieldPolicy::INTERRUPT_ONLY: {
-            return std::make_unique<PlanYieldPolicyImpl>(exec, policy);
+            return std::make_unique<PlanYieldPolicyImpl>(exec, policy, yieldable);
         }
         case PlanYieldPolicy::YieldPolicy::ALWAYS_TIME_OUT: {
             return std::make_unique<AlwaysTimeOutYieldPolicy>(exec);
@@ -130,7 +131,7 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx,
                                    unique_ptr<QuerySolution> qs,
                                    unique_ptr<CanonicalQuery> cq,
                                    const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                                   const Collection* collection,
+                                   const CollectionPtr& collection,
                                    NamespaceString nss,
                                    PlanYieldPolicy::YieldPolicy yieldPolicy)
     : _opCtx(opCtx),
@@ -141,8 +142,10 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx,
       _root(std::move(rt)),
       _nss(std::move(nss)),
       // There's no point in yielding if the collection doesn't exist.
-      _yieldPolicy(makeYieldPolicy(
-          this, collection ? yieldPolicy : PlanYieldPolicy::YieldPolicy::NO_YIELD)) {
+      _yieldPolicy(
+          makeYieldPolicy(this,
+                          collection ? yieldPolicy : PlanYieldPolicy::YieldPolicy::NO_YIELD,
+                          collection ? &collection : nullptr)) {
     invariant(!_expCtx || _expCtx->opCtx == _opCtx);
     invariant(!_cq || !_expCtx || _cq->getExpCtx() == _expCtx);
 
@@ -243,12 +246,13 @@ void PlanExecutorImpl::saveState() {
     if (!isMarkedAsKilled()) {
         _root->saveState();
     }
+    _yieldPolicy->setYieldable(nullptr);
     _currentState = kSaved;
 }
 
-void PlanExecutorImpl::restoreState() {
+void PlanExecutorImpl::restoreState(const Yieldable* yieldable) {
     try {
-        restoreStateWithoutRetrying();
+        restoreStateWithoutRetrying(yieldable);
     } catch (const WriteConflictException&) {
         if (!_yieldPolicy->canAutoYield())
             throw;
@@ -258,9 +262,10 @@ void PlanExecutorImpl::restoreState() {
     }
 }
 
-void PlanExecutorImpl::restoreStateWithoutRetrying() {
+void PlanExecutorImpl::restoreStateWithoutRetrying(const Yieldable* yieldable) {
     invariant(_currentState == kSaved);
 
+    _yieldPolicy->setYieldable(yieldable);
     if (!isMarkedAsKilled()) {
         _root->restoreState();
     }
diff --git a/src/mongo/db/query/plan_executor_impl.h b/src/mongo/db/query/plan_executor_impl.h
index 865a20fc515..651937eeb3f 100644
--- a/src/mongo/db/query/plan_executor_impl.h
+++ b/src/mongo/db/query/plan_executor_impl.h
@@ -59,7 +59,7 @@ public:
                      std::unique_ptr<QuerySolution> qs,
                      std::unique_ptr<CanonicalQuery> cq,
                      const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                     const Collection* collection,
+                     const CollectionPtr& collection,
                      NamespaceString nss,
                      PlanYieldPolicy::YieldPolicy yieldPolicy);
 
@@ -68,7 +68,7 @@ public:
     const NamespaceString& nss() const final;
     OperationContext* getOpCtx() const final;
     void saveState() final;
-    void restoreState() final;
+    void restoreState(const Yieldable* yieldable) final;
     void detachFromOperationContext() final;
     void reattachToOperationContext(OperationContext* opCtx) final;
     ExecState getNextDocument(Document* objOut, RecordId* dlOut) final;
@@ -96,7 +96,7 @@ public:
      *
      * This is only public for PlanYieldPolicy. DO NOT CALL ANYWHERE ELSE.
      */
-    void restoreStateWithoutRetrying();
+    void restoreStateWithoutRetrying(const Yieldable* yieldable);
 
     /**
      * Return a pointer to this executor's MultiPlanStage, or nullptr if it does not have one.
diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp
index 3e65c0c46e7..4e8931a00f1 100644
--- a/src/mongo/db/query/plan_executor_sbe.cpp
+++ b/src/mongo/db/query/plan_executor_sbe.cpp
@@ -42,7 +42,7 @@ PlanExecutorSBE::PlanExecutorSBE(
     OperationContext* opCtx,
     std::unique_ptr<CanonicalQuery> cq,
     std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-    const Collection* collection,
+    const CollectionPtr& collection,
     NamespaceString nss,
     bool isOpen,
     boost::optional<std::queue<std::pair<BSONObj, boost::optional<RecordId>>>> stash,
@@ -105,7 +105,7 @@ void PlanExecutorSBE::saveState() {
     _root->saveState();
 }
 
-void PlanExecutorSBE::restoreState() {
+void PlanExecutorSBE::restoreState(const Yieldable* yieldable) {
     invariant(_root);
     _root->restoreState();
 }
diff --git a/src/mongo/db/query/plan_executor_sbe.h b/src/mongo/db/query/plan_executor_sbe.h
index 680f76c578e..0e3ebb3505c 100644
--- a/src/mongo/db/query/plan_executor_sbe.h
+++ b/src/mongo/db/query/plan_executor_sbe.h
@@ -43,7 +43,7 @@ public:
         OperationContext* opCtx,
         std::unique_ptr<CanonicalQuery> cq,
         std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData> root,
-        const Collection* collection,
+        const CollectionPtr& collection,
         NamespaceString nss,
         bool isOpen,
         boost::optional<std::queue<std::pair<BSONObj, boost::optional<RecordId>>>> stash,
@@ -62,7 +62,7 @@ public:
     }
 
     void saveState();
-    void restoreState();
+    void restoreState(const Yieldable* yieldable);
 
     void detachFromOperationContext();
     void reattachToOperationContext(OperationContext* opCtx);
diff --git a/src/mongo/db/query/plan_yield_policy.h b/src/mongo/db/query/plan_yield_policy.h
index 1a10961baa5..8fe09adb8e2 100644
--- a/src/mongo/db/query/plan_yield_policy.h
+++ b/src/mongo/db/query/plan_yield_policy.h
@@ -38,6 +38,7 @@
 namespace mongo {
 
 class ClockSource;
+class Yieldable;
 
 class PlanYieldPolicy {
 public:
@@ -238,6 +239,11 @@ public:
         return _policy;
     }
 
+    /**
+     * Set new yieldable instance if policy supports it.
+     */
+    virtual void setYieldable(const Yieldable* yieldable) {}
+
 private:
     /**
      * Yields locks and calls 'abandonSnapshot()'. Calls 'whileYieldingFn()', if provided, while
diff --git a/src/mongo/db/query/plan_yield_policy_impl.cpp b/src/mongo/db/query/plan_yield_policy_impl.cpp
index 4fc3af37052..a88378d58f3 100644
--- a/src/mongo/db/query/plan_yield_policy_impl.cpp
+++ b/src/mongo/db/query/plan_yield_policy_impl.cpp
@@ -45,20 +45,24 @@ MONGO_FAIL_POINT_DEFINE(setInterruptOnlyPlansCheckForInterruptHang);
 }  // namespace
 
 PlanYieldPolicyImpl::PlanYieldPolicyImpl(PlanExecutorImpl* exec,
-                                         PlanYieldPolicy::YieldPolicy policy)
+                                         PlanYieldPolicy::YieldPolicy policy,
+                                         const Yieldable* yieldable)
     : PlanYieldPolicy(exec->getOpCtx()->lockState()->isGlobalLockedRecursively()
                           ? PlanYieldPolicy::YieldPolicy::NO_YIELD
                           : policy,
                       exec->getOpCtx()->getServiceContext()->getFastClockSource(),
                       internalQueryExecYieldIterations.load(),
                       Milliseconds{internalQueryExecYieldPeriodMS.load()}),
-      _planYielding(exec) {}
+      _planYielding(exec),
+      _yieldable(yieldable) {}
 
 Status PlanYieldPolicyImpl::yield(OperationContext* opCtx, std::function<void()> whileYieldingFn) {
     // Can't use writeConflictRetry since we need to call saveState before reseting the
     // transaction.
     for (int attempt = 1; true; attempt++) {
         try {
+            // Saving and restoring state modifies _yieldable so make a copy before we start
+            const Yieldable* yieldable = _yieldable;
             try {
                 _planYielding->saveState();
             } catch (const WriteConflictException&) {
@@ -70,10 +74,10 @@ Status PlanYieldPolicyImpl::yield(OperationContext* opCtx, std::function<void()>
                 opCtx->recoveryUnit()->abandonSnapshot();
             } else {
                 // Release and reacquire locks.
-                _yieldAllLocks(opCtx, whileYieldingFn, _planYielding->nss());
+                _yieldAllLocks(opCtx, yieldable, whileYieldingFn, _planYielding->nss());
             }
 
-            _planYielding->restoreStateWithoutRetrying();
+            _planYielding->restoreStateWithoutRetrying(yieldable);
             return Status::OK();
         } catch (const WriteConflictException&) {
             CurOp::get(opCtx)->debug().additiveMetrics.incrementWriteConflicts(1);
@@ -89,6 +93,7 @@ Status PlanYieldPolicyImpl::yield(OperationContext* opCtx, std::function<void()>
 }
 
 void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
+                                         const Yieldable* yieldable,
                                          std::function<void()> whileYieldingFn,
                                          const NamespaceString& planExecNS) {
     // Things have to happen here in a specific order:
@@ -101,6 +106,10 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
 
     Locker::LockSnapshot snapshot;
 
+    if (yieldable) {
+        yieldable->yield();
+    }
+
     auto unlocked = locker->saveLockStateAndUnlock(&snapshot);
 
     // Attempt to check for interrupt while locks are not held, in order to discourage the
@@ -129,6 +138,10 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
 
     locker->restoreLockState(opCtx, snapshot);
 
+    if (yieldable) {
+        yieldable->restore();
+    }
+
     // After yielding and reacquiring locks, the preconditions that were used to select our
     // ReadSource initially need to be checked again. Queries hold an AutoGetCollectionForRead RAII
     // lock for their lifetime, which may select a ReadSource based on state (e.g. replication
diff --git a/src/mongo/db/query/plan_yield_policy_impl.h b/src/mongo/db/query/plan_yield_policy_impl.h
index be4163eedec..812c17638d0 100644
--- a/src/mongo/db/query/plan_yield_policy_impl.h
+++ b/src/mongo/db/query/plan_yield_policy_impl.h
@@ -31,14 +31,20 @@
 
 #include "mongo/db/query/plan_executor_impl.h"
 #include "mongo/db/query/plan_yield_policy.h"
+#include "mongo/db/yieldable.h"
 
 namespace mongo {
 
 class PlanYieldPolicyImpl final : public PlanYieldPolicy {
 public:
-    PlanYieldPolicyImpl(PlanExecutorImpl* exec, PlanYieldPolicy::YieldPolicy policy);
+    PlanYieldPolicyImpl(PlanExecutorImpl* exec,
+                        PlanYieldPolicy::YieldPolicy policy,
+                        const Yieldable* yieldable);
 
 private:
+    void setYieldable(const Yieldable* yieldable) override {
+        _yieldable = yieldable;
+    }
     Status yield(OperationContext* opCtx, std::function<void()> whileYieldingFn = nullptr) override;
 
     void preCheckInterruptOnly(OperationContext* opCtx) override;
@@ -52,12 +58,14 @@ private:
      * The whileYieldingFn will be executed after unlocking the locks and before re-acquiring them.
      */
     void _yieldAllLocks(OperationContext* opCtx,
+                        const Yieldable* yieldable,
                         std::function<void()> whileYieldingFn,
                         const NamespaceString& planExecNS);
 
     // The plan executor which this yield policy is responsible for yielding. Must not outlive the
     // plan executor.
     PlanExecutorImpl* const _planYielding;
+    const Yieldable* _yieldable;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/planner_analysis.h b/src/mongo/db/query/planner_analysis.h
index ba1e6a8ac47..e5700483d30 100644
--- a/src/mongo/db/query/planner_analysis.h
+++ b/src/mongo/db/query/planner_analysis.h
@@ -36,6 +36,7 @@
 namespace mongo {
 
 class Collection;
+class CollectionPtr;
 
 class QueryPlannerAnalysis {
 public:
diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 18d7f3543a7..13a9fe6129b 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -1124,7 +1124,7 @@ StatusWith<std::vector<std::unique_ptr<QuerySolution>>> QueryPlanner::plan(
 
 StatusWith<QueryPlanner::SubqueriesPlanningResult> QueryPlanner::planSubqueries(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const PlanCache* planCache,
     const CanonicalQuery& query,
     const QueryPlannerParams& params) {
diff --git a/src/mongo/db/query/query_planner.h b/src/mongo/db/query/query_planner.h
index e9d27bea0f5..236ec4f0269 100644
--- a/src/mongo/db/query/query_planner.h
+++ b/src/mongo/db/query/query_planner.h
@@ -38,6 +38,7 @@ namespace mongo {
 
 class CachedSolution;
 class Collection;
+class CollectionPtr;
 
 /**
  * QueryPlanner's job is to provide an entry point to the query planning and optimization
@@ -106,7 +107,7 @@ public:
      * lists of query solutions in 'SubqueriesPlanningResult'.
      */
     static StatusWith<SubqueriesPlanningResult> planSubqueries(OperationContext* opCtx,
-                                                               const Collection* collection,
+                                                               const CollectionPtr& collection,
                                                                const PlanCache* planCache,
                                                                const CanonicalQuery& query,
                                                                const QueryPlannerParams& params);
diff --git a/src/mongo/db/query/sbe_cached_solution_planner.h b/src/mongo/db/query/sbe_cached_solution_planner.h
index 264ddfbbbfa..9474cca0073 100644
--- a/src/mongo/db/query/sbe_cached_solution_planner.h
+++ b/src/mongo/db/query/sbe_cached_solution_planner.h
@@ -43,7 +43,7 @@ namespace mongo::sbe {
 class CachedSolutionPlanner final : public BaseRuntimePlanner {
 public:
     CachedSolutionPlanner(OperationContext* opCtx,
-                          const Collection* collection,
+                          const CollectionPtr& collection,
                           const CanonicalQuery& cq,
                           const QueryPlannerParams& queryParams,
                           size_t decisionReads,
diff --git a/src/mongo/db/query/sbe_multi_planner.h b/src/mongo/db/query/sbe_multi_planner.h
index 48dca7081bf..80ac325c3b3 100644
--- a/src/mongo/db/query/sbe_multi_planner.h
+++ b/src/mongo/db/query/sbe_multi_planner.h
@@ -43,7 +43,7 @@ namespace mongo::sbe {
 class MultiPlanner final : public BaseRuntimePlanner {
 public:
     MultiPlanner(OperationContext* opCtx,
-                 const Collection* collection,
+                 const CollectionPtr& collection,
                  const CanonicalQuery& cq,
                  PlanCachingMode cachingMode,
                  PlanYieldPolicySBE* yieldPolicy)
diff --git a/src/mongo/db/query/sbe_runtime_planner.cpp b/src/mongo/db/query/sbe_runtime_planner.cpp
index 254cfeb84de..1da89caf8e1 100644
--- a/src/mongo/db/query/sbe_runtime_planner.cpp
+++ b/src/mongo/db/query/sbe_runtime_planner.cpp
@@ -30,6 +30,7 @@
 
 #include "mongo/db/query/sbe_runtime_planner.h"
 
+#include "mongo/db/catalog/collection.h"
 #include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/trial_period_utils.h"
 #include "mongo/db/query/plan_executor_sbe.h"
diff --git a/src/mongo/db/query/sbe_runtime_planner.h b/src/mongo/db/query/sbe_runtime_planner.h
index d399b096877..c7dbfbd73d2 100644
--- a/src/mongo/db/query/sbe_runtime_planner.h
+++ b/src/mongo/db/query/sbe_runtime_planner.h
@@ -60,7 +60,7 @@ public:
 class BaseRuntimePlanner : public RuntimePlanner {
 public:
     BaseRuntimePlanner(OperationContext* opCtx,
-                       const Collection* collection,
+                       const CollectionPtr& collection,
                        const CanonicalQuery& cq,
                        PlanYieldPolicySBE* yieldPolicy)
         : _opCtx(opCtx), _collection(collection), _cq(cq), _yieldPolicy(yieldPolicy) {
@@ -95,7 +95,7 @@ protected:
         std::vector<std::pair<std::unique_ptr<PlanStage>, stage_builder::PlanStageData>> roots);
 
     OperationContext* const _opCtx;
-    const Collection* const _collection;
+    const CollectionPtr& _collection;
     const CanonicalQuery& _cq;
     PlanYieldPolicySBE* const _yieldPolicy;
 };
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index 203b04ef8ff..568f45d2b3f 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -89,7 +89,7 @@ struct PlanStageData {
 class SlotBasedStageBuilder final : public StageBuilder<sbe::PlanStage> {
 public:
     SlotBasedStageBuilder(OperationContext* opCtx,
-                          const Collection* collection,
+                          const CollectionPtr& collection,
                           const CanonicalQuery& cq,
                           const QuerySolution& solution,
                           PlanYieldPolicySBE* yieldPolicy,
diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp
index 8b44fb09c89..bf099dee54d 100644
--- a/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp
@@ -54,13 +54,13 @@ namespace {
  * Checks whether a callback function should be created for a ScanStage and returns it, if so. The
  * logic in the provided callback will be executed when the ScanStage is opened or reopened.
  */
-sbe::ScanOpenCallback makeOpenCallbackIfNeeded(const Collection* collection,
+sbe::ScanOpenCallback makeOpenCallbackIfNeeded(const CollectionPtr& collection,
                                                const CollectionScanNode* csn) {
     if (csn->direction == CollectionScanParams::FORWARD && csn->shouldWaitForOplogVisibility) {
         invariant(!csn->tailable);
         invariant(collection->ns().isOplog());
 
-        return [](OperationContext* opCtx, const Collection* collection, bool reOpen) {
+        return [](OperationContext* opCtx, const CollectionPtr& collection, bool reOpen) {
             if (!reOpen) {
                 // Forward, non-tailable scans from the oplog need to wait until all oplog entries
                 // before the read begins to be visible. This isn't needed for reverse scans because
@@ -87,7 +87,7 @@ sbe::ScanOpenCallback makeOpenCallbackIfNeeded(const Collection* collection,
  * of the same SlotId (the latter is returned purely for convenience purposes).
  */
 std::tuple<std::vector<std::string>, sbe::value::SlotVector, boost::optional<sbe::value::SlotId>>
-makeOplogTimestampSlotsIfNeeded(const Collection* collection,
+makeOplogTimestampSlotsIfNeeded(const CollectionPtr& collection,
                                 sbe::value::SlotIdGenerator* slotIdGenerator,
                                 bool shouldTrackLatestOplogTimestamp) {
     if (shouldTrackLatestOplogTimestamp) {
@@ -118,7 +118,7 @@ std::tuple<sbe::value::SlotId,
            boost::optional<sbe::value::SlotId>,
            std::unique_ptr<sbe::PlanStage>>
 generateOptimizedOplogScan(OperationContext* opCtx,
-                           const Collection* collection,
+                           const CollectionPtr& collection,
                            const CollectionScanNode* csn,
                            sbe::value::SlotIdGenerator* slotIdGenerator,
                            sbe::value::FrameIdGenerator* frameIdGenerator,
@@ -300,7 +300,7 @@ std::tuple<sbe::value::SlotId,
            boost::optional<sbe::value::SlotId>,
            std::unique_ptr<sbe::PlanStage>>
 generateGenericCollScan(OperationContext* opCtx,
-                        const Collection* collection,
+                        const CollectionPtr& collection,
                         const CollectionScanNode* csn,
                         sbe::value::SlotIdGenerator* slotIdGenerator,
                         sbe::value::FrameIdGenerator* frameIdGenerator,
@@ -447,7 +447,7 @@ std::tuple<sbe::value::SlotId,
            boost::optional<sbe::value::SlotId>,
            std::unique_ptr<sbe::PlanStage>>
 generateCollScan(OperationContext* opCtx,
-                 const Collection* collection,
+                 const CollectionPtr& collection,
                  const CollectionScanNode* csn,
                  sbe::value::SlotIdGenerator* slotIdGenerator,
                  sbe::value::FrameIdGenerator* frameIdGenerator,
diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.h b/src/mongo/db/query/sbe_stage_builder_coll_scan.h
index 0fff4cf6212..338c5cff894 100644
--- a/src/mongo/db/query/sbe_stage_builder_coll_scan.h
+++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.h
@@ -53,7 +53,7 @@ std::tuple<sbe::value::SlotId,
            boost::optional<sbe::value::SlotId>,
            std::unique_ptr<sbe::PlanStage>>
 generateCollScan(OperationContext* opCtx,
-                 const Collection* collection,
+                 const CollectionPtr& collection,
                  const CollectionScanNode* csn,
                  sbe::value::SlotIdGenerator* slotIdGenerator,
                  sbe::value::FrameIdGenerator* frameIdGenerator,
diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp
index b7da05ca14b..2e98673bb93 100644
--- a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp
@@ -266,7 +266,7 @@ makeIntervalsFromIndexBounds(const IndexBounds& bounds,
  */
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>
 generateOptimizedMultiIntervalIndexScan(
-    const Collection* collection,
+    const CollectionPtr& collection,
     const std::string& indexName,
     bool forward,
     std::vector<std::pair<std::unique_ptr<KeyString::Value>, std::unique_ptr<KeyString::Value>>>
@@ -386,7 +386,7 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> makeAnchorBranchF
  * consisting of valid recordId's and index seek keys to restart the index scan from.
  */
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>
-makeRecursiveBranchForGenericIndexScan(const Collection* collection,
+makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection,
                                        const std::string& indexName,
                                        const sbe::CheckBoundsParams& params,
                                        sbe::SpoolId spoolId,
@@ -513,7 +513,7 @@ makeRecursiveBranchForGenericIndexScan(const Collection* collection,
  *   - The recursion is terminated when the sspool becomes empty.
  */
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>>
-generateGenericMultiIntervalIndexScan(const Collection* collection,
+generateGenericMultiIntervalIndexScan(const CollectionPtr& collection,
                                       const IndexScanNode* ixn,
                                       KeyString::Version version,
                                       Ordering ordering,
@@ -614,7 +614,7 @@ generateGenericMultiIntervalIndexScan(const Collection* collection,
 }  // namespace
 
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleIntervalIndexScan(
-    const Collection* collection,
+    const CollectionPtr& collection,
     const std::string& indexName,
     bool forward,
     std::unique_ptr<KeyString::Value> lowKey,
@@ -673,7 +673,7 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleInt
 
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateIndexScan(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const IndexScanNode* ixn,
     boost::optional<sbe::value::SlotId> returnKeySlot,
     sbe::value::SlotIdGenerator* slotIdGenerator,
diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.h b/src/mongo/db/query/sbe_stage_builder_index_scan.h
index beb1d983482..538e54dff77 100644
--- a/src/mongo/db/query/sbe_stage_builder_index_scan.h
+++ b/src/mongo/db/query/sbe_stage_builder_index_scan.h
@@ -40,7 +40,7 @@ namespace mongo::stage_builder {
  */
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateIndexScan(
     OperationContext* opCtx,
-    const Collection* collection,
+    const CollectionPtr& collection,
     const IndexScanNode* ixn,
     boost::optional<sbe::value::SlotId> returnKeySlot,
     sbe::value::SlotIdGenerator* slotIdGenerator,
@@ -67,7 +67,7 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateIndexScan
  * in the index.
  */
 std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleIntervalIndexScan(
-    const Collection* collection,
+    const CollectionPtr& collection,
     const std::string& indexName,
     bool forward,
     std::unique_ptr<KeyString::Value> lowKey,
diff --git a/src/mongo/db/query/sbe_sub_planner.h b/src/mongo/db/query/sbe_sub_planner.h
index 1bb465df0f4..5f708a2cd4d 100644
--- a/src/mongo/db/query/sbe_sub_planner.h
+++ b/src/mongo/db/query/sbe_sub_planner.h
@@ -43,7 +43,7 @@ namespace mongo::sbe {
 class SubPlanner final : public BaseRuntimePlanner {
 public:
     SubPlanner(OperationContext* opCtx,
-               const Collection* collection,
+               const CollectionPtr& collection,
                const CanonicalQuery& cq,
                const QueryPlannerParams& queryParams,
                PlanYieldPolicySBE* yieldPolicy)
diff --git a/src/mongo/db/query/stage_builder.h b/src/mongo/db/query/stage_builder.h
index dbee0b6b8b6..7e20618cd83 100644
--- a/src/mongo/db/query/stage_builder.h
+++ b/src/mongo/db/query/stage_builder.h
@@ -41,7 +41,7 @@ template <typename PlanStageType>
 class StageBuilder {
 public:
     StageBuilder(OperationContext* opCtx,
-                 const Collection* collection,
+                 const CollectionPtr& collection,
                  const CanonicalQuery& cq,
                  const QuerySolution& solution)
         : _opCtx(opCtx), _collection(collection), _cq(cq), _solution(solution) {}
@@ -56,7 +56,7 @@ public:
 
 protected:
     OperationContext* _opCtx;
-    const Collection* _collection;
+    const CollectionPtr& _collection;
     const CanonicalQuery& _cq;
     const QuerySolution& _solution;
 };
diff --git a/src/mongo/db/query/stage_builder_util.cpp b/src/mongo/db/query/stage_builder_util.cpp
index 9af708df67b..9b535c20b17 100644
--- a/src/mongo/db/query/stage_builder_util.cpp
+++ b/src/mongo/db/query/stage_builder_util.cpp
@@ -37,7 +37,7 @@
 
 namespace mongo::stage_builder {
 std::unique_ptr<PlanStage> buildClassicExecutableTree(OperationContext* opCtx,
-                                                      const Collection* collection,
+                                                      const CollectionPtr& collection,
                                                       const CanonicalQuery& cq,
                                                       const QuerySolution& solution,
                                                       WorkingSet* ws) {
@@ -54,7 +54,7 @@ std::unique_ptr<PlanStage> buildClassicExecutableTree(OperationContext* opCtx,
 
 std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>
 buildSlotBasedExecutableTree(OperationContext* opCtx,
-                             const Collection* collection,
+                             const CollectionPtr& collection,
                              const CanonicalQuery& cq,
                              const QuerySolution& solution,
                              PlanYieldPolicy* yieldPolicy,
diff --git a/src/mongo/db/query/stage_builder_util.h b/src/mongo/db/query/stage_builder_util.h
index aa6c0abfad2..cd1f77594eb 100644
--- a/src/mongo/db/query/stage_builder_util.h
+++ b/src/mongo/db/query/stage_builder_util.h
@@ -45,14 +45,14 @@ namespace mongo::stage_builder {
  * will consist of.
  */
 std::unique_ptr<PlanStage> buildClassicExecutableTree(OperationContext* opCtx,
-                                                      const Collection* collection,
+                                                      const CollectionPtr& collection,
                                                       const CanonicalQuery& cq,
                                                       const QuerySolution& solution,
                                                       WorkingSet* ws);
 
 std::pair<std::unique_ptr<sbe::PlanStage>, stage_builder::PlanStageData>
 buildSlotBasedExecutableTree(OperationContext* opCtx,
-                             const Collection* collection,
+                             const CollectionPtr& collection,
                              const CanonicalQuery& cq,
                              const QuerySolution& solution,
                              PlanYieldPolicy* yieldPolicy,