SERVER-33610 Recycle 2-phase-locks pending unlock during lock acquisition

5 files changed, 188 insertions, 5 deletions

diff --git a/src/mongo/db/concurrency/lock_manager_defs.h b/src/mongo/db/concurrency/lock_manager_defs.h
index ae8ed410c94..dc36a7bbe32 100644
--- a/src/mongo/db/concurrency/lock_manager_defs.h
+++ b/src/mongo/db/concurrency/lock_manager_defs.h
@@ -418,6 +418,9 @@ struct LockRequest {
     // Written by LockManager on any thread
     // Read by LockManager on any thread
     // Protected by LockHead bucket's mutex
+    // Read by Locker on Locker thread
+    // It is safe for the Locker to read this without taking the bucket mutex provided that the
+    // LockRequest status is not WAITING or CONVERTING.
     LockMode mode;
 
     // This value is different from MODE_NONE only if a conversion is requested for a lock and that
@@ -427,6 +430,18 @@ struct LockRequest {
     // Read by LockManager on any thread
     // Protected by LockHead bucket's mutex
     LockMode convertMode;
+
+    // This unsigned represents the number of pending unlocks for this LockRequest. It is greater
+    // than 0 when the LockRequest is participating in two-phase lock and unlock() is called on it.
+    // It can be greater than 1 if this lock is participating in two-phase-lock and has been
+    // converted to a different mode that also participates in two-phase-lock. unlock() may be
+    // called multiple times on the same resourceId within the same WriteUnitOfWork in this case, so
+    // the number of unlocks() to execute at the end of this WUOW is tracked with this unsigned.
+    //
+    // Written by Locker on Locker thread
+    // Read by Locker on Locker thread
+    // No synchronization
+    unsigned unlockPending = 0;
 };
 
 /**
diff --git a/src/mongo/db/concurrency/lock_state.cpp b/src/mongo/db/concurrency/lock_state.cpp
index 25d3edc9657..17e79b24e61 100644
--- a/src/mongo/db/concurrency/lock_state.cpp
+++ b/src/mongo/db/concurrency/lock_state.cpp
@@ -272,7 +272,7 @@ LockerImpl<IsForMMAPV1>::~LockerImpl() {
     // LockManager may attempt to access deleted memory. Besides it is probably incorrect
     // to delete with unaccounted locks anyways.
     invariant(!inAWriteUnitOfWork());
-    invariant(_resourcesToUnlockAtEndOfUnitOfWork.empty());
+    invariant(_numResourcesToUnlockAtEndUnitOfWork == 0);
     invariant(_requests.empty());
     invariant(_modeForTicket == MODE_NONE);
 
@@ -443,9 +443,19 @@ void LockerImpl<IsForMMAPV1>::endWriteUnitOfWork() {
         return;
     }
 
-    while (!_resourcesToUnlockAtEndOfUnitOfWork.empty()) {
-        unlock(_resourcesToUnlockAtEndOfUnitOfWork.front());
-        _resourcesToUnlockAtEndOfUnitOfWork.pop();
+    LockRequestsMap::Iterator it = _requests.begin();
+    while (_numResourcesToUnlockAtEndUnitOfWork > 0) {
+        if (it->unlockPending) {
+            invariant(!it.finished());
+            _numResourcesToUnlockAtEndUnitOfWork--;
+        }
+        while (it->unlockPending > 0) {
+            // If a lock is converted, unlock() may be called multiple times on a resource within
+            // the same WriteUnitOfWork. All such unlock() requests must thus be fulfilled here.
+            it->unlockPending--;
+            unlock(it.key());
+        }
+        it.next();
     }
 
     // For MMAP V1, we need to yield the flush lock so that the flush thread can run
@@ -482,7 +492,14 @@ template <bool IsForMMAPV1>
 bool LockerImpl<IsForMMAPV1>::unlock(ResourceId resId) {
     LockRequestsMap::Iterator it = _requests.find(resId);
     if (inAWriteUnitOfWork() && _shouldDelayUnlock(it.key(), (it->mode))) {
-        _resourcesToUnlockAtEndOfUnitOfWork.push(it.key());
+        if (!it->unlockPending) {
+            _numResourcesToUnlockAtEndUnitOfWork++;
+        }
+        it->unlockPending++;
+        // unlockPending will only be incremented if a lock is converted and unlock() is called
+        // multiple times on one ResourceId.
+        invariant(it->unlockPending < LockModesCount);
+
         return false;
     }
 
@@ -714,6 +731,17 @@ LockResult LockerImpl<IsForMMAPV1>::lockBegin(ResourceId resId, LockMode mode) {
         isNew = false;
     }
 
+    // If unlockPending is nonzero, that means a LockRequest already exists for this resource but
+    // is planned to be released at the end of this WUOW due to two-phase locking. Rather than
+    // unlocking the existing request, we can reuse it if the existing mode matches the new mode.
+    if (request->unlockPending && isModeCovered(mode, request->mode)) {
+        request->unlockPending--;
+        if (!request->unlockPending) {
+            _numResourcesToUnlockAtEndUnitOfWork--;
+        }
+        return LOCK_OK;
+    }
+
     // Making this call here will record lock re-acquisitions and conversions as well.
     globalStats.recordAcquisition(_id, resId, mode);
     _stats.recordAcquisition(resId, mode);
diff --git a/src/mongo/db/concurrency/lock_state.h b/src/mongo/db/concurrency/lock_state.h
index 02d5d58e3f9..0702e3fe8aa 100644
--- a/src/mongo/db/concurrency/lock_state.h
+++ b/src/mongo/db/concurrency/lock_state.h
@@ -229,6 +229,14 @@ public:
         return lockComplete(nullptr, resId, mode, deadline, checkDeadlock);
     }
 
+    /**
+     * This function is for unit testing only.
+     */
+    FastMapNoAlloc<ResourceId, LockRequest> getRequestsForTest() const {
+        scoped_spinlock scopedLock(_lock);
+        return _requests;
+    }
+
 private:
     friend class AutoYieldFlushLockForMMAPV1Commit;
 
diff --git a/src/mongo/db/concurrency/lock_state_test.cpp b/src/mongo/db/concurrency/lock_state_test.cpp
index a244378fbfe..5e37fdf5e35 100644
--- a/src/mongo/db/concurrency/lock_state_test.cpp
+++ b/src/mongo/db/concurrency/lock_state_test.cpp
@@ -472,4 +472,125 @@ TEST(LockerImpl, GetLockerInfoShouldReportPendingLocks) {
     ASSERT(conflictingLocker.unlockGlobal());
 }
 
+TEST(LockerImpl, ReaquireLockPendingUnlock) {
+    const ResourceId resId(RESOURCE_COLLECTION, "TestDB.collection"_sd);
+
+    DefaultLockerImpl locker;
+    locker.lockGlobal(MODE_IS);
+
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_X));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+
+    locker.beginWriteUnitOfWork();
+
+    ASSERT_FALSE(locker.unlock(resId));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    // Reacquire lock pending unlock.
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_X));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 0);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 0);
+
+    locker.endWriteUnitOfWork();
+
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+
+    locker.unlockGlobal();
+}
+
+TEST(LockerImpl, AcquireLockPendingUnlockWithCoveredMode) {
+    const ResourceId resId(RESOURCE_COLLECTION, "TestDB.collection"_sd);
+
+    DefaultLockerImpl locker;
+    locker.lockGlobal(MODE_IS);
+
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_X));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+
+    locker.beginWriteUnitOfWork();
+
+    ASSERT_FALSE(locker.unlock(resId));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    // Attempt to lock the resource with a mode that is covered by the existing mode.
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_IX));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 0);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 0);
+
+    locker.endWriteUnitOfWork();
+
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_IX));
+
+    locker.unlockGlobal();
+}
+
+TEST(LockerImpl, ConvertLockPendingUnlock) {
+    const ResourceId resId(RESOURCE_COLLECTION, "TestDB.collection"_sd);
+
+    DefaultLockerImpl locker;
+    locker.lockGlobal(MODE_IS);
+
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_IX));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_IX));
+
+    locker.beginWriteUnitOfWork();
+
+    ASSERT_FALSE(locker.unlock(resId));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_IX));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    // Convert lock pending unlock.
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_X));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    locker.endWriteUnitOfWork();
+
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 0);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 0);
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_X));
+
+    locker.unlockGlobal();
+}
+
+TEST(LockerImpl, ConvertLockPendingUnlockAndUnlock) {
+    const ResourceId resId(RESOURCE_COLLECTION, "TestDB.collection"_sd);
+
+    DefaultLockerImpl locker;
+    locker.lockGlobal(MODE_IS);
+
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_IX));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_IX));
+
+    locker.beginWriteUnitOfWork();
+
+    ASSERT_FALSE(locker.unlock(resId));
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_IX));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    // Convert lock pending unlock.
+    ASSERT_EQ(LOCK_OK, locker.lock(resId, MODE_X));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 1);
+
+    // Unlock the lock conversion.
+    ASSERT_FALSE(locker.unlock(resId));
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 1);
+    ASSERT(locker.getRequestsForTest().find(resId).objAddr()->unlockPending == 2);
+
+    locker.endWriteUnitOfWork();
+
+    ASSERT(locker.numResourcesToUnlockAtEndUnitOfWorkForTest() == 0);
+    ASSERT(locker.getRequestsForTest().find(resId).finished());
+    ASSERT_TRUE(locker.isLockHeldForMode(resId, MODE_NONE));
+
+    locker.unlockGlobal();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/concurrency/locker.h b/src/mongo/db/concurrency/locker.h
index 8d98b123f72..8fd9684671e 100644
--- a/src/mongo/db/concurrency/locker.h
+++ b/src/mongo/db/concurrency/locker.h
@@ -390,6 +390,12 @@ public:
     bool shouldAcquireTicket() const {
         return _shouldAcquireTicket;
     }
+    /**
+     * This function is for unit testing only.
+     */
+    unsigned numResourcesToUnlockAtEndUnitOfWorkForTest() const {
+        return _numResourcesToUnlockAtEndUnitOfWork;
+    }
 
 
 protected:
@@ -402,6 +408,11 @@ protected:
      * never interruptible.
      */
     int _uninterruptibleLocksRequested = 0;
+    /**
+     * The number of LockRequests to unlock at the end of this WUOW. This is used for locks
+     * participating in two-phase locking.
+     */
+    unsigned _numResourcesToUnlockAtEndUnitOfWork = 0;
 
 private:
     bool _shouldConflictWithSecondaryBatchApplication = true;