1 files changed, 57 insertions, 0 deletions
diff --git a/src/mongo/db/catalog/collection.cpp b/src/mongo/db/catalog/collection.cpp
index e02689c7622..36f73778483 100644
--- a/src/mongo/db/catalog/collection.cpp
+++ b/src/mongo/db/catalog/collection.cpp
@@ -67,6 +67,63 @@ 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, uint64_t)>& _catalogLookup() {
+    static std::function<CollectionPtr(OperationContext*, CollectionUUID, uint64_t)> func;
+    return func;
+}
+
+void CollectionPtr::installCatalogLookupImpl(
+    std::function<CollectionPtr(OperationContext*, CollectionUUID, uint64_t)> impl) {
+    _catalogLookup() = std::move(impl);
+}
+
+CollectionPtr CollectionPtr::null;
+
+CollectionPtr::CollectionPtr() : _collection(nullptr), _opCtx(nullptr) {}
+CollectionPtr::CollectionPtr(OperationContext* opCtx,
+                             const Collection* collection,
+                             uint64_t catalogEpoch)
+    : _collection(collection), _opCtx(opCtx), _catalogEpoch(catalogEpoch) {}
+CollectionPtr::CollectionPtr(const Collection* collection, NoYieldTag)
+    : CollectionPtr(nullptr, collection, 0) {}
+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, _catalogEpoch);
+}
+
+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 {
+    if (_canYield()) {
+        _uuid = _collection->uuid();
+        _ns = _collection->ns();
+        _collection = nullptr;
+    }
+}
+void CollectionPtr::restore() const {
+    if (_canYield()) {
+        // 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, _catalogEpoch);
+        if (coll && coll->ns() == _ns) {
+            _collection = coll.get();
+        }
+    }
+}
+
 // ----
 
 namespace {