SERVER-49190 Expose a function to iterate over all the entries in the ReadThroughCache

3 files changed, 93 insertions, 3 deletions

diff --git a/src/mongo/db/auth/authorization_manager_impl.cpp b/src/mongo/db/auth/authorization_manager_impl.cpp
index c7c27d3c461..c4bca48c570 100644
--- a/src/mongo/db/auth/authorization_manager_impl.cpp
+++ b/src/mongo/db/auth/authorization_manager_impl.cpp
@@ -608,7 +608,7 @@ void AuthorizationManagerImpl::invalidateUsersFromDB(OperationContext* opCtx, St
     LOGV2_DEBUG(20236, 2, "Invalidating all users from database", "database"_attr = dbname);
     _updateCacheGeneration();
     _authSchemaVersionCache.invalidateAll();
-    _userCache.invalidateIf(
+    _userCache.invalidateIfKey(
         [&](const UserRequest& userRequest) { return userRequest.name.getDB() == dbname; });
 }
 
diff --git a/src/mongo/util/read_through_cache.h b/src/mongo/util/read_through_cache.h
index 656f5c04a42..ffa4d7295e2 100644
--- a/src/mongo/util/read_through_cache.h
+++ b/src/mongo/util/read_through_cache.h
@@ -331,8 +331,11 @@ public:
         _cache.invalidate(key);
     }
 
+    /**
+     * Invalidates all cached entries and in progress lookups with keys that matches the preidcate.
+     */
     template <typename Pred>
-    void invalidateIf(const Pred& predicate) {
+    void invalidateIfKey(const Pred& predicate) {
         stdx::lock_guard lg(_mutex);
         for (auto& entry : _inProgressLookups) {
             if (predicate(entry.first))
@@ -341,8 +344,18 @@ public:
         _cache.invalidateIf([&](const Key& key, const StoredValue*) { return predicate(key); });
     }
 
+    /**
+     * Invalidates all cached entries with stored values that matches the preidcate.
+     */
+    template <typename Pred>
+    void invalidateIfCachedValue(const Pred& predicate) {
+        stdx::lock_guard lg(_mutex);
+        _cache.invalidateIf(
+            [&](const Key&, const StoredValue* value) { return predicate(value->value); });
+    }
+
     void invalidateAll() {
-        invalidateIf([](const Key&) { return true; });
+        invalidateIfKey([](const Key&) { return true; });
     }
 
     /**
diff --git a/src/mongo/util/read_through_cache_test.cpp b/src/mongo/util/read_through_cache_test.cpp
index b2b4148c744..0093f6aa8dc 100644
--- a/src/mongo/util/read_through_cache_test.cpp
+++ b/src/mongo/util/read_through_cache_test.cpp
@@ -172,6 +172,83 @@ TEST_F(ReadThroughCacheTest, FetchInvalidateAndRefetch) {
         }));
 }
 
+TEST_F(ReadThroughCacheTest, FetchInvalidateKeyAndRefetch) {
+    auto fnTest = [&](auto cache) {
+        for (int i = 1; i <= 3; i++) {
+            auto value = cache.acquire(_opCtx, "TestKey");
+            ASSERT(value);
+            ASSERT_EQ(100 * i, value->counter);
+            ASSERT_EQ(i, cache.countLookups);
+
+            ASSERT(cache.acquire(_opCtx, "TestKey"));
+            ASSERT_EQ(i, cache.countLookups);
+
+            cache.invalidateIfKey([](const std::string& key) { return key == "TestKey"; });
+        }
+    };
+
+    fnTest(CacheWithThreadPool<Cache>(
+        getServiceContext(),
+        1,
+        [&, nextValue = 0](
+            OperationContext*, const std::string& key, const Cache::ValueHandle&) mutable {
+            ASSERT_EQ("TestKey", key);
+            return Cache::LookupResult(CachedValue(100 * ++nextValue));
+        }));
+
+    fnTest(CacheWithThreadPool<CausallyConsistentCache>(
+        getServiceContext(),
+        1,
+        [&, nextValue = 0](OperationContext*,
+                           const std::string& key,
+                           const CausallyConsistentCache::ValueHandle&,
+                           const Timestamp& timeInStore) mutable {
+            ASSERT_EQ("TestKey", key);
+            ++nextValue;
+            return CausallyConsistentCache::LookupResult(CachedValue(100 * nextValue),
+                                                         Timestamp(nextValue));
+        }));
+}
+
+TEST_F(ReadThroughCacheTest, FetchInvalidateValueAndRefetch) {
+    auto fnTest = [&](auto cache) {
+        for (int i = 1; i <= 3; i++) {
+            auto value = cache.acquire(_opCtx, "TestKey");
+            ASSERT(value);
+            ASSERT_EQ(100 * i, value->counter);
+            ASSERT_EQ(i, cache.countLookups);
+
+            ASSERT(cache.acquire(_opCtx, "TestKey"));
+            ASSERT_EQ(i, cache.countLookups);
+
+            cache.invalidateIfCachedValue(
+                [i](const CachedValue& value) { return value.counter == 100 * i; });
+        }
+    };
+
+    fnTest(CacheWithThreadPool<Cache>(
+        getServiceContext(),
+        1,
+        [&, nextValue = 0](
+            OperationContext*, const std::string& key, const Cache::ValueHandle&) mutable {
+            ASSERT_EQ("TestKey", key);
+            return Cache::LookupResult(CachedValue(100 * ++nextValue));
+        }));
+
+    fnTest(CacheWithThreadPool<CausallyConsistentCache>(
+        getServiceContext(),
+        1,
+        [&, nextValue = 0](OperationContext*,
+                           const std::string& key,
+                           const CausallyConsistentCache::ValueHandle&,
+                           const Timestamp& timeInStore) mutable {
+            ASSERT_EQ("TestKey", key);
+            ++nextValue;
+            return CausallyConsistentCache::LookupResult(CachedValue(100 * nextValue),
+                                                         Timestamp(nextValue));
+        }));
+}
+
 TEST_F(ReadThroughCacheTest, FailedLookup) {
     auto fnTest = [&](auto cache) {
         ASSERT_THROWS_CODE(