SERVER-57494 Clear the locker before re-attempting a command using the same opCtx as we used to refresh sharding state

(cherry picked from commit bffbd983cdb1db991c06d46196f9c008d3980411)

4 files changed, 35 insertions, 2 deletions

diff --git a/src/mongo/db/service_entry_point_common.cpp b/src/mongo/db/service_entry_point_common.cpp
index 6f2b84eb9aa..d1626fb0c57 100644
--- a/src/mongo/db/service_entry_point_common.cpp
+++ b/src/mongo/db/service_entry_point_common.cpp
@@ -692,6 +692,20 @@ private:
         }
     }
 
+    // Do any initialization of the lock state required for a transaction.
+    void _setLockStateForTransaction(OperationContext* opCtx) {
+        opCtx->lockState()->setSharedLocksShouldTwoPhaseLock(true);
+        opCtx->lockState()->setShouldConflictWithSecondaryBatchApplication(false);
+    }
+
+    // Clear any lock state which may have changed after the locker update.
+    void _resetLockerStateAfterShardingUpdate(OperationContext* opCtx) {
+        dassert(!opCtx->isContinuingMultiDocumentTransaction());
+        _execContext->behaviors->resetLockerState(opCtx);
+        if (opCtx->isStartingMultiDocumentTransaction())
+            _setLockStateForTransaction(opCtx);
+    }
+
     // Any logic, such as authorization and auditing, that must precede execution of the command.
     void _initiateCommand();
 
@@ -1581,8 +1595,7 @@ void ExecCommandDatabase::_initiateCommand() {
     }
 
     if (startTransaction) {
-        opCtx->lockState()->setSharedLocksShouldTwoPhaseLock(true);
-        opCtx->lockState()->setShouldConflictWithSecondaryBatchApplication(false);
+        _setLockStateForTransaction(opCtx);
     }
 
     // Remember whether or not this operation is starting a transaction, in case something later in
@@ -1673,6 +1686,7 @@ Future<void> ExecCommandDatabase::_commandExec() {
                 if (refreshed) {
                     _refreshedDatabase = true;
                     if (!opCtx->isContinuingMultiDocumentTransaction()) {
+                        _resetLockerStateAfterShardingUpdate(opCtx);
                         return _commandExec();
                     }
                 }
@@ -1693,6 +1707,7 @@ Future<void> ExecCommandDatabase::_commandExec() {
                     if (refreshed) {
                         _refreshedCollection = true;
                         if (!opCtx->isContinuingMultiDocumentTransaction()) {
+                            _resetLockerStateAfterShardingUpdate(opCtx);
                             return _commandExec();
                         }
                     }
@@ -1719,6 +1734,7 @@ Future<void> ExecCommandDatabase::_commandExec() {
                 if (refreshed) {
                     _refreshedCatalogCache = true;
                     if (!opCtx->isContinuingMultiDocumentTransaction()) {
+                        _resetLockerStateAfterShardingUpdate(opCtx);
                         return _commandExec();
                     }
                 }
diff --git a/src/mongo/db/service_entry_point_common.h b/src/mongo/db/service_entry_point_common.h
index 9500752fd4a..53de36e4829 100644
--- a/src/mongo/db/service_entry_point_common.h
+++ b/src/mongo/db/service_entry_point_common.h
@@ -94,6 +94,8 @@ struct ServiceEntryPointCommon {
             OperationContext* opCtx, const ShardCannotRefreshDueToLocksHeldInfo& refreshInfo) const
             noexcept = 0;
 
+        virtual void resetLockerState(OperationContext* opCtx) const noexcept = 0;
+
         virtual void advanceConfigOpTimeFromRequestMetadata(OperationContext* opCtx) const = 0;
 
         MONGO_WARN_UNUSED_RESULT_FUNCTION virtual std::unique_ptr<PolymorphicScoped>
diff --git a/src/mongo/db/service_entry_point_mongod.cpp b/src/mongo/db/service_entry_point_mongod.cpp
index f6b22b26899..b2b7fa48944 100644
--- a/src/mongo/db/service_entry_point_mongod.cpp
+++ b/src/mongo/db/service_entry_point_mongod.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/service_entry_point_mongod.h"
 
 #include "mongo/db/commands/fsync_locked.h"
+#include "mongo/db/concurrency/lock_state.h"
 #include "mongo/db/curop.h"
 #include "mongo/db/read_concern.h"
 #include "mongo/db/repl/repl_client_info.h"
@@ -253,6 +254,18 @@ public:
             .isOK();
     }
 
+    // The refreshDatabase, refreshCollection, and refreshCatalogCache methods may have modified the
+    // locker state, in particular the flags which say if the operation took a write lock or shared
+    // lock.  This will cause mongod to perhaps erroneously check for write concern when no writes
+    // were done, or unnecessarily kill a read operation.  If we re-use the opCtx to retry command
+    // execution, we must reset the locker state.
+    void resetLockerState(OperationContext* opCtx) const noexcept override {
+        // It is necessary to lock the client to change the Locker on the OperationContext.
+        stdx::lock_guard<Client> lk(*opCtx->getClient());
+        invariant(!opCtx->lockState()->isLocked());
+        opCtx->swapLockState(std::make_unique<LockerImpl>(), lk);
+    }
+
     void advanceConfigOpTimeFromRequestMetadata(OperationContext* opCtx) const override {
         // Handle config optime information that may have been sent along with the command.
         rpc::advanceConfigOpTimeFromRequestMetadata(opCtx);
diff --git a/src/mongo/embedded/service_entry_point_embedded.cpp b/src/mongo/embedded/service_entry_point_embedded.cpp
index 428d15e2d8a..56626f0c87b 100644
--- a/src/mongo/embedded/service_entry_point_embedded.cpp
+++ b/src/mongo/embedded/service_entry_point_embedded.cpp
@@ -137,6 +137,8 @@ public:
         return false;
     }
 
+    void resetLockerState(OperationContext* opCtx) const noexcept override {}
+
     void advanceConfigOpTimeFromRequestMetadata(OperationContext* opCtx) const override {}
 
     std::unique_ptr<PolymorphicScoped> scopedOperationCompletionShardingActions(