SERVER-33081 Reset KeysCollectionManager during rollback properly

1 files changed, 19 insertions, 8 deletions

diff --git a/src/mongo/db/keys_collection_manager.cpp b/src/mongo/db/keys_collection_manager.cpp
index 889b822477c..c3c66b2b2c0 100644
--- a/src/mongo/db/keys_collection_manager.cpp
+++ b/src/mongo/db/keys_collection_manager.cpp
@@ -172,6 +172,9 @@ void KeysCollectionManager::enableKeyGenerator(OperationContext* opCtx, bool doE
                 return StatusWith<KeysCollectionDocument>(keyGenerationStatus);
             }
 
+            // Set a small deadline so the refresh will not wait forever to satisfy readConcern.
+            opCtx->setDeadlineAfterNowBy(Milliseconds(500));
+
             // An error encountered by the keyGenerator should not prevent refreshing the cache
             auto cacheRefreshStatus = _keysCache.refresh(opCtx);
 
@@ -182,8 +185,11 @@ void KeysCollectionManager::enableKeyGenerator(OperationContext* opCtx, bool doE
             return cacheRefreshStatus;
         });
     } else {
-        _refresher.switchFunc(
-            opCtx, [this](OperationContext* opCtx) { return _keysCache.refresh(opCtx); });
+        _refresher.switchFunc(opCtx, [this](OperationContext* opCtx) {
+            // Set a small deadline so the refresh will not wait forever to satisfy readConcern.
+            opCtx->setDeadlineAfterNowBy(Milliseconds(500));
+            return _keysCache.refresh(opCtx);
+        });
     }
 }
 
@@ -195,7 +201,7 @@ void KeysCollectionManager::PeriodicRunner::refreshNow(OperationContext* opCtx)
     auto refreshRequest = [this]() {
         stdx::lock_guard<stdx::mutex> lk(_mutex);
 
-        if (_inShutdown) {
+        if (_isStopping) {
             uasserted(ErrorCodes::ShutdownInProgress,
                       "aborting keys cache refresh because node is shutting down");
         }
@@ -229,7 +235,7 @@ void KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh(ServiceContext* s
         {
             stdx::lock_guard<stdx::mutex> lock(_mutex);
 
-            if (_inShutdown) {
+            if (_isStopping) {
                 break;
             }
 
@@ -242,7 +248,6 @@ void KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh(ServiceContext* s
 
         {
             auto opCtx = cc().makeOperationContext();
-
             auto latestKeyStatusWith = (*doRefresh)(opCtx.get());
             if (latestKeyStatusWith.getStatus().isOK()) {
                 errorCount = 0;
@@ -285,7 +290,7 @@ void KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh(ServiceContext* s
             _refreshRequest.reset();
         }
 
-        if (_inShutdown) {
+        if (_isStopping) {
             break;
         }
 
@@ -324,7 +329,7 @@ void KeysCollectionManager::PeriodicRunner::start(ServiceContext* service,
                                                   Milliseconds refreshInterval) {
     stdx::lock_guard<stdx::mutex> lock(_mutex);
     invariant(!_backgroundThread.joinable());
-    invariant(!_inShutdown);
+    invariant(!_isStopping);
 
     _backgroundThread = stdx::thread([this, service, threadName, refreshInterval] {
         _doPeriodicRefresh(service, threadName, refreshInterval);
@@ -338,11 +343,17 @@ void KeysCollectionManager::PeriodicRunner::stop() {
             return;
         }
 
-        _inShutdown = true;
+        _isStopping = true;
         _refreshNeededCV.notify_all();
     }
 
     _backgroundThread.join();
+
+    {
+        stdx::lock_guard<stdx::mutex> lock(_mutex);
+        invariant(!_backgroundThread.joinable());
+        _isStopping = false;
+    }
 }
 
 bool KeysCollectionManager::PeriodicRunner::hasSeenKeys() {