SERVER-51201 CollectionPtr returned from AutoGetCollectionLockFree is yieldable

9 files changed, 70 insertions, 73 deletions

diff --git a/src/mongo/db/catalog/collection.cpp b/src/mongo/db/catalog/collection.cpp
index 806112ae28b..5a4eeb5c584 100644
--- a/src/mongo/db/catalog/collection.cpp
+++ b/src/mongo/db/catalog/collection.cpp
@@ -67,60 +67,39 @@ bool CappedInsertNotifier::isDead() {
     return _dead;
 }
 
-// We can't reference the catalog from this library as it would create a cyclic library dependency.
-// Setup a weak dependency using a std::function that is installed by the catalog lib
-std::function<CollectionPtr(OperationContext*, CollectionUUID)>& _catalogLookup() {
-    static std::function<CollectionPtr(OperationContext*, CollectionUUID)> func;
-    return func;
-}
-
-void CollectionPtr::installCatalogLookupImpl(
-    std::function<CollectionPtr(OperationContext*, CollectionUUID)> impl) {
-    _catalogLookup() = std::move(impl);
-}
-
 CollectionPtr CollectionPtr::null;
 
 CollectionPtr::CollectionPtr() : _collection(nullptr), _opCtx(nullptr) {}
-CollectionPtr::CollectionPtr(OperationContext* opCtx, const Collection* collection)
-    : _collection(collection), _opCtx(opCtx) {}
+CollectionPtr::CollectionPtr(OperationContext* opCtx,
+                             const Collection* collection,
+                             RestoreFn restoreFn)
+    : _collection(collection), _opCtx(opCtx), _restoreFn(std::move(restoreFn)) {}
 CollectionPtr::CollectionPtr(const Collection* collection, NoYieldTag)
-    : CollectionPtr(nullptr, collection) {}
+    : CollectionPtr(nullptr, collection, nullptr) {}
 CollectionPtr::CollectionPtr(Collection* collection) : CollectionPtr(collection, NoYieldTag{}) {}
-CollectionPtr::CollectionPtr(const std::shared_ptr<const Collection>& collection)
-    : CollectionPtr(collection.get(), NoYieldTag{}) {}
 CollectionPtr::CollectionPtr(CollectionPtr&&) = default;
 CollectionPtr::~CollectionPtr() {}
 CollectionPtr& CollectionPtr::operator=(CollectionPtr&&) = default;
 
