1 files changed, 28 insertions, 10 deletions

diff --git a/src/mongo/db/s/database_sharding_state.h b/src/mongo/db/s/database_sharding_state.h
index 3033fb27ea4..f00157183f0 100644
--- a/src/mongo/db/s/database_sharding_state.h
+++ b/src/mongo/db/s/database_sharding_state.h
@@ -33,6 +33,7 @@
 #include "mongo/base/disallow_copying.h"
 #include "mongo/db/catalog/database.h"
 #include "mongo/db/s/sharding_migration_critical_section.h"
+#include "mongo/db/s/sharding_state_lock.h"
 #include "mongo/s/database_version_gen.h"
 
 namespace mongo {
@@ -47,6 +48,12 @@ class DatabaseShardingState {
     MONGO_DISALLOW_COPYING(DatabaseShardingState);
 
 public:
+    /**
+     * A ShardingStateLock is used on DatabaseShardingState operations in order to ensure
+     * synchronization across operations.
+     */
+    using DSSLock = ShardingStateLock<DatabaseShardingState>;
+
     static const Database::Decoration<DatabaseShardingState> get;
 
     DatabaseShardingState();
@@ -56,12 +63,13 @@ public:
      * Methods to control the databases's critical section. Must be called with the database X lock
      * held.
      */
-    void enterCriticalSectionCatchUpPhase(OperationContext* opCtx);
-    void enterCriticalSectionCommitPhase(OperationContext* opCtx);
+    void enterCriticalSectionCatchUpPhase(OperationContext* opCtx, DSSLock&);
+    void enterCriticalSectionCommitPhase(OperationContext* opCtx, DSSLock&);
     void exitCriticalSection(OperationContext* opCtx,
-                             boost::optional<DatabaseVersion> newDbVersion);
+                             boost::optional<DatabaseVersion> newDbVersion,
+                             DSSLock&);
 
-    auto getCriticalSectionSignal(ShardingMigrationCriticalSection::Operation op) const {
+    auto getCriticalSectionSignal(ShardingMigrationCriticalSection::Operation op, DSSLock&) const {
         return _critSec.getSignal(op);
     }
 
@@ -70,42 +78,52 @@ public:
      *
      * Invariants that the caller holds the DBLock in X or IS.
      */
-    boost::optional<DatabaseVersion> getDbVersion(OperationContext* opCtx) const;
+    boost::optional<DatabaseVersion> getDbVersion(OperationContext* opCtx, DSSLock&) const;
 
     /**
      * Sets this shard server's cached dbVersion to newVersion.
      *
      * Invariants that the caller holds the DBLock in X mode.
      */
-    void setDbVersion(OperationContext* opCtx, boost::optional<DatabaseVersion> newVersion);
+    void setDbVersion(OperationContext* opCtx,
+                      boost::optional<DatabaseVersion> newVersion,
+                      DSSLock&);
 
     /**
      * If _critSecSignal is non-null, always throws StaleDbVersion.
      * Otherwise, if there is a client dbVersion on the OperationContext, compares it with this
      * shard server's cached dbVersion and throws StaleDbVersion if they do not match.
      */
-    void checkDbVersion(OperationContext* opCtx) const;
+    void checkDbVersion(OperationContext* opCtx, DSSLock&) const;
 
     /**
      * Returns the active movePrimary source manager, if one is available.
      */
-    MovePrimarySourceManager* getMovePrimarySourceManager();
+    MovePrimarySourceManager* getMovePrimarySourceManager(DSSLock&);
 
     /**
      * Attaches a movePrimary source manager to this database's sharding state. Must be called with
      * the database lock in X mode. May not be called if there is a movePrimary source manager
      * already installed. Must be followed by a call to clearMovePrimarySourceManager.
      */
-    void setMovePrimarySourceManager(OperationContext* opCtx, MovePrimarySourceManager* sourceMgr);
+    void setMovePrimarySourceManager(OperationContext* opCtx,
+                                     MovePrimarySourceManager* sourceMgr,
+                                     DSSLock&);
 
     /**
      * Removes a movePrimary source manager from this database's sharding state. Must be called with
      * with the database lock in X mode. May not be called if there isn't a movePrimary source
      * manager installed already through a previous call to setMovePrimarySourceManager.
      */
-    void clearMovePrimarySourceManager(OperationContext* opCtx);
+    void clearMovePrimarySourceManager(OperationContext* opCtx, DSSLock&);
 
 private:
+    friend DSSLock;
+
+    // Object-wide ResourceMutex to protect changes to the DatabaseShardingState or objects held
+    // within.
+    Lock::ResourceMutex _stateChangeMutex{"DatabaseShardingState"};
+
     // Modifying the state below requires holding the DBLock in X mode; holding the DBLock in any
     // mode is acceptable for reading it. (Note: accessing this class at all requires holding the
     // DBLock in some mode, since it requires having a pointer to the Database).