SERVER-33081 Reset KeysCollectionManager during rollback properly

8 files changed, 48 insertions, 5 deletions

diff --git a/src/mongo/db/keys_collection_cache.cpp b/src/mongo/db/keys_collection_cache.cpp
index 0491d69397f..4dd2c28b31d 100644
--- a/src/mongo/db/keys_collection_cache.cpp
+++ b/src/mongo/db/keys_collection_cache.cpp
@@ -42,12 +42,16 @@ KeysCollectionCache::KeysCollectionCache(std::string purpose, KeysCollectionClie
 StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext* opCtx) {
     LogicalTime newerThanThis;
 
+    decltype(_cache)::size_type originalSize = 0;
+
     {
         stdx::lock_guard<stdx::mutex> lk(_cacheMutex);
         auto iter = _cache.crbegin();
         if (iter != _cache.crend()) {
             newerThanThis = iter->second.getExpiresAt();
         }
+
+        originalSize = _cache.size();
     }
 
     auto refreshStatus = _client->getNewKeys(opCtx, _purpose, newerThanThis);
@@ -59,6 +63,13 @@ StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext
     auto& newKeys = refreshStatus.getValue();
 
     stdx::lock_guard<stdx::mutex> lk(_cacheMutex);
+    if (originalSize > _cache.size()) {
+        // _cache cleared while we getting the new keys, just return the newest key without
+        // touching the _cache so the next refresh will populate it properly.
+        // Note: newKeys are sorted.
+        return std::move(newKeys.back());
+    }
+
     for (auto&& key : newKeys) {
         _cache.emplace(std::make_pair(key.getExpiresAt(), std::move(key)));
     }
@@ -104,6 +115,7 @@ StatusWith<KeysCollectionDocument> KeysCollectionCache::getKey(const LogicalTime
 void KeysCollectionCache::resetCache() {
     // keys that read with non majority readConcern level can be rolled back.
     if (!_client->supportsMajorityReads()) {
+        stdx::lock_guard<stdx::mutex> lk(_cacheMutex);
         _cache.clear();
     }
 }
diff --git a/src/mongo/db/keys_collection_manager.cpp b/src/mongo/db/keys_collection_manager.cpp
index 889b822477c..4f5bd5b1643 100644
--- a/src/mongo/db/keys_collection_manager.cpp
+++ b/src/mongo/db/keys_collection_manager.cpp
@@ -191,6 +191,10 @@ bool KeysCollectionManager::hasSeenKeys() {
     return _refresher.hasSeenKeys();
 }
 
+void KeysCollectionManager::clearCache() {
+    _keysCache.resetCache();
+}
+
 void KeysCollectionManager::PeriodicRunner::refreshNow(OperationContext* opCtx) {
     auto refreshRequest = [this]() {
         stdx::lock_guard<stdx::mutex> lk(_mutex);
diff --git a/src/mongo/db/keys_collection_manager.h b/src/mongo/db/keys_collection_manager.h
index 5826e2228ee..2f6bec1df36 100644
--- a/src/mongo/db/keys_collection_manager.h
+++ b/src/mongo/db/keys_collection_manager.h
@@ -110,6 +110,11 @@ public:
      */
     bool hasSeenKeys();
 
+    /**
+     * Clears the in memory cache of the keys.
+     */
+    void clearCache();
+
 private:
     /**
      * This is responsible for periodically performing refresh in the background.
diff --git a/src/mongo/db/logical_time_validator.cpp b/src/mongo/db/logical_time_validator.cpp
index 4e3f3249171..085aef29e50 100644
--- a/src/mongo/db/logical_time_validator.cpp
+++ b/src/mongo/db/logical_time_validator.cpp
@@ -190,11 +190,10 @@ bool LogicalTimeValidator::shouldGossipLogicalTime() {
     return _getKeyManagerCopy()->hasSeenKeys();
 }
 
-void LogicalTimeValidator::resetKeyManagerCache(ServiceContext* service) {
+void LogicalTimeValidator::resetKeyManagerCache() {
     log() << "Resetting key manager cache";
     if (auto keyManager = _getKeyManagerCopy()) {
-        keyManager->stopMonitoring();
-        keyManager->startMonitoring(service);
+        keyManager->clearCache();
         _lastSeenValidTime = SignedLogicalTime();
         _timeProofService.resetCache();
     }
diff --git a/src/mongo/db/logical_time_validator.h b/src/mongo/db/logical_time_validator.h
index c05a0e5811b..17ea7379d33 100644
--- a/src/mongo/db/logical_time_validator.h
+++ b/src/mongo/db/logical_time_validator.h
@@ -110,7 +110,7 @@ public:
     /**
      * Reset the key manager cache of keys.
      */
-    void resetKeyManagerCache(ServiceContext* service);
+    void resetKeyManagerCache();
 
 private:
     /**
diff --git a/src/mongo/db/logical_time_validator_test.cpp b/src/mongo/db/logical_time_validator_test.cpp
index 461689a61e6..90eddb5ec98 100644
--- a/src/mongo/db/logical_time_validator_test.cpp
+++ b/src/mongo/db/logical_time_validator_test.cpp
@@ -181,5 +181,23 @@ TEST_F(LogicalTimeValidatorTest, ShouldGossipLogicalTimeIsFalseUntilKeysAreFound
     ASSERT_OK(validator()->validate(operationContext(), newTime));
 }
 
+TEST_F(LogicalTimeValidatorTest, CanSignTimesAfterReset) {
+    validator()->enableKeyGenerator(operationContext(), true);
+
+    LogicalTime t1(Timestamp(10, 0));
+    auto newTime = validator()->trySignLogicalTime(t1);
+
+    ASSERT_EQ(t1.asTimestamp(), newTime.getTime().asTimestamp());
+    ASSERT_TRUE(newTime.getProof());
+
+    validator()->resetKeyManagerCache();
+
+    LogicalTime t2(Timestamp(20, 0));
+    auto newTime2 = validator()->trySignLogicalTime(t2);
+
+    ASSERT_EQ(t2.asTimestamp(), newTime2.getTime().asTimestamp());
+    ASSERT_TRUE(newTime2.getProof());
+}
+
 }  // unnamed namespace
 }  // namespace mongo
diff --git a/src/mongo/db/repl/rollback_impl.cpp b/src/mongo/db/repl/rollback_impl.cpp
index bfed02e8d60..741249607ee 100644
--- a/src/mongo/db/repl/rollback_impl.cpp
+++ b/src/mongo/db/repl/rollback_impl.cpp
@@ -148,7 +148,7 @@ Status RollbackImpl::runRollback(OperationContext* opCtx) {
     ON_BLOCK_EXIT([this, opCtx] {
         auto validator = LogicalTimeValidator::get(opCtx);
         if (validator) {
-            validator->resetKeyManagerCache(opCtx->getClient()->getServiceContext());
+            validator->resetKeyManagerCache();
         }
 
         _checkShardIdentityRollback(opCtx);
diff --git a/src/mongo/db/repl/rs_rollback.cpp b/src/mongo/db/repl/rs_rollback.cpp
index 5b76eb1c42f..729348e07bb 100644
--- a/src/mongo/db/repl/rs_rollback.cpp
+++ b/src/mongo/db/repl/rs_rollback.cpp
@@ -53,6 +53,7 @@
 #include "mongo/db/dbhelpers.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/logical_session_id.h"
+#include "mongo/db/logical_time_validator.h"
 #include "mongo/db/ops/delete.h"
 #include "mongo/db/ops/update.h"
 #include "mongo/db/ops/update_lifecycle_impl.h"
@@ -1421,6 +1422,10 @@ void rollback_internal::syncFixUp(OperationContext* opCtx,
         SessionCatalog::get(opCtx)->invalidateSessions(opCtx, boost::none);
     }
 
+    if (auto validator = LogicalTimeValidator::get(opCtx)) {
+        validator->resetKeyManagerCache();
+    }
+
     // Reload the lastAppliedOpTime and lastDurableOpTime value in the replcoord and the
     // lastAppliedHash value in bgsync to reflect our new last op. The rollback common point does
     // not necessarily represent a consistent database state. For example, on a secondary, we may