SERVER-33485 Active clients should always have a lock state

7 files changed, 48 insertions, 50 deletions

diff --git a/src/mongo/db/catalog/catalog_raii_test.cpp b/src/mongo/db/catalog/catalog_raii_test.cpp
index 3066d2aa1e6..46fc00a4758 100644
--- a/src/mongo/db/catalog/catalog_raii_test.cpp
+++ b/src/mongo/db/catalog/catalog_raii_test.cpp
@@ -51,8 +51,7 @@ public:
     ClientAndCtx makeClientWithLocker(const std::string& clientName) {
         auto client = getGlobalServiceContext()->makeClient(clientName);
         auto opCtx = client->makeOperationContext();
-        opCtx->releaseLockState();
-        opCtx->setLockState(stdx::make_unique<DefaultLockerImpl>());
+        opCtx->swapLockState(stdx::make_unique<DefaultLockerImpl>());
         return std::make_pair(std::move(client), std::move(opCtx));
     }
 
diff --git a/src/mongo/db/concurrency/d_concurrency_test.cpp b/src/mongo/db/concurrency/d_concurrency_test.cpp
index ca6f47643e7..95177c687bc 100644
--- a/src/mongo/db/concurrency/d_concurrency_test.cpp
+++ b/src/mongo/db/concurrency/d_concurrency_test.cpp
@@ -87,7 +87,6 @@ public:
      */
     ServiceContext::UniqueOperationContext makeOpCtx() const {
         auto opCtx = _client->makeOperationContext();
-        opCtx->releaseLockState();
         return opCtx;
     }
 
@@ -105,8 +104,7 @@ public:
             auto client = getGlobalServiceContext()->makeClient(
                 str::stream() << "test client for thread " << i);
             auto opCtx = client->makeOperationContext();
-            opCtx->releaseLockState();
-            opCtx->setLockState(stdx::make_unique<LockerType>());
+            opCtx->swapLockState(stdx::make_unique<LockerType>());
             clients.emplace_back(std::move(client), std::move(opCtx));
         }
         return clients;
