SERVER-23667 can overtake locks by using the same lockSessionID as the lock owner, via DistLockManager::lockWithSessionID

6 files changed, 167 insertions, 31 deletions

diff --git a/src/mongo/s/catalog/dist_lock_manager.h b/src/mongo/s/catalog/dist_lock_manager.h
index e06b2b9eac8..3b60c4a09bd 100644
--- a/src/mongo/s/catalog/dist_lock_manager.h
+++ b/src/mongo/s/catalog/dist_lock_manager.h
@@ -138,6 +138,22 @@ public:
         stdx::chrono::milliseconds lockTryInterval = kDefaultLockRetryInterval) = 0;
 
     /**
+     * Same behavior as lock(...) above, except takes a specific lock session ID "lockSessionID"
+     * instead of randomly generating one internally.
+     *
+     * This is useful for a process running on the config primary after a failover. A lock can be
+     * immediately reacquired if "lockSessionID" matches that of the lock, rather than waiting for
+     * the inactive lock to expire.
+     */
+    virtual StatusWith<ScopedDistLock> lockWithSessionID(
+        OperationContext* txn,
+        StringData name,
+        StringData whyMessage,
+        const OID lockSessionID,
+        stdx::chrono::milliseconds waitFor = kDefaultLockTimeout,
+        stdx::chrono::milliseconds lockTryInterval = kDefaultLockRetryInterval) = 0;
+
+    /**
      * Makes a best-effort attempt to unlock all locks owned by the given processID.
      * Only implemented for the ReplSetDistLockManager and only used after catalog manager swap
      * during upgrade to CSRS.
diff --git a/src/mongo/s/catalog/dist_lock_manager_mock.cpp b/src/mongo/s/catalog/dist_lock_manager_mock.cpp
index 9a493708c54..28ac22cb411 100644
--- a/src/mongo/s/catalog/dist_lock_manager_mock.cpp
+++ b/src/mongo/s/catalog/dist_lock_manager_mock.cpp
@@ -73,6 +73,17 @@ StatusWith<DistLockManager::ScopedDistLock> DistLockManagerMock::lock(
     StringData whyMessage,
     milliseconds waitFor,
     milliseconds lockTryInterval) {
+    return lockWithSessionID(
+        txn, name, whyMessage, DistLockHandle::gen(), waitFor, lockTryInterval);
+}
+
+StatusWith<DistLockManager::ScopedDistLock> DistLockManagerMock::lockWithSessionID(
+    OperationContext* txn,
+    StringData name,
+    StringData whyMessage,
+    const OID lockSessionID,
+    milliseconds waitFor,
+    milliseconds lockTryInterval) {
     _lockChecker(name, whyMessage, waitFor, lockTryInterval);
     _lockChecker = NoLockFuncSet;
 
@@ -89,7 +100,7 @@ StatusWith<DistLockManager::ScopedDistLock> DistLockManagerMock::lock(
 
     LockInfo info;
     info.name = name.toString();
-    info.lockID = DistLockHandle::gen();
+    info.lockID = lockSessionID;
     _locks.push_back(info);
 
     return DistLockManager::ScopedDistLock(nullptr, info.lockID, this);
diff --git a/src/mongo/s/catalog/dist_lock_manager_mock.h b/src/mongo/s/catalog/dist_lock_manager_mock.h
index 75b7b8a720f..21ddef99f93 100644
--- a/src/mongo/s/catalog/dist_lock_manager_mock.h
+++ b/src/mongo/s/catalog/dist_lock_manager_mock.h
@@ -54,6 +54,14 @@ public:
         stdx::chrono::milliseconds waitFor,
         stdx::chrono::milliseconds lockTryInterval) override;
 
+    virtual StatusWith<DistLockManager::ScopedDistLock> lockWithSessionID(
+        OperationContext* txn,
+        StringData name,
+        StringData whyMessage,
+        const OID lockSessionID,
+        stdx::chrono::milliseconds waitFor,
+        stdx::chrono::milliseconds lockTryInterval) override;
+
     virtual void unlockAll(OperationContext* txn, const std::string& processID) override;
 
     using LockFunc = stdx::function<void(StringData name,
diff --git a/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp b/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp
index 49defc445d9..4f4457c793b 100644
--- a/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp
+++ b/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp
@@ -168,9 +168,9 @@ void ReplSetDistLockManager::doTask() {
     }
 }
 
-StatusWith<bool> ReplSetDistLockManager::canOvertakeLock(OperationContext* txn,
-                                                         LocksType lockDoc,
-                                                         const milliseconds& lockExpiration) {
+StatusWith<bool> ReplSetDistLockManager::isLockExpired(OperationContext* txn,
+                                                       LocksType lockDoc,
+                                                       const milliseconds& lockExpiration) {
     const auto& processID = lockDoc.getProcess();
     auto pingStatus = _catalog->getPing(txn, processID);
 
@@ -267,6 +267,16 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lock(
     StringData whyMessage,
     milliseconds waitFor,
     milliseconds lockTryInterval) {
+    return lockWithSessionID(txn, name, whyMessage, OID::gen(), waitFor, lockTryInterval);
+}
+
+StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lockWithSessionID(
+    OperationContext* txn,
+    StringData name,
+    StringData whyMessage,
+    const OID lockSessionID,
+    milliseconds waitFor,
+    milliseconds lockTryInterval) {
     Timer timer(_serviceContext->getTickSource());
     Timer msgTimer(_serviceContext->getTickSource());
 
@@ -280,7 +290,6 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lock(
     // until the lockTryInterval has been reached. If a network error occurs at each lock
     // acquisition attempt, the lock acquisition will be retried immediately.
     while (waitFor <= milliseconds::zero() || milliseconds(timer.millis()) < waitFor) {
-        const OID lockSessionID = OID::gen();
         const string who = str::stream() << _processID << ":" << getThreadName();
 
         auto lockExpiration = _lockExpiration;
@@ -350,13 +359,13 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lock(
         // found, use the normal grab lock path to acquire it.
         if (getLockStatusResult.isOK()) {
             auto currentLock = getLockStatusResult.getValue();
-            auto canOvertakeResult = canOvertakeLock(txn, currentLock, lockExpiration);
+            auto isLockExpiredResult = isLockExpired(txn, currentLock, lockExpiration);
 
-            if (!canOvertakeResult.isOK()) {
-                return canOvertakeResult.getStatus();
+            if (!isLockExpiredResult.isOK()) {
+                return isLockExpiredResult.getStatus();
             }
 
-            if (canOvertakeResult.getValue()) {
+            if (isLockExpiredResult.getValue() || (lockSessionID == currentLock.getLockID())) {
                 auto overtakeResult = _catalog->overtakeLock(txn,
                                                              name,
                                                              lockSessionID,
diff --git a/src/mongo/s/catalog/replset/replset_dist_lock_manager.h b/src/mongo/s/catalog/replset/replset_dist_lock_manager.h
index df3d8c6b855..2f876f8e693 100644
--- a/src/mongo/s/catalog/replset/replset_dist_lock_manager.h
+++ b/src/mongo/s/catalog/replset/replset_dist_lock_manager.h
@@ -75,6 +75,14 @@ public:
         stdx::chrono::milliseconds waitFor,
         stdx::chrono::milliseconds lockTryInterval) override;
 
+    virtual StatusWith<ScopedDistLock> lockWithSessionID(
+        OperationContext* txn,
+        StringData name,
+        StringData whyMessage,
+        const OID lockSessionID,
+        stdx::chrono::milliseconds waitFor,
+        stdx::chrono::milliseconds lockTryInterval) override;
+
     virtual void unlockAll(OperationContext* txn, const std::string& processID) override;
 
 protected:
@@ -102,9 +110,9 @@ private:
      * Returns true if the current process that owns the lock has no fresh pings since
      * the lock expiration threshold.
      */
