SERVER-25202 Expose lock/unlock control on DistLockManager

This change exposes the lock/unlock methods as public on the
DistLockManager. This enables the scenario of being able to pass locks
between threads, which is necessary for the parallel balancer.

Also makes the sharding balancer acquire the balancer distributed lock
permanently.

8 files changed, 74 insertions, 112 deletions

diff --git a/src/mongo/s/balancer/balancer.cpp b/src/mongo/s/balancer/balancer.cpp
index a5268a81337..598a810d463 100644
--- a/src/mongo/s/balancer/balancer.cpp
+++ b/src/mongo/s/balancer/balancer.cpp
@@ -308,46 +308,42 @@ void Balancer::report(OperationContext* txn, BSONObjBuilder* builder) {
 void Balancer::_mainThread() {
     Client::initThread("Balancer");
 
-    // TODO (SERVER-24754): Balancer thread should only keep the operation context alive while it is
-    // doing balancing
-    const auto txn = cc().makeOperationContext();
-
     log() << "CSRS balancer is starting";
 
     const Seconds kInitBackoffInterval(60);
 
-    // The balancer thread is holding the balancer during its entire lifetime
-    boost::optional<DistLockManager::ScopedDistLock> scopedBalancerLock;
-
-    // Take the balancer distributed lock
-    while (!_stopRequested() && !scopedBalancerLock) {
+    // Take the balancer distributed lock and hold it permanently
+    while (!_stopRequested()) {
+        auto txn = cc().makeOperationContext();
         auto shardingContext = Grid::get(txn.get());
-        auto scopedDistLock = shardingContext->catalogClient(txn.get())->distLock(
-            txn.get(), "balancer", "CSRS Balancer");
-        if (!scopedDistLock.isOK()) {
-            warning() << "Balancer distributed lock could not be acquired and will be retried in "
-                         "one minute"
-                      << causedBy(scopedDistLock.getStatus());
-            _sleepFor(txn.get(), kInitBackoffInterval);
-            continue;
+
+        auto distLockHandleStatus =
+            shardingContext->catalogClient(txn.get())->getDistLockManager()->lockWithSessionID(
+                txn.get(), "balancer", "CSRS Balancer", OID::gen());
+        if (distLockHandleStatus.isOK()) {
+            break;
         }
 
-        // Initialization and distributed lock acquisition succeeded
-        scopedBalancerLock = std::move(scopedDistLock.getValue());
+        warning() << "Balancer distributed lock could not be acquired and will be retried in "
+                  << durationCount<Seconds>(kInitBackoffInterval) << " seconds"
+                  << causedBy(distLockHandleStatus.getStatus());
+
+        _sleepFor(txn.get(), kInitBackoffInterval);
     }
 
     log() << "CSRS balancer thread is now running";
 
     // Main balancer loop
     while (!_stopRequested()) {
+        auto txn = cc().makeOperationContext();
         auto shardingContext = Grid::get(txn.get());
         auto balancerConfig = shardingContext->getBalancerConfiguration();
 
         BalanceRoundDetails roundDetails;
 
-        try {
-            _beginRound(txn.get());
+        _beginRound(txn.get());
 
+        try {
             shardingContext->shardRegistry()->reload(txn.get());
 
             uassert(13258, "oids broken after resetting!", _checkOIDs(txn.get()));
diff --git a/src/mongo/s/catalog/dist_lock_manager.cpp b/src/mongo/s/catalog/dist_lock_manager.cpp
index 732e2fc664a..7b5911e78f6 100644
--- a/src/mongo/s/catalog/dist_lock_manager.cpp
+++ b/src/mongo/s/catalog/dist_lock_manager.cpp
@@ -71,6 +71,20 @@ DistLockManager::ScopedDistLock& DistLockManager::ScopedDistLock::operator=(
     return *this;
 }
 
+StatusWith<DistLockManager::ScopedDistLock> DistLockManager::lock(OperationContext* txn,
+                                                                  StringData name,
+                                                                  StringData whyMessage,
+                                                                  Milliseconds waitFor,
+                                                                  Milliseconds lockTryInterval) {
+    auto distLockHandleStatus =
+        lockWithSessionID(txn, name, whyMessage, OID::gen(), waitFor, lockTryInterval);
+    if (!distLockHandleStatus.isOK()) {
+        return distLockHandleStatus.getStatus();
+    }
+
+    return DistLockManager::ScopedDistLock(txn, std::move(distLockHandleStatus.getValue()), this);
+}
+
 Status DistLockManager::ScopedDistLock::checkStatus() {
     if (!_lockManager) {
         return Status(ErrorCodes::IllegalOperation, "no lock manager, lock was not acquired");
diff --git a/src/mongo/s/catalog/dist_lock_manager.h b/src/mongo/s/catalog/dist_lock_manager.h
index 9622fae428e..7b957a0221a 100644
--- a/src/mongo/s/catalog/dist_lock_manager.h
+++ b/src/mongo/s/catalog/dist_lock_manager.h
@@ -131,12 +131,11 @@ public:
      * Returns ErrorCodes::DistributedClockSkewed when a clock skew is detected.
      * Returns ErrorCodes::LockBusy if the lock is being held.
      */
-    virtual StatusWith<ScopedDistLock> lock(
-        OperationContext* txn,
-        StringData name,
-        StringData whyMessage,
-        Milliseconds waitFor = kDefaultLockTimeout,
-        Milliseconds lockTryInterval = kDefaultLockRetryInterval) = 0;
+    StatusWith<ScopedDistLock> lock(OperationContext* txn,
+                                    StringData name,
+                                    StringData whyMessage,
+                                    Milliseconds waitFor = kDefaultLockTimeout,
+                                    Milliseconds lockTryInterval = kDefaultLockRetryInterval);
 
     /**
      * Same behavior as lock(...) above, except takes a specific lock session ID "lockSessionID"
@@ -146,7 +145,7 @@ public:
      * immediately reacquired if "lockSessionID" matches that of the lock, rather than waiting for
      * the inactive lock to expire.
      */
-    virtual StatusWith<ScopedDistLock> lockWithSessionID(
+    virtual StatusWith<DistLockHandle> lockWithSessionID(
         OperationContext* txn,
         StringData name,
         StringData whyMessage,
@@ -155,20 +154,21 @@ public:
         Milliseconds lockTryInterval = kDefaultLockRetryInterval) = 0;
 
     /**
-     * Makes a best-effort attempt to unlock all locks owned by the given processID.
-     */
-    virtual void unlockAll(OperationContext* txn, const std::string& processID) = 0;
-
-protected:
-    /**
      * Unlocks the given lockHandle. Will attempt to retry again later if the config
      * server is not reachable.
      */
     virtual void unlock(OperationContext* txn, const DistLockHandle& lockHandle) = 0;
 
     /**
+     * Makes a best-effort attempt to unlock all locks owned by the given processID.
+     */
+    virtual void unlockAll(OperationContext* txn, const std::string& processID) = 0;
+
+protected:
+    /**
      * Checks if the lockHandle still exists in the config server.
      */
     virtual Status checkStatus(OperationContext* txn, const DistLockHandle& lockHandle) = 0;
 };
-}
+
+}  // namespace mongo
diff --git a/src/mongo/s/catalog/dist_lock_manager_mock.cpp b/src/mongo/s/catalog/dist_lock_manager_mock.cpp
index 9325caf29d7..7346ae5962a 100644
--- a/src/mongo/s/catalog/dist_lock_manager_mock.cpp
+++ b/src/mongo/s/catalog/dist_lock_manager_mock.cpp
@@ -72,23 +72,12 @@ std::string DistLockManagerMock::getProcessID() {
     return "Mock dist lock manager process id";
 }
 
-StatusWith<DistLockManager::ScopedDistLock> DistLockManagerMock::lock(
-    OperationContext* txn,
-    StringData name,
-    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) {
+StatusWith<DistLockHandle> DistLockManagerMock::lockWithSessionID(OperationContext* txn,
+                                                                  StringData name,
+                                                                  StringData whyMessage,
+                                                                  const OID lockSessionID,
+                                                                  Milliseconds waitFor,
+                                                                  Milliseconds lockTryInterval) {
     _lockChecker(name, whyMessage, waitFor, lockTryInterval);
     _lockChecker = NoLockFuncSet;
 
@@ -108,7 +97,7 @@ StatusWith<DistLockManager::ScopedDistLock> DistLockManagerMock::lockWithSession
     info.lockID = lockSessionID;
     _locks.push_back(info);
 
-    return DistLockManager::ScopedDistLock(nullptr, info.lockID, this);
+    return info.lockID;
 }
 
 void DistLockManagerMock::unlockAll(OperationContext* txn, const std::string& processID) {
diff --git a/src/mongo/s/catalog/dist_lock_manager_mock.h b/src/mongo/s/catalog/dist_lock_manager_mock.h
index 5d32acda6b3..57fa13cca3a 100644
--- a/src/mongo/s/catalog/dist_lock_manager_mock.h
+++ b/src/mongo/s/catalog/dist_lock_manager_mock.h
@@ -47,19 +47,12 @@ public:
 
     virtual std::string getProcessID() override;
 
-    virtual StatusWith<DistLockManager::ScopedDistLock> lock(OperationContext* txn,
-                                                             StringData name,
-                                                             StringData whyMessage,
-                                                             Milliseconds waitFor,
-                                                             Milliseconds lockTryInterval) override;
-
-    virtual StatusWith<DistLockManager::ScopedDistLock> lockWithSessionID(
-        OperationContext* txn,
-        StringData name,
-        StringData whyMessage,
-        const OID lockSessionID,
-        Milliseconds waitFor,
-        Milliseconds lockTryInterval) override;
+    virtual StatusWith<DistLockHandle> lockWithSessionID(OperationContext* txn,
+                                                         StringData name,
+                                                         StringData whyMessage,
+                                                         const OID lockSessionID,
+                                                         Milliseconds waitFor,
+                                                         Milliseconds lockTryInterval) override;
 
     virtual void unlockAll(OperationContext* txn, const std::string& processID) override;
 
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 18809505027..fef9d2d65d3 100644
--- a/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp
+++ b/src/mongo/s/catalog/replset/replset_dist_lock_manager.cpp
@@ -261,22 +261,12 @@ StatusWith<bool> ReplSetDistLockManager::isLockExpired(OperationContext* txn,
     return false;
 }
 
-StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lock(
-    OperationContext* txn,
-    StringData name,
-    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) {
+StatusWith<DistLockHandle> ReplSetDistLockManager::lockWithSessionID(OperationContext* txn,
+                                                                     StringData name,
+                                                                     StringData whyMessage,
+                                                                     const OID lockSessionID,
+                                                                     Milliseconds waitFor,
+                                                                     Milliseconds lockTryInterval) {
     Timer timer(_serviceContext->getTickSource());
     Timer msgTimer(_serviceContext->getTickSource());
 
@@ -286,6 +276,7 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lockWithSess
     int networkErrorRetries = 0;
 
     auto configShard = Grid::get(txn)->shardRegistry()->getConfigShard();
+
     // Distributed lock acquisition works by tring to update the state of the lock to 'taken'. If
     // the lock is currently taken, we will back off and try the acquisition again, repeating this
     // until the lockTryInterval has been reached. If a network error occurs at each lock
@@ -321,7 +312,7 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lockWithSess
             // the lock document.
             log() << "distributed lock '" << name << "' acquired for '"
                   << redact(whyMessage.toString()) << "', ts : " << lockSessionID;
-            return ScopedDistLock(txn, lockSessionID, this);
+            return lockSessionID;
         }
 
         // If a network error occurred, unlock the lock synchronously and try again
@@ -390,7 +381,7 @@ StatusWith<DistLockManager::ScopedDistLock> ReplSetDistLockManager::lockWithSess
 
                     LOG(0) << "lock '" << name << "' successfully forced";
                     LOG(0) << "distributed lock '" << name << "' acquired, ts : " << lockSessionID;
-                    return ScopedDistLock(txn, lockSessionID, this);
+                    return lockSessionID;
                 }
 
                 if (overtakeStatus != ErrorCodes::LockStateChangeFailed) {
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 c29482fbe54..c4f020944b1 100644
--- a/src/mongo/s/catalog/replset/replset_dist_lock_manager.h
+++ b/src/mongo/s/catalog/replset/replset_dist_lock_manager.h
@@ -68,24 +68,18 @@ public:
 
     virtual std::string getProcessID() override;
 
-    virtual StatusWith<ScopedDistLock> lock(OperationContext* txn,
-                                            StringData name,
-                                            StringData whyMessage,
-                                            Milliseconds waitFor,
-                                            Milliseconds lockTryInterval) override;
-
-    virtual StatusWith<ScopedDistLock> lockWithSessionID(OperationContext* txn,
+    virtual StatusWith<DistLockHandle> lockWithSessionID(OperationContext* txn,
                                                          StringData name,
                                                          StringData whyMessage,
                                                          const OID lockSessionID,
                                                          Milliseconds waitFor,
                                                          Milliseconds lockTryInterval) override;
 
+    virtual void unlock(OperationContext* txn, const DistLockHandle& lockSessionID) override;
+
     virtual void unlockAll(OperationContext* txn, const std::string& processID) override;
 
 protected:
-    virtual void unlock(OperationContext* txn, const DistLockHandle& lockSessionID) override;
-
     virtual Status checkStatus(OperationContext* txn, const DistLockHandle& lockSessionID) override;
 
 private:
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 c353ee05e82..537bfd124b7 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
@@ -1186,26 +1186,11 @@ TEST_F(ReplSetDistLockManagerFixture, LockOvertakingWithSessionID) {
         },
         currentLockDoc);
 
-    int unlockCallCount = 0;
-    OID unlockSessionIDPassed;
+    auto distLockHandleStatus = getMgr()->lockWithSessionID(
+        txn(), "bar", "foo", passedLockSessionID, Milliseconds(0), Milliseconds(0));
+    ASSERT_OK(distLockHandleStatus.getStatus());
 
-    {
-        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);
+    getMockCatalog()->expectNoGrabLock();
 }
 
 TEST_F(ReplSetDistLockManagerFixture, CannotOvertakeIfExpirationHasNotElapsed) {