@@ -166,13 +164,13 @@ private:
 
 TEST_F(DConcurrencyTestFixture, WriteConflictRetryInstantiatesOK) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     writeConflictRetry(opCtx.get(), "", "", [] {});
 }
 
 TEST_F(DConcurrencyTestFixture, WriteConflictRetryRetriesFunctionOnWriteConflictException) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto&& opDebug = CurOp::get(opCtx.get())->debug();
     ASSERT_EQUALS(0LL, opDebug.writeConflicts);
     ASSERT_EQUALS(100, writeConflictRetry(opCtx.get(), "", "", [&opDebug] {
@@ -186,7 +184,7 @@ TEST_F(DConcurrencyTestFixture, WriteConflictRetryRetriesFunctionOnWriteConflict
 
 TEST_F(DConcurrencyTestFixture, WriteConflictRetryPropagatesNonWriteConflictException) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     ASSERT_THROWS_CODE(writeConflictRetry(opCtx.get(),
                                           "",
                                           "",
@@ -201,7 +199,7 @@ TEST_F(DConcurrencyTestFixture, WriteConflictRetryPropagatesNonWriteConflictExce
 TEST_F(DConcurrencyTestFixture,
        WriteConflictRetryPropagatesWriteConflictExceptionIfAlreadyInAWriteUnitOfWork) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::GlobalWrite globalWrite(opCtx.get());
     WriteUnitOfWork wuow(opCtx.get());
     ASSERT_THROWS(writeConflictRetry(opCtx.get(), "", "", [] { throw WriteConflictException(); }),
@@ -292,21 +290,21 @@ TEST_F(DConcurrencyTestFixture, ResourceMutex) {
 
 TEST_F(DConcurrencyTestFixture, GlobalRead) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::GlobalRead globalRead(opCtx.get());
     ASSERT(opCtx->lockState()->isR());
 }
 
 TEST_F(DConcurrencyTestFixture, GlobalWrite) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::GlobalWrite globalWrite(opCtx.get());
     ASSERT(opCtx->lockState()->isW());
 }
 
 TEST_F(DConcurrencyTestFixture, GlobalWriteAndGlobalRead) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     Lock::GlobalWrite globalWrite(opCtx.get());
@@ -323,7 +321,7 @@ TEST_F(DConcurrencyTestFixture, GlobalWriteAndGlobalRead) {
 TEST_F(DConcurrencyTestFixture,
        GlobalWriteRequiresExplicitDowngradeToIntentWriteModeIfDestroyedWhileHoldingDatabaseLock) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     const ResourceId globalId(RESOURCE_GLOBAL, ResourceId::SINGLETON_GLOBAL);
@@ -371,7 +369,7 @@ TEST_F(DConcurrencyTestFixture,
 TEST_F(DConcurrencyTestFixture,
        GlobalWriteRequiresSupportsDowngradeToIntentWriteModeWhileHoldingDatabaseLock) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     const ResourceId globalId(RESOURCE_GLOBAL, ResourceId::SINGLETON_GLOBAL);
@@ -418,7 +416,7 @@ TEST_F(DConcurrencyTestFixture,
 TEST_F(DConcurrencyTestFixture,
        NestedGlobalWriteSupportsDowngradeToIntentWriteModeWhileHoldingDatabaseLock) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     const ResourceId globalId(RESOURCE_GLOBAL, ResourceId::SINGLETON_GLOBAL);
@@ -608,7 +606,7 @@ TEST_F(DConcurrencyTestFixture, GlobalLockX_TimeoutDueToGlobalLockX) {
 
 TEST_F(DConcurrencyTestFixture, TempReleaseGlobalWrite) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
     Lock::GlobalWrite globalWrite(opCtx.get());
 
@@ -622,7 +620,7 @@ TEST_F(DConcurrencyTestFixture, TempReleaseGlobalWrite) {
 
 TEST_F(DConcurrencyTestFixture, TempReleaseRecursive) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
     Lock::GlobalWrite globalWrite(opCtx.get());
     Lock::DBLock lk(opCtx.get(), "SomeDBName", MODE_X);
@@ -638,7 +636,7 @@ TEST_F(DConcurrencyTestFixture, TempReleaseRecursive) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesS) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbRead(opCtx.get(), "db", MODE_S);
 
     const ResourceId resIdDb(RESOURCE_DATABASE, std::string("db"));
@@ -647,7 +645,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesS) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesX) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbWrite(opCtx.get(), "db", MODE_X);
 
     const ResourceId resIdDb(RESOURCE_DATABASE, std::string("db"));
@@ -656,7 +654,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesX) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesISForAdminIS) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbRead(opCtx.get(), "admin", MODE_IS);
 
     ASSERT(opCtx->lockState()->getLockMode(resourceIdAdminDB) == MODE_IS);
@@ -664,7 +662,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesISForAdminIS) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesSForAdminS) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbRead(opCtx.get(), "admin", MODE_S);
 
     ASSERT(opCtx->lockState()->getLockMode(resourceIdAdminDB) == MODE_S);
@@ -672,7 +670,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesSForAdminS) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesXForAdminIX) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbWrite(opCtx.get(), "admin", MODE_IX);
 
     ASSERT(opCtx->lockState()->getLockMode(resourceIdAdminDB) == MODE_X);
@@ -680,7 +678,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesXForAdminIX) {
 
 TEST_F(DConcurrencyTestFixture, DBLockTakesXForAdminX) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock dbWrite(opCtx.get(), "admin", MODE_X);
 
     ASSERT(opCtx->lockState()->getLockMode(resourceIdAdminDB) == MODE_X);
@@ -688,7 +686,7 @@ TEST_F(DConcurrencyTestFixture, DBLockTakesXForAdminX) {
 
 TEST_F(DConcurrencyTestFixture, MultipleWriteDBLocksOnSameThread) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     Lock::DBLock r1(opCtx.get(), "db1", MODE_X);
     Lock::DBLock r2(opCtx.get(), "db1", MODE_X);
 
@@ -697,7 +695,7 @@ TEST_F(DConcurrencyTestFixture, MultipleWriteDBLocksOnSameThread) {
 
 TEST_F(DConcurrencyTestFixture, MultipleConflictingDBLocksOnSameThread) {
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
     Lock::DBLock r1(opCtx.get(), "db1", MODE_X);
     Lock::DBLock r2(opCtx.get(), "db1", MODE_S);
@@ -710,7 +708,7 @@ TEST_F(DConcurrencyTestFixture, IsDbLockedForSMode) {
     const std::string dbName("db");
 
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
     Lock::DBLock dbLock(opCtx.get(), dbName, MODE_S);
 
@@ -724,7 +722,7 @@ TEST_F(DConcurrencyTestFixture, IsDbLockedForXMode) {
     const std::string dbName("db");
 
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
     Lock::DBLock dbLock(opCtx.get(), dbName, MODE_X);
 
@@ -738,7 +736,7 @@ TEST_F(DConcurrencyTestFixture, IsCollectionLocked_DB_Locked_IS) {
     const std::string ns("db1.coll");
 
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     Lock::DBLock dbLock(opCtx.get(), "db1", MODE_IS);
@@ -769,7 +767,7 @@ TEST_F(DConcurrencyTestFixture, IsCollectionLocked_DB_Locked_IX) {
     const std::string ns("db1.coll");
 
     auto opCtx = makeOpCtx();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     auto lockState = opCtx->lockState();
 
     Lock::DBLock dbLock(opCtx.get(), "db1", MODE_IX);
diff --git a/src/mongo/db/op_observer_impl_test.cpp b/src/mongo/db/op_observer_impl_test.cpp
index 72917228aaf..37449d5f177 100644
--- a/src/mongo/db/op_observer_impl_test.cpp
+++ b/src/mongo/db/op_observer_impl_test.cpp
@@ -388,8 +388,7 @@ TEST_F(OpObserverTest, MultipleAboutToDeleteAndOnDelete) {
 DEATH_TEST_F(OpObserverTest, AboutToDeleteMustPreceedOnDelete, "invariant") {
     OpObserverImpl opObserver;
     auto opCtx = cc().makeOperationContext();
-    opCtx->releaseLockState();
-    opCtx->setLockState(stdx::make_unique<LockerNoop>());
+    opCtx->swapLockState(stdx::make_unique<LockerNoop>());
     NamespaceString nss = {"test", "coll"};
     opObserver.onDelete(opCtx.get(), nss, {}, {}, false, {});
 }
@@ -397,8 +396,7 @@ DEATH_TEST_F(OpObserverTest, AboutToDeleteMustPreceedOnDelete, "invariant") {
 DEATH_TEST_F(OpObserverTest, EachOnDeleteRequiresAboutToDelete, "invariant") {
     OpObserverImpl opObserver;
     auto opCtx = cc().makeOperationContext();
-    opCtx->releaseLockState();
-    opCtx->setLockState(stdx::make_unique<LockerNoop>());
+    opCtx->swapLockState(stdx::make_unique<LockerNoop>());
     NamespaceString nss = {"test", "coll"};
     opObserver.aboutToDelete(opCtx.get(), nss, {});
     opObserver.onDelete(opCtx.get(), nss, {}, {}, false, {});
diff --git a/src/mongo/db/operation_context.cpp b/src/mongo/db/operation_context.cpp
index a647d0fcbd0..e503c2e071d 100644
--- a/src/mongo/db/operation_context.cpp
+++ b/src/mongo/db/operation_context.cpp
@@ -382,17 +382,19 @@ OperationContext::RecoveryUnitState OperationContext::setRecoveryUnit(RecoveryUn
     return oldState;
 }
 
-std::unique_ptr<Locker> OperationContext::releaseLockState() {
-    dassert(_locker);
-    return std::move(_locker);
-}
-
 void OperationContext::setLockState(std::unique_ptr<Locker> locker) {
-    dassert(!_locker);
-    dassert(locker);
+    invariant(!_locker);
+    invariant(locker);
     _locker = std::move(locker);
 }
 
+std::unique_ptr<Locker> OperationContext::swapLockState(std::unique_ptr<Locker> locker) {
+    invariant(_locker);
+    invariant(locker);
+    _locker.swap(locker);
+    return locker;
+}
+
 Date_t OperationContext::getExpirationDateForWaitForValue(Milliseconds waitFor) {
     return getServiceContext()->getPreciseClockSource()->now() + waitFor;
 }
diff --git a/src/mongo/db/operation_context.h b/src/mongo/db/operation_context.h
index 77029acd1c7..cb39e4cbec6 100644
--- a/src/mongo/db/operation_context.h
+++ b/src/mongo/db/operation_context.h
@@ -131,10 +131,9 @@ public:
     void setLockState(std::unique_ptr<Locker> locker);
 
     /**
-     * Releases the locker to the caller. Call during OperationContext cleanup or initialization,
-     * only.
+     * Swaps the locker, releasing the old locker to the caller.
      */
-    std::unique_ptr<Locker> releaseLockState();
+    std::unique_ptr<Locker> swapLockState(std::unique_ptr<Locker> locker);
 
     /**
      * Raises a AssertionException if this operation is in a killed state.
diff --git a/src/mongo/db/repl/service_context_repl_mock.cpp b/src/mongo/db/repl/service_context_repl_mock.cpp
index 31dcf83c3ba..c8d87716e8d 100644
--- a/src/mongo/db/repl/service_context_repl_mock.cpp
+++ b/src/mongo/db/repl/service_context_repl_mock.cpp
@@ -39,8 +39,7 @@ namespace repl {
 
 std::unique_ptr<OperationContext> ServiceContextReplMock::_newOpCtx(Client* client, unsigned opId) {
     auto opCtx = stdx::make_unique<OperationContextNoop>(client, opId);
-    opCtx->releaseLockState();
-    opCtx->setLockState(stdx::make_unique<MMAPV1LockerImpl>());
+    opCtx->swapLockState(stdx::make_unique<MMAPV1LockerImpl>());
     return {std::move(opCtx)};
 }
 
diff --git a/src/mongo/db/session.cpp b/src/mongo/db/session.cpp
index f1773c3a682..1310b261932 100644
--- a/src/mongo/db/session.cpp
+++ b/src/mongo/db/session.cpp
@@ -475,11 +475,9 @@ void Session::stashTransactionResources(OperationContext* opCtx) {
     opCtx->getWriteUnitOfWork()->release();
     opCtx->setWriteUnitOfWork(nullptr);
 
-    _stashedLocker = opCtx->releaseLockState();
+    _stashedLocker = opCtx->swapLockState(stdx::make_unique<DefaultLockerImpl>());
     _stashedLocker->releaseTicket();
     _stashedRecoveryUnit.reset(opCtx->releaseRecoveryUnit());
-
-    opCtx->setLockState(stdx::make_unique<DefaultLockerImpl>());
     opCtx->setRecoveryUnit(opCtx->getServiceContext()->getGlobalStorageEngine()->newRecoveryUnit(),
                            OperationContext::kNotInUnitOfWork);
 }
@@ -502,9 +500,14 @@ void Session::unstashTransactionResources(OperationContext* opCtx) {
 
     if (_stashedLocker) {
         invariant(_stashedRecoveryUnit);
-        opCtx->releaseLockState();
         _stashedLocker->reacquireTicket();
-        opCtx->setLockState(std::move(_stashedLocker));
+
+        // We intentionally do not capture the return value of swapLockState(), which is just an
+        // empty locker. At the end of the operation, if the transaction is not complete, we will
+        // stash the operation context's locker and replace it with a new empty locker.
+        invariant(opCtx->lockState()->getClientState() == Locker::ClientState::kInactive);
+        opCtx->swapLockState(std::move(_stashedLocker));
+
         opCtx->setRecoveryUnit(_stashedRecoveryUnit.release(),
                                OperationContext::RecoveryUnitState::kNotInUnitOfWork);
         opCtx->setWriteUnitOfWork(WriteUnitOfWork::createForSnapshotResume(opCtx));