-    StatusWith<bool> canOvertakeLock(OperationContext* txn,
-                                     const LocksType lockDoc,
-                                     const stdx::chrono::milliseconds& lockExpiration);
+    StatusWith<bool> isLockExpired(OperationContext* txn,
+                                   const LocksType lockDoc,
+                                   const stdx::chrono::milliseconds& lockExpiration);
 
     //
     // All member variables are labeled with one of the following codes indicating the
diff --git a/src/mongo/s/catalog/replset/replset_dist_lock_manager_test.cpp b/src/mongo/s/catalog/replset/replset_dist_lock_manager_test.cpp
index 83d90fa1e20..999386de0ff 100644
--- a/src/mongo/s/catalog/replset/replset_dist_lock_manager_test.cpp
+++ b/src/mongo/s/catalog/replset/replset_dist_lock_manager_test.cpp
@@ -34,6 +34,8 @@
 #include <string>
 #include <type_traits>
 #include <vector>
+#include <boost/optional.hpp>
+#include <boost/optional/optional_io.hpp>
 
 #include "mongo/base/status.h"
 #include "mongo/base/status_with.h"
@@ -104,7 +106,7 @@ public:
     }
 
     /**
-     * Get the process id that was initialiezd with the lock manager being tested.
+     * Get the process id that was initialized with the lock manager being tested.
      */
     string getProcessID() const {
         return _processID;
@@ -236,14 +238,14 @@ TEST_F(ReplSetDistLockManagerFixture, BasicLockLifeCycle) {
 /**
  * Test scenario:
  * 1. Grab lock fails up to 3 times.
- * 2. Check that each attempt uses a unique lock session id.
+ * 2. Check that each subsequent attempt uses the same lock session id.
  * 3. Unlock (on destructor of ScopedDistLock).
  * 4. Check lock id used in lock and unlock are the same.
  */
 TEST_F(RSDistLockMgrWithMockTickSource, LockSuccessAfterRetry) {
     string lockName("test");
     string me("me");
-    OID lastTS;
+    boost::optional<OID> lastTS;
     Date_t lastTime(Date_t::now());
     string whyMsg("because");
 
@@ -274,8 +276,10 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockSuccessAfterRetry) {
                        Date_t time,
                        StringData why) {
             ASSERT_EQUALS(lockName, lockID);
-            // Every attempt should have a unique sesssion ID.
-            ASSERT_NOT_EQUALS(lastTS, lockSessionID);
+            // Lock session ID should be the same after first attempt.
+            if (lastTS) {
+                ASSERT_EQUALS(lastTS, lockSessionID);
+            }
             ASSERT_EQUALS(getProcessID(), processId);
             ASSERT_GREATER_THAN_OR_EQUALS(time, lastTime);
             ASSERT_EQUALS(whyMsg, why);
@@ -294,9 +298,10 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockSuccessAfterRetry) {
                                                                         Date_t time,
                                                                         StringData why) {
                         ASSERT_EQUALS(lockName, lockID);
-                        // Every attempt should have a unique sesssion ID.
-                        ASSERT_NOT_EQUALS(lastTS, lockSessionID);
-                        lastTS = lockSessionID;
+                        // Lock session ID should be the same after first attempt.
+                        if (lastTS) {
+                            ASSERT_EQUALS(lastTS, lockSessionID);
+                        }
                         ASSERT_TRUE(lockSessionID.isSet());
                         ASSERT_EQUALS(getProcessID(), processId);
                         ASSERT_GREATER_THAN_OR_EQUALS(time, lastTime);
@@ -366,14 +371,14 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockSuccessAfterRetry) {
 /**
  * Test scenario:
  * 1. Grab lock fails up to 3 times.
- * 2. Check that each attempt uses a unique lock session id.
+ * 2. Check that each subsequent attempt uses the same lock session id.
  * 3. Grab lock errors out on the fourth try.
  * 4. Make sure that unlock is called to cleanup the last lock attempted that error out.
  */
 TEST_F(RSDistLockMgrWithMockTickSource, LockFailsAfterRetry) {
     string lockName("test");
     string me("me");
-    OID lastTS;
+    boost::optional<OID> lastTS;
     Date_t lastTime(Date_t::now());
     string whyMsg("because");
 
@@ -389,8 +394,10 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockFailsAfterRetry) {
             Date_t time,
             StringData why) {
             ASSERT_EQUALS(lockName, lockID);
-            // Every attempt should have a unique sesssion ID.
-            ASSERT_NOT_EQUALS(lastTS, lockSessionID);
+            // Lock session ID should be the same after first attempt.
+            if (lastTS) {
+                ASSERT_EQUALS(lastTS, lockSessionID);
+            }
             ASSERT_EQUALS(getProcessID(), processId);
             ASSERT_GREATER_THAN_OR_EQUALS(time, lastTime);
             ASSERT_EQUALS(whyMsg, why);
@@ -409,8 +416,10 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockFailsAfterRetry) {
                                                                         Date_t time,
                                                                         StringData why) {
                         ASSERT_EQUALS(lockName, lockID);
-                        // Every attempt should have a unique sesssion ID.
-                        ASSERT_NOT_EQUALS(lastTS, lockSessionID);
+                        // Lock session ID should be the same after first attempt.
+                        if (lastTS) {
+                            ASSERT_EQUALS(lastTS, lockSessionID);
+                        }
                         lastTS = lockSessionID;
                         ASSERT_TRUE(lockSessionID.isSet());
                         ASSERT_EQUALS(getProcessID(), processId);
@@ -491,7 +500,7 @@ TEST_F(ReplSetDistLockManagerFixture, LockBusyNoRetry) {
 /**
  * Test scenario:
  * 1. Attempt to grab lock.
- * 2. Check that each attempt uses a unique lock session id.
+ * 2. Check that each subsequent attempt uses the same lock session id.
  * 3. Times out trying.
  * 4. Checks result is error.
  * 5. Implicitly check that unlock is not called (default setting of mock catalog).
@@ -499,7 +508,7 @@ TEST_F(ReplSetDistLockManagerFixture, LockBusyNoRetry) {
 TEST_F(RSDistLockMgrWithMockTickSource, LockRetryTimeout) {
     string lockName("test");
     string me("me");
-    OID lastTS;
+    boost::optional<OID> lastTS;
     Date_t lastTime(Date_t::now());
     string whyMsg("because");
 
@@ -513,8 +522,10 @@ TEST_F(RSDistLockMgrWithMockTickSource, LockRetryTimeout) {
                                                                            Date_t time,
                                                                            StringData why) {
             ASSERT_EQUALS(lockName, lockID);
-            // Every attempt should have a unique sesssion ID.
-            ASSERT_NOT_EQUALS(lastTS, lockSessionID);
+            // Lock session ID should be the same after first attempt.
+            if (lastTS) {
+                ASSERT_EQUALS(lastTS, lockSessionID);
+            }
             ASSERT_EQUALS(getProcessID(), processId);
             ASSERT_GREATER_THAN_OR_EQUALS(time, lastTime);
             ASSERT_EQUALS(whyMsg, why);
@@ -974,7 +985,7 @@ TEST_F(ReplSetDistLockManagerFixture, CheckLockStatusError) {
  * 5. 2nd attempt to grab lock still fails for the same reason.
  * 6. But since the ping is not fresh anymore, dist lock manager should overtake lock.
  */
-TEST_F(ReplSetDistLockManagerFixture, BasicLockOvertaking) {
+TEST_F(ReplSetDistLockManagerFixture, LockOvertakingAfterLockExpiration) {
     OID lastTS;
 
     getMockCatalog()->expectGrabLock(
@@ -1054,6 +1065,79 @@ TEST_F(ReplSetDistLockManagerFixture, BasicLockOvertaking) {
     ASSERT_EQUALS(lastTS, unlockSessionIDPassed);
 }
 
+/**
+ * Test scenario:
+ * 1. Attempt to grab lock with lockSessionID fails because lock is already owned.
+ * 2. Then the the lock is overtaken because the lockSessionID matches the lock owner.
+ */
+TEST_F(ReplSetDistLockManagerFixture, LockOvertakingWithSessionID) {
+    OID passedLockSessionID("5572007fda9e476582bf3716");
+
+    LocksType currentLockDoc;
+    currentLockDoc.setName("bar");
+    currentLockDoc.setState(LocksType::LOCKED);
+    currentLockDoc.setProcess("otherProcess");
+    currentLockDoc.setLockID(passedLockSessionID);
+    currentLockDoc.setWho("me");
+    currentLockDoc.setWhy("why");
+
+    getMockCatalog()->expectGrabLock(
+        [&passedLockSessionID, &currentLockDoc](
+            StringData, const OID& lockSessionID, StringData, StringData, Date_t, StringData) {
+            ASSERT_EQUALS(passedLockSessionID, lockSessionID);
+        },
+        {ErrorCodes::LockStateChangeFailed, "nMod 0"});
+
+    getMockCatalog()->expectGetLockByName([](StringData name) { ASSERT_EQUALS("bar", name); },
+                                          currentLockDoc);
+
+    LockpingsType pingDoc;
+    pingDoc.setProcess("otherProcess");
+    pingDoc.setPing(Date_t());
+
+    getMockCatalog()->expectGetPing(
+        [](StringData process) { ASSERT_EQUALS("otherProcess", process); }, pingDoc);
+
+    getMockCatalog()->expectGetServerInfo([]() {}, DistLockCatalog::ServerInfo(Date_t(), OID()));
+
+    getMockCatalog()->expectOvertakeLock(
+        [this, &passedLockSessionID, &currentLockDoc](StringData lockID,
+                                                      const OID& lockSessionID,
+                                                      const OID& currentHolderTS,
+                                                      StringData who,
+                                                      StringData processId,
+                                                      Date_t time,
+                                                      StringData why) {
+            ASSERT_EQUALS("bar", lockID);
+            ASSERT_EQUALS(passedLockSessionID, lockSessionID);
+            ASSERT_EQUALS(currentLockDoc.getLockID(), currentHolderTS);
+            ASSERT_EQUALS(getProcessID(), processId);
+            ASSERT_EQUALS("foo", why);
+        },
+        currentLockDoc);
+
+    int unlockCallCount = 0;
+    OID unlockSessionIDPassed;
+
+    {
+        auto lockStatus = getMgr()->lockWithSessionID(
+            txn(), "bar", "foo", passedLockSessionID, Milliseconds(0), Milliseconds(0));
+
+        ASSERT_OK(lockStatus.getStatus());
+
+        getMockCatalog()->expectNoGrabLock();
+        getMockCatalog()->expectUnLock(
+            [&unlockCallCount, &unlockSessionIDPassed](const OID& lockSessionID) {
+                unlockCallCount++;
+                unlockSessionIDPassed = lockSessionID;
+            },
+            Status::OK());
+    }
+
+    ASSERT_EQUALS(1, unlockCallCount);
+    ASSERT_EQUALS(passedLockSessionID, unlockSessionIDPassed);
+}
+
 TEST_F(ReplSetDistLockManagerFixture, CannotOvertakeIfExpirationHasNotElapsed) {
     getMockCatalog()->expectGrabLock(
         [](StringData, const OID&, StringData, StringData, Date_t, StringData) {