SERVER-53427 Infrastructure changes to support nested LFR operations

1) Create a nested lock helper to run lock-free if a higher level lock-free operation is already running.
2) Change LockFreeReadsBlock to use a counter rather than a boolean to accommodate out of order lock helper destructors.
3) Only yield lock-free read state in query yield when NOT recursively locked.
4) Change query stages and plan executor to use new nested lock-free lock helper.

9 files changed, 87 insertions, 23 deletions

diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index 002fbaeda87..841b4a8c43f 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -27,6 +27,8 @@
  *    it in the license file.
  */
 
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kStorage
+
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/catalog_raii.h"
@@ -36,6 +38,7 @@
 #include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/database_sharding_state.h"
 #include "mongo/db/views/view_catalog.h"
+#include "mongo/logv2/log.h"
 #include "mongo/util/fail_point.h"
 
 namespace mongo {
@@ -279,6 +282,30 @@ AutoGetCollectionLockFree::AutoGetCollectionLockFree(OperationContext* opCtx,
             !_view || viewMode == AutoGetCollectionViewMode::kViewsPermitted);
 }
 
+AutoGetCollectionMaybeLockFree::AutoGetCollectionMaybeLockFree(
+    OperationContext* opCtx,
+    const NamespaceStringOrUUID& nsOrUUID,
+    LockMode modeColl,
+    AutoGetCollectionViewMode viewMode,
+    Date_t deadline) {
+    if (opCtx->isLockFreeReadsOp()) {
+        _autoGetLockFree.emplace(opCtx,
+                                 nsOrUUID,
+                                 [](std::shared_ptr<const Collection>& collection,
+                                    OperationContext* opCtx,
+                                    CollectionUUID uuid) {
+                                     LOGV2_FATAL(
+                                         5342700,
+                                         "This is a nested lock helper and there was an attempt to "
+                                         "yield locks, which should be impossible");
+                                 },
+                                 viewMode,
+                                 deadline);
+    } else {
+        _autoGet.emplace(opCtx, nsOrUUID, modeColl, viewMode, deadline);
+    }
+}
+
 struct CollectionWriter::SharedImpl {
     SharedImpl(CollectionWriter* parent) : _parent(parent) {}
 
diff --git a/src/mongo/db/catalog_raii.h b/src/mongo/db/catalog_raii.h
index df94dbc2f4b..eddb77fc384 100644
--- a/src/mongo/db/catalog_raii.h
+++ b/src/mongo/db/catalog_raii.h
@@ -296,6 +296,32 @@ private:
 };
 
 /**
+ * This is a nested lock helper. If a higher level operation is running a lock-free read, then this
+ * helper will follow suite and instantiate a AutoGetCollectionLockFree. Otherwise, it will
+ * instantiate a regular AutoGetCollection helper.
+ */
+class AutoGetCollectionMaybeLockFree {
+    AutoGetCollectionMaybeLockFree(const AutoGetCollectionMaybeLockFree&) = delete;
+    AutoGetCollectionMaybeLockFree& operator=(const AutoGetCollectionMaybeLockFree&) = delete;
+
+public:
+    /**
+     * Decides whether to instantiate a lock-free or locked helper based on whether a lock-free
+     * operation is set on the opCtx.
+     */
+    AutoGetCollectionMaybeLockFree(
+        OperationContext* opCtx,
+        const NamespaceStringOrUUID& nsOrUUID,
+        LockMode modeColl,
+        AutoGetCollectionViewMode viewMode = AutoGetCollectionViewMode::kViewsForbidden,
+        Date_t deadline = Date_t::max());
+
+private:
+    boost::optional<AutoGetCollection> _autoGet;
+    boost::optional<AutoGetCollectionLockFree> _autoGetLockFree;
+};
+
+/**
  * RAII-style class to handle the lifetime of writable Collections.
  * It does not take any locks, concurrency needs to be handled separately using explicit locks or
  * AutoGetCollection. This class can serve as an adaptor to unify different methods of acquiring a
diff --git a/src/mongo/db/db_raii.cpp b/src/mongo/db/db_raii.cpp
index 1da574a91c1..ec09f87cfe8 100644
--- a/src/mongo/db/db_raii.cpp
+++ b/src/mongo/db/db_raii.cpp
@@ -418,9 +418,11 @@ AutoGetCollectionForReadLockFree::AutoGetCollectionForReadLockFree(
     AutoGetCollectionViewMode viewMode,
     Date_t deadline)
     : _catalogStash(opCtx) {
-    // Supported lock-free reads should never have an open storage snapshot prior to calling
-    // this helper. The storage snapshot and in-memory state fetched here must be consistent.
-    invariant(supportsLockFreeRead(opCtx) && !opCtx->recoveryUnit()->isActive());
+    // Supported lock-free reads should only ever have an open storage snapshot prior to calling
+    // this helper if it is a nested lock-free operation. The storage snapshot and in-memory state
+    // used across lock=free reads must be consistent.
+    invariant(supportsLockFreeRead(opCtx) &&
+              (!opCtx->recoveryUnit()->isActive() || opCtx->isLockFreeReadsOp()));
 
     EmplaceHelper emplaceFunc(opCtx, _catalogStash, nsOrUUID, viewMode, deadline);
     aquireCollectionAndConsistentSnapshot(
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h
index 92279f05132..6023957a3d2 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.h
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.h
@@ -120,7 +120,7 @@ private:
     std::unique_ptr<SortedDataInterface::Cursor> _cursor;
     std::weak_ptr<const IndexCatalogEntry> _weakIndexCatalogEntry;
     boost::optional<Ordering> _ordering{boost::none};
-    boost::optional<AutoGetCollectionForRead> _coll;
+    boost::optional<AutoGetCollectionForReadMaybeLockFree> _coll;
     boost::optional<KeyStringEntry> _nextRecord;
 
     // This buffer stores values that are projected out of the index entry. Values in the
diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h
index 935366c7da7..71a3a8438e3 100644
--- a/src/mongo/db/exec/sbe/stages/scan.h
+++ b/src/mongo/db/exec/sbe/stages/scan.h
@@ -96,7 +96,7 @@ private:
     bool _open{false};
 
     std::unique_ptr<SeekableRecordCursor> _cursor;
-    boost::optional<AutoGetCollectionForRead> _coll;
+    boost::optional<AutoGetCollectionForReadMaybeLockFree> _coll;
     RecordId _key;
     bool _firstGetNext{false};
 
@@ -179,7 +179,7 @@ private:
     bool _open{false};
 
     std::unique_ptr<SeekableRecordCursor> _cursor;
-    boost::optional<AutoGetCollectionForRead> _coll;
+    boost::optional<AutoGetCollectionForReadMaybeLockFree> _coll;
 };
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index 21c8c777cdf..5868ad76347 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -286,7 +286,7 @@ public:
      * Returns true if the operation is running lock-free.
      */
     bool isLockFreeReadsOp() const {
-        return _isLockFreeReadsOp;
+        return _lockFreeReadOpCount;
     }
 
     /**
@@ -587,10 +587,13 @@ private:
     }
 
     /**
-     * Set whether or not the operation is running lock-free.
+     * Increment a count to indicate that the operation is running lock-free.
      */
