1 files changed, 26 insertions, 1 deletions
diff --git a/src/mongo/db/s/database_sharding_state.cpp b/src/mongo/db/s/database_sharding_state.cpp
index a4c0dff2667..aed788ee1fb 100644
--- a/src/mongo/db/s/database_sharding_state.cpp
+++ b/src/mongo/db/s/database_sharding_state.cpp
@@ -121,10 +121,12 @@ std::shared_ptr<DatabaseShardingState> DatabaseShardingState::getSharedForLockFr
 }
 
 void DatabaseShardingState::enterCriticalSectionCatchUpPhase(OperationContext* opCtx,
-                                                             DSSLock&,
+                                                             DSSLock& dssLock,
                                                              const BSONObj& reason) {
     invariant(opCtx->lockState()->isDbLockedForMode(_dbName, MODE_X));
     _critSec.enterCriticalSectionCatchUpPhase(reason);
+
+    cancelDbMetadataRefresh(dssLock);
 }
 
 void DatabaseShardingState::enterCriticalSectionCommitPhase(OperationContext* opCtx,
@@ -220,4 +222,27 @@ void DatabaseShardingState::clearMovePrimarySourceManager(OperationContext* opCt
     _sourceMgr = nullptr;
 }
 
+void DatabaseShardingState::setDbMetadataRefreshFuture(SharedSemiFuture<void> future,
+                                                       CancellationSource cancellationSource,
+                                                       const DSSLock&) {
+    invariant(!_dbMetadataRefresh);
+    _dbMetadataRefresh.emplace(std::move(future), std::move(cancellationSource));
+}
+
+boost::optional<SharedSemiFuture<void>> DatabaseShardingState::getDbMetadataRefreshFuture(
+    const DSSLock&) const {
+    return _dbMetadataRefresh ? boost::optional<SharedSemiFuture<void>>(_dbMetadataRefresh->future)
+                              : boost::none;
+}
+
+void DatabaseShardingState::resetDbMetadataRefreshFuture(const DSSLock&) {
+    _dbMetadataRefresh = boost::none;
+}
+
+void DatabaseShardingState::cancelDbMetadataRefresh(const DSSLock&) {
+    if (_dbMetadataRefresh) {
+        _dbMetadataRefresh->cancellationSource.cancel();
+    }
+}
+
 }  // namespace mongo