-CollectionPtr CollectionPtr::detached() const {
-    return CollectionPtr(_opCtx, _collection);
-}
-
 bool CollectionPtr::_canYield() const {
     // We only set the opCtx when we use a constructor that allows yielding.
-    // When we are doing a lock free read or having a writable pointer in a WUOW it is not allowed
-    // to yield.
     return _opCtx;
 }
 
 void CollectionPtr::yield() const {
     // Yield if we are yieldable and have a valid collection
     if (_canYield() && _collection) {
-        _uuid = _collection->uuid();
+        _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.
-    if (_canYield() && _uuid) {
+    if (_canYield() && _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.
-        auto coll = _catalogLookup()(_opCtx, *_uuid);
-        if (coll) {
-            _collection = coll.get();
-        }
-        _uuid.reset();
+        _collection = _restoreFn(_opCtx, *_yieldedUUID);
+        _yieldedUUID.reset();
     }
 }
 
diff --git a/src/mongo/db/catalog/collection.h b/src/mongo/db/catalog/collection.h
index 58e4f5d75cb..72d1497f362 100644
--- a/src/mongo/db/catalog/collection.h
+++ b/src/mongo/db/catalog/collection.h
@@ -591,24 +591,26 @@ public:
 /**
  * Smart-pointer'esque type to handle yielding of Collection lock that may invalidate pointers when
  * resuming. CollectionPtr will re-load the Collection from the Catalog when restoring from a yield
- * that dropped Collection locks. If this is constructed from a Lock-Free Reads context (shared_ptr)
- * or writable pointer, then it is not allowed to reload from the catalog and the yield operations
- * are no-ops.
+ * that dropped locks. The yield and restore behavior can be disabled by constructing this type from
+ * a writable Collection or by specifying NoYieldTag.
  */
 class CollectionPtr : public Yieldable {
 public:
     static CollectionPtr null;
 
+    // Function for the implementation on how we load a new Collection pointer when restoring from
+    // yield
+    using RestoreFn = std::function<const Collection*(OperationContext*, CollectionUUID)>;
+
     CollectionPtr();
 
     // Creates a Yieldable CollectionPtr that reloads the Collection pointer from the catalog when
     // restoring from yield
-    CollectionPtr(OperationContext* opCtx, const Collection* collection);
+    CollectionPtr(OperationContext* opCtx, const Collection* collection, RestoreFn restoreFn);
 
     // Creates non-yieldable CollectionPtr, performing yield/restore will be a no-op.
     struct NoYieldTag {};
     CollectionPtr(const Collection* collection, NoYieldTag);
-    CollectionPtr(const std::shared_ptr<const Collection>& collection);
     CollectionPtr(Collection* collection);
 
     CollectionPtr(const CollectionPtr&) = delete;
@@ -635,11 +637,6 @@ public:
         return _collection;
     }
 
-    // Creates a new CollectionPtr that is detached from the current, if the current instance is
-    // yieldable the new CollectionPtr will also be. Yielding on the new instance may cause the
-    // instance we detached from to dangle.
-    CollectionPtr detached() const;
-
     void reset() {
         *this = CollectionPtr();
     }
@@ -647,9 +644,6 @@ public:
     void yield() const override;
     void restore() const override;
 
-    static void installCatalogLookupImpl(
-        std::function<CollectionPtr(OperationContext*, CollectionUUID)> impl);
-
     friend std::ostream& operator<<(std::ostream& os, const CollectionPtr& coll);
 
 private:
@@ -658,8 +652,10 @@ private:
     // 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;
-    mutable OptionalCollectionUUID _uuid;
+    mutable OptionalCollectionUUID _yieldedUUID;
+
     OperationContext* _opCtx;
+    RestoreFn _restoreFn;
 };
 
 inline std::ostream& operator<<(std::ostream& os, const CollectionPtr& coll) {
diff --git a/src/mongo/db/catalog/collection_catalog.cpp b/src/mongo/db/catalog/collection_catalog.cpp
index 4db4b069794..ac3104627a1 100644
--- a/src/mongo/db/catalog/collection_catalog.cpp
+++ b/src/mongo/db/catalog/collection_catalog.cpp
@@ -155,15 +155,6 @@ const OperationContext::Decoration<UncommittedWritableCollections>
     getUncommittedWritableCollections =
         OperationContext::declareDecoration<UncommittedWritableCollections>();
 
-struct installCatalogLookupFn {
-    installCatalogLookupFn() {
-        CollectionPtr::installCatalogLookupImpl([](OperationContext* opCtx, CollectionUUID uuid) {
-            const auto& catalog = CollectionCatalog::get(opCtx);
-            return catalog.lookupCollectionByUUID(opCtx, uuid);
-        });
-    }
-} inst;
-
 class FinishDropCollectionChange : public RecoveryUnit::Change {
 public:
     FinishDropCollectionChange(CollectionCatalog* catalog,
@@ -221,7 +212,7 @@ CollectionCatalog::iterator::value_type CollectionCatalog::iterator::operator*()
         return CollectionPtr();
     }
 
-    return {_opCtx, _mapIter->second.get()};
+    return {_opCtx, _mapIter->second.get(), LookupCollectionForYieldRestore()};
 }
 
 Collection* CollectionCatalog::iterator::getWritableCollection(OperationContext* opCtx,
@@ -469,12 +460,14 @@ CollectionPtr CollectionCatalog::lookupCollectionByUUID(OperationContext* opCtx,
     }
 
     if (auto coll = UncommittedCollections::getForTxn(opCtx, uuid)) {
-        return {opCtx, coll.get()};
+        return {opCtx, coll.get(), LookupCollectionForYieldRestore()};
     }
 
     stdx::lock_guard<Latch> lock(_catalogLock);
     auto coll = _lookupCollectionByUUID(lock, uuid);
-    return (coll && coll->isCommitted()) ? CollectionPtr(opCtx, coll.get()) : CollectionPtr();
+    return (coll && coll->isCommitted())
+        ? CollectionPtr(opCtx, coll.get(), LookupCollectionForYieldRestore())
+        : CollectionPtr();
 }
 
 void CollectionCatalog::makeCollectionVisible(CollectionUUID uuid) {
@@ -565,13 +558,15 @@ CollectionPtr CollectionCatalog::lookupCollectionByNamespace(OperationContext* o
     }
 
     if (auto coll = UncommittedCollections::getForTxn(opCtx, nss)) {
-        return {opCtx, coll.get()};
+        return {opCtx, coll.get(), LookupCollectionForYieldRestore()};
     }
 
     stdx::lock_guard<Latch> lock(_catalogLock);
     auto it = _collections.find(nss);
     auto coll = (it == _collections.end() ? nullptr : it->second);
-    return (coll && coll->isCommitted()) ? CollectionPtr(opCtx, coll.get()) : nullptr;
+    return (coll && coll->isCommitted())
+        ? CollectionPtr(opCtx, coll.get(), LookupCollectionForYieldRestore())
+        : nullptr;
 }
 
 boost::optional<NamespaceString> CollectionCatalog::lookupNSSByUUID(OperationContext* opCtx,
@@ -937,4 +932,9 @@ void CollectionCatalog::discardUnmanagedClone(OperationContext* opCtx, Collectio
     uncommittedWritableCollections.remove(collection);
 }
 
+const Collection* LookupCollectionForYieldRestore::operator()(OperationContext* opCtx,
+                                                              CollectionUUID uuid) const {
+    return CollectionCatalog::get(opCtx).lookupCollectionByUUID(opCtx, uuid).get();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/catalog/collection_catalog.h b/src/mongo/db/catalog/collection_catalog.h
index 6d36d220e0d..bf32a3c0de5 100644
--- a/src/mongo/db/catalog/collection_catalog.h
+++ b/src/mongo/db/catalog/collection_catalog.h
@@ -436,4 +436,13 @@ private:
      */
     DatabaseProfileSettingsMap _databaseProfileSettings;
 };
+
+/**
+ * Functor for looking up Collection by UUID from the Collection Catalog. This is the default yield
+ * restore implementation for CollectionPtr when acquired from the catalog.
+ */
+struct LookupCollectionForYieldRestore {
+    const Collection* operator()(OperationContext* opCtx, CollectionUUID uuid) const;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index e8bb4e48e74..d0a91958cc3 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -151,26 +151,31 @@ Collection* AutoGetCollection::getWritableCollection(CollectionCatalog::Lifetime
     // Acquire writable instance if not already available
     if (!_writableColl) {
 
-        // Resets the writable Collection when the write unit of work finishes so we re-fetches and
-        // re-clones the Collection if a new write unit of work is opened.
+        // Makes the internal CollectionPtr Yieldable and resets the writable Collection when the
+        // write unit of work finishes so we re-fetches and re-clones the Collection if a new write
+        // unit of work is opened.
         class WritableCollectionReset : public RecoveryUnit::Change {
         public:
             WritableCollectionReset(AutoGetCollection& autoColl,
-                                    const CollectionPtr& rollbackCollection)
-                : _autoColl(autoColl), _rollbackCollection(rollbackCollection.get()) {}
+                                    const Collection* originalCollection)
+                : _autoColl(autoColl), _originalCollection(originalCollection) {}
             void commit(boost::optional<Timestamp> commitTime) final {
-                // Restore coll to a yieldable collection
-                _autoColl._coll = {_autoColl.getOperationContext(), _autoColl._coll.get()};
+                _autoColl._coll = CollectionPtr(_autoColl.getOperationContext(),
+                                                _autoColl._coll.get(),
+                                                LookupCollectionForYieldRestore());
                 _autoColl._writableColl = nullptr;
             }
             void rollback() final {
-                _autoColl._coll = {_autoColl.getOperationContext(), _rollbackCollection};
+                _autoColl._coll = CollectionPtr(_autoColl.getOperationContext(),
+                                                _originalCollection,
+                                                LookupCollectionForYieldRestore());
                 _autoColl._writableColl = nullptr;
             }
 
         private:
             AutoGetCollection& _autoColl;
-            const Collection* _rollbackCollection;
+            // Used to be able to restore to the original pointer in case of a rollback
+            const Collection* _originalCollection;
         };
 
         auto& catalog = CollectionCatalog::get(_opCtx);
@@ -178,7 +183,7 @@ Collection* AutoGetCollection::getWritableCollection(CollectionCatalog::Lifetime
             catalog.lookupCollectionByNamespaceForMetadataWrite(_opCtx, mode, _resolvedNss);
         if (mode == CollectionCatalog::LifetimeMode::kManagedInWriteUnitOfWork) {
             _opCtx->recoveryUnit()->registerChange(
-                std::make_unique<WritableCollectionReset>(*this, _coll));
+                std::make_unique<WritableCollectionReset>(*this, _coll.get()));
         }
 
         // Set to writable collection. We are no longer yieldable.
@@ -201,7 +206,14 @@ AutoGetCollectionLockFree::AutoGetCollectionLockFree(OperationContext* opCtx,
     _resolvedNss = CollectionCatalog::get(opCtx).resolveNamespaceStringOrUUID(opCtx, nsOrUUID);
     _collection =
         CollectionCatalog::get(opCtx).lookupCollectionByNamespaceForRead(opCtx, _resolvedNss);
-    _collectionPtr = CollectionPtr(_collection);
+
+    // When we restore from yield on this CollectionPtr we will update _collection above and use its
+    // new pointer in the CollectionPtr
+    _collectionPtr = CollectionPtr(
+        opCtx, _collection.get(), [this](OperationContext* opCtx, CollectionUUID uuid) {
+            _collection = CollectionCatalog::get(opCtx).lookupCollectionByUUIDForRead(opCtx, uuid);
+            return _collection.get();
+        });
 
     // TODO (SERVER-51319): add DatabaseShardingState::checkDbVersion somewhere.
 
diff --git a/src/mongo/db/catalog_raii.h b/src/mongo/db/catalog_raii.h
index 4b0ddc7bf0c..a4ba77ee797 100644
--- a/src/mongo/db/catalog_raii.h
+++ b/src/mongo/db/catalog_raii.h
@@ -113,7 +113,7 @@ public:
         Date_t deadline = Date_t::max());
 
     explicit operator bool() const {
-        return static_cast<bool>(_coll);
+        return static_cast<bool>(getCollection());
     }
 
     /**
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index 0c5a689182b..74d1f2ddcab 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -2453,8 +2453,9 @@ void IndexBuildsCoordinator::_runIndexBuildInner(
     // dropped while the index build is still registered for the collection -- until abortIndexBuild
     // is called. The collection can be renamed, but it is OK for the name to be stale just for
     // logging purposes.
-    auto collection = CollectionCatalog::get(opCtx).lookupCollectionByUUIDForRead(
+    auto collectionSharedPtr = CollectionCatalog::get(opCtx).lookupCollectionByUUIDForRead(
         opCtx, replState->collectionUUID);
+    CollectionPtr collection(collectionSharedPtr.get(), CollectionPtr::NoYieldTag{});
     invariant(collection,
               str::stream() << "Collection with UUID " << replState->collectionUUID
                             << " should exist because an index build is in progress: "
diff --git a/src/mongo/db/repl/local_oplog_info.cpp b/src/mongo/db/repl/local_oplog_info.cpp
index e73199c00db..eb7de479b6d 100644
--- a/src/mongo/db/repl/local_oplog_info.cpp
+++ b/src/mongo/db/repl/local_oplog_info.cpp
@@ -86,11 +86,11 @@ const CollectionPtr& LocalOplogInfo::getCollection() const {
 }
 
 void LocalOplogInfo::setCollection(const CollectionPtr& oplog) {
-    _oplog = oplog.detached();
+    _oplog = CollectionPtr(oplog.get(), CollectionPtr::NoYieldTag{});
 }
 
 void LocalOplogInfo::resetCollection() {
-    _oplog = CollectionPtr();
+    _oplog.reset();
 }
 
 void LocalOplogInfo::setNewTimestamp(ServiceContext* service, const Timestamp& newTime) {
diff --git a/src/mongo/dbtests/querytests.cpp b/src/mongo/dbtests/querytests.cpp
index 2bd9ba0a18c..377aee1f160 100644
--- a/src/mongo/dbtests/querytests.cpp
+++ b/src/mongo/dbtests/querytests.cpp
@@ -129,7 +129,7 @@ protected:
             wunit.commit();
         }
         abortOnExit.dismiss();
-        _collection = collection.get().detached();
+        _collection = CollectionPtr(collection.get().get(), CollectionPtr::NoYieldTag{});
     }
 
     void insert(const char* s) {