1 files changed, 98 insertions, 29 deletions

diff --git a/src/mongo/db/keys_collection_cache.cpp b/src/mongo/db/keys_collection_cache.cpp
index 5cf62f8da12..87f9e798a33 100644
--- a/src/mongo/db/keys_collection_cache.cpp
+++ b/src/mongo/db/keys_collection_cache.cpp
@@ -42,20 +42,6 @@ KeysCollectionCache::KeysCollectionCache(std::string purpose, KeysCollectionClie
     : _purpose(std::move(purpose)), _client(client) {}
 
 StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext* opCtx) {
-    LogicalTime newerThanThis;
-
-    decltype(_cache)::size_type originalSize = 0;
-
-    {
-        stdx::lock_guard<Latch> lk(_cacheMutex);
-        auto iter = _cache.crbegin();
-        if (iter != _cache.crend()) {
-            newerThanThis = iter->second.getExpiresAt();
-        }
-
-        originalSize = _cache.size();
-    }
-
     // Don't allow this to read during initial sync because it will read at the initialDataTimestamp
     // and that could conflict with reconstructing prepared transactions using the
     // initialDataTimestamp as the prepareTimestamp.
@@ -65,7 +51,32 @@ StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext
                 "Cannot refresh keys collection cache during initial sync"};
     }
 
-    auto refreshStatus = _client->getNewKeys(opCtx, _purpose, newerThanThis, true);
+    auto refreshStatus = _refreshExternalKeys(opCtx);
+
+    if (!refreshStatus.isOK()) {
+        return refreshStatus;
+    }
+
+    return _refreshInternalKeys(opCtx);
+}
+
+
+StatusWith<KeysCollectionDocument> KeysCollectionCache::_refreshInternalKeys(
+    OperationContext* opCtx) {
+    LogicalTime newerThanThis;
+    decltype(_internalKeysCache)::size_type originalSize = 0;
+
+    {
+        stdx::lock_guard<Latch> lk(_cacheMutex);
+        auto iter = _internalKeysCache.crbegin();
+        if (iter != _internalKeysCache.crend()) {
+            newerThanThis = iter->second.getExpiresAt();
+        }
+
+        originalSize = _internalKeysCache.size();
+    }
+
+    auto refreshStatus = _client->getNewInternalKeys(opCtx, _purpose, newerThanThis, true);
 
     if (!refreshStatus.isOK()) {
         return refreshStatus.getStatus();
@@ -74,9 +85,9 @@ StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext
     auto& newKeys = refreshStatus.getValue();
 
     stdx::lock_guard<Latch> 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.
+    if (originalSize > _internalKeysCache.size()) {
+        // _internalKeysCache cleared while we were getting the new keys, just return the newest key
+        // without touching the _internalKeysCache so the next refresh will populate it properly.
         // Note: newKeys are sorted.
         if (!newKeys.empty()) {
             return std::move(newKeys.back());
@@ -84,38 +95,95 @@ StatusWith<KeysCollectionDocument> KeysCollectionCache::refresh(OperationContext
     }
 
     for (auto&& key : newKeys) {
-        _cache.emplace(std::make_pair(key.getExpiresAt(), std::move(key)));
+        _internalKeysCache.emplace(std::make_pair(key.getExpiresAt(), std::move(key)));
     }
 
-    if (_cache.empty()) {
+    if (_internalKeysCache.empty()) {
         return {ErrorCodes::KeyNotFound, "No keys found after refresh"};
     }
 
-    return _cache.crbegin()->second;
+    return _internalKeysCache.crbegin()->second;
+}
+
+Status KeysCollectionCache::_refreshExternalKeys(OperationContext* opCtx) {
+    decltype(_externalKeysCache)::size_type originalSize = 0;
+
+    {
+        stdx::lock_guard<Latch> lk(_cacheMutex);
+        originalSize = _externalKeysCache.size();
+    }
+
+    auto refreshStatus = _client->getNewExternalKeys(opCtx, _purpose, LogicalTime(), true);
+
+    if (!refreshStatus.isOK()) {
+        return refreshStatus.getStatus();
+    }
+
+    auto& newKeys = refreshStatus.getValue();
+
+    stdx::lock_guard<Latch> lk(_cacheMutex);
+    if (originalSize > _externalKeysCache.size()) {
+        // _externalKeysCache cleared while we were getting the new keys, just return so the next
+        // refresh will populate it properly.
+        return Status::OK();
+    }
+
+    for (auto&& key : newKeys) {
+        _externalKeysCache[key.getKeyId()].emplace(key.getReplicaSetName(), std::move(key));
+    }
+
+    return Status::OK();
 }
 
-StatusWith<KeysCollectionDocument> KeysCollectionCache::getKeyById(long long keyId,
-                                                                   const LogicalTime& forThisTime) {
+StatusWith<KeysCollectionDocument> KeysCollectionCache::getInternalKeyById(
+    long long keyId, const LogicalTime& forThisTime) {
     stdx::lock_guard<Latch> lk(_cacheMutex);
 
-    for (auto iter = _cache.lower_bound(forThisTime); iter != _cache.cend(); ++iter) {
+    for (auto iter = _internalKeysCache.lower_bound(forThisTime); iter != _internalKeysCache.cend();
+         ++iter) {
         if (iter->second.getKeyId() == keyId) {
             return iter->second;
         }
     }
 
     return {ErrorCodes::KeyNotFound,
-            str::stream() << "Cache Reader No keys found for " << _purpose
+            str::stream() << "Cache Reader No internal keys found for " << _purpose
                           << " that is valid for time: " << forThisTime.toString()
                           << " with id: " << keyId};
 }
 
-StatusWith<KeysCollectionDocument> KeysCollectionCache::getKey(const LogicalTime& forThisTime) {
+StatusWith<std::vector<ExternalKeysCollectionDocument>> KeysCollectionCache::getExternalKeysById(
+    long long keyId, const LogicalTime& forThisTime) {
+    stdx::lock_guard<Latch> lk(_cacheMutex);
+    std::vector<ExternalKeysCollectionDocument> keys;
+
+    auto keysIter = _externalKeysCache.find(keyId);
+
+    if (keysIter == _externalKeysCache.end()) {
+        return {ErrorCodes::KeyNotFound,
+                str::stream() << "Cache Reader No external keys found for " << _purpose
+                              << " with id: " << keyId};
+    }
+
+    invariant(!keysIter->second.empty());
+
+    for (auto keyIter = keysIter->second.begin(); keyIter != keysIter->second.end(); keyIter++) {
+        auto key = keyIter->second;
+        if (key.getExpiresAt() > forThisTime) {
+            keys.push_back(key);
+        }
+    }
+
+    return std::move(keys);
+}
+
+StatusWith<KeysCollectionDocument> KeysCollectionCache::getInternalKey(
+    const LogicalTime& forThisTime) {
     stdx::lock_guard<Latch> lk(_cacheMutex);
 
-    auto iter = _cache.upper_bound(forThisTime);
+    auto iter = _internalKeysCache.upper_bound(forThisTime);
 
-    if (iter == _cache.cend()) {
+    if (iter == _internalKeysCache.cend()) {
         return {ErrorCodes::KeyNotFound,
                 str::stream() << "No key found that is valid for " << forThisTime.toString()};
     }
@@ -127,7 +195,8 @@ void KeysCollectionCache::resetCache() {
     // keys that read with non majority readConcern level can be rolled back.
     if (!_client->supportsMajorityReads()) {
         stdx::lock_guard<Latch> lk(_cacheMutex);
-        _cache.clear();
+        _internalKeysCache.clear();
+        _externalKeysCache.clear();
     }
 }