-    void setIsLockFreeReadsOp(bool isLockFreeReadsOp) {
-        _isLockFreeReadsOp = isLockFreeReadsOp;
+    void incrementLockFreeReadOpCount() {
+        ++_lockFreeReadOpCount;
+    }
+    void decrementLockFreeReadOpCount() {
+        --_lockFreeReadOpCount;
     }
 
     friend class WriteUnitOfWork;
@@ -648,12 +651,16 @@ private:
     Timer _elapsedTime;
 
     bool _writesAreReplicated = true;
-    bool _isLockFreeReadsOp = false;
     bool _shouldIncrementLatencyStats = true;
     bool _shouldParticipateInFlowControl = true;
     bool _inMultiDocumentTransaction = false;
     bool _isStartingMultiDocumentTransaction = false;
 
+    // Counts how many lock-free read operations are running nested.
+    // Necessary to use a counter rather than a boolean because there is existing code that
+    // destructs lock helpers out of order.
+    int _lockFreeReadOpCount = 0;
+
     // If true, this OpCtx will get interrupted during replica set stepUp and stepDown, regardless
     // of what locks it's taken.
     AtomicWord<bool> _alwaysInterruptAtStepDownOrUp{false};
@@ -705,20 +712,16 @@ class LockFreeReadsBlock {
     LockFreeReadsBlock& operator=(const LockFreeReadsBlock&) = delete;
 
 public:
-    LockFreeReadsBlock(OperationContext* opCtx)
-        : _opCtx(opCtx), _previousLockFreeReadsSetting(opCtx->isLockFreeReadsOp()) {
-        opCtx->setIsLockFreeReadsOp(true);
+    LockFreeReadsBlock(OperationContext* opCtx) : _opCtx(opCtx) {
+        _opCtx->incrementLockFreeReadOpCount();
     }
 
     ~LockFreeReadsBlock() {
-        _opCtx->setIsLockFreeReadsOp(_previousLockFreeReadsSetting);
+        _opCtx->decrementLockFreeReadOpCount();
     }
 
 private:
     OperationContext* _opCtx;
-
-    // Used to re-set the value on the operation context upon destruction that was originally set.
-    const bool _previousLockFreeReadsSetting;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp
index cf87c410051..61ed81b90d8 100644
--- a/src/mongo/db/query/plan_executor_sbe.cpp
+++ b/src/mongo/db/query/plan_executor_sbe.cpp
@@ -195,7 +195,7 @@ PlanExecutor::ExecState PlanExecutorSBE::getNext(BSONObj* out, RecordId* dlOut)
     // insert notifier is necessary for the notifierVersion to advance.
     //
     // Note that we need to hold a database intent lock before acquiring a notifier.
-    boost::optional<AutoGetCollectionForRead> coll;
+    boost::optional<AutoGetCollectionForReadMaybeLockFree> coll;
     insert_listener::CappedInsertNotifierData cappedInsertNotifierData;
     if (insert_listener::shouldListenForInserts(_opCtx, _cq.get())) {
         if (!_opCtx->lockState()->isCollectionLockedForMode(_nss, MODE_IS)) {
diff --git a/src/mongo/db/query/plan_yield_policy_impl.cpp b/src/mongo/db/query/plan_yield_policy_impl.cpp
index c1fdd445472..a6e54317bab 100644
--- a/src/mongo/db/query/plan_yield_policy_impl.cpp
+++ b/src/mongo/db/query/plan_yield_policy_impl.cpp
@@ -104,12 +104,18 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
 
     Locker* locker = opCtx->lockState();
 
-    Locker::LockSnapshot snapshot;
+    if (locker->isGlobalLockedRecursively()) {
+        // No purpose in yielding if the locks are recursively held and cannot be released.
+        return;
+    }
 
+    // Since the locks are not recursively held, this is a top level operation and we can safely
+    // clear the 'yieldable' state before unlocking and then re-establish it after re-locking.
     if (yieldable) {
         yieldable->yield();
     }
 
+    Locker::LockSnapshot snapshot;
     auto unlocked = locker->saveLockStateAndUnlock(&snapshot);
 
     // Attempt to check for interrupt while locks are not held, in order to discourage the
@@ -119,8 +125,8 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
     }
 
     if (!unlocked) {
-        // Nothing was unlocked, just return, yielding is pointless. Restore the yieldable before
-        // returning.
+        // Nothing was unlocked. Recursively held locks are not the only reason locks cannot be
+        // released. Restore the 'yieldable' state before returning.
         if (yieldable) {
             yieldable->restore();
         }
diff --git a/src/mongo/dbtests/query_stage_near.cpp b/src/mongo/dbtests/query_stage_near.cpp
index 3c99ad3dfcf..598b6f9ca2d 100644
--- a/src/mongo/dbtests/query_stage_near.cpp
+++ b/src/mongo/dbtests/query_stage_near.cpp
@@ -91,7 +91,7 @@ protected:
 
     boost::intrusive_ptr<ExpressionContext> _expCtx;
 
-    boost::optional<AutoGetCollectionForRead> _autoColl;
+    boost::optional<AutoGetCollectionForReadMaybeLockFree> _autoColl;
     const IndexDescriptor* _mockGeoIndex;
 };