1 files changed, 44 insertions, 3 deletions

diff --git a/src/mongo/db/keys_collection_cache.h b/src/mongo/db/keys_collection_cache.h
index 17c532b72fa..a37b4e45c65 100644
--- a/src/mongo/db/keys_collection_cache.h
+++ b/src/mongo/db/keys_collection_cache.h
@@ -55,8 +55,27 @@ public:
      */
     StatusWith<KeysCollectionDocument> refresh(OperationContext* opCtx);
 
-    StatusWith<KeysCollectionDocument> getKey(const LogicalTime& forThisTime);
-    StatusWith<KeysCollectionDocument> getKeyById(long long keyId, const LogicalTime& forThisTime);
+    /**
+     * Returns the internal key (see definition below) with an expiresAt value greater than
+     * forThisTime. Returns KeyNotFound if there is no such key.
+     */
+    StatusWith<KeysCollectionDocument> getInternalKey(const LogicalTime& forThisTime);
+
+    /**
+     * Returns the internal key (see definition below) with the given keyId and an expiresAt value
+     * greater than forThisTime. There should only be one matching key since keyId is unique for
+     * keys generated within a cluster. Returns KeyNotFound if there is no such key.
+     */
+    StatusWith<KeysCollectionDocument> getInternalKeyById(long long keyId,
+                                                          const LogicalTime& forThisTime);
+
+    /**
+     * Returns the external keys (see definition below) with the given keyId and an expiresAt value
+     * greater than forThisTime. There are a variable number of matching keys since keyId is not
+     * necessarily unique across clusters. Returns KeyNotFound if there are no such keys.
+     */
+    StatusWith<std::vector<ExternalKeysCollectionDocument>> getExternalKeysById(
+        long long keyId, const LogicalTime& forThisTime);
 
     /**
      * Resets the cache of keys if the client doesnt allow readConcern level:majority reads.
@@ -65,11 +84,33 @@ public:
     void resetCache();
 
 private:
+    /**
+     * Checks if there are new internal key documents (see definition below) with expiresAt greater
+     * than the latest internal key document's expiresAt. Returns KeyNotFound if _internalKeysCache
+     * is empty after refresh.
+     */
+    StatusWith<KeysCollectionDocument> _refreshInternalKeys(OperationContext* opCtx);
+
+    /**
+     * Checks if there are new external key documents (see definition below). Does not return
+     * KeyNotFound if _externalKeysCache is empty after refresh.
+     */
+    Status _refreshExternalKeys(OperationContext* opCtx);
+
     const std::string _purpose;
     KeysCollectionClient* const _client;
 
     Mutex _cacheMutex = MONGO_MAKE_LATCH("KeysCollectionCache::_cacheMutex");
-    std::map<LogicalTime, KeysCollectionDocument> _cache;  // expiresAt -> KeysDocument
+
+    // Stores keys for signing and validating cluster times created by the cluster that this node
+    // is in.
+    std::map<LogicalTime, KeysCollectionDocument> _internalKeysCache;  // expiresAt -> KeysDocument
+
+    // Stores keys for validating cluster times created by other clusters. These key documents
+    // cannot be stored in a regular map like _internalKeysCache since expiresAt and keyId are not
+    // necessarily unique across clusters so there is chance of collision.
+    stdx::unordered_map<long long, StringMap<ExternalKeysCollectionDocument>>
+        _externalKeysCache;  // keyId -> (replicaSetName -> ExternalKeysDocument)
 };
 
 }  // namespace mongo