[ML] Fix concurrent LFU cache count invariant under lock timeouts

docs/changelog/3090.yaml

@@ -0,0 +1,5 @@
+pr: 3090
+summary: "Fix flaky concurrent LFU cache count invariant under lock timeouts"
+area: Machine Learning
+type: bug
+issues: []

include/core/CCompressedLfuCache.h

@@ -125,8 +125,14 @@ class CCompressedLfuCache {
 
         auto compressedKey = m_CompressKey(m_Dictionary, key);
 
+        // Count every lookup exactly once, even if we fail to acquire the read
+        // lock below. The fall-through miss path can still add a count (or record
+        // a lost update), so counting only inside the read guard would let the
+        // bookkeeping totals exceed the lookup count when the read lock times out
+        // under contention.
+        ++m_NumberLookups;
+
         if (this->guardRead(TIME_OUT, [&] {
-                ++m_NumberLookups;
                 auto hit = this->hit(compressedKey);
                 if (hit != nullptr) {
                     ++m_NumberHits;
@@ -282,9 +288,19 @@ class CCompressedLfuCache {
                           << m_ItemsMemoryUsage << ".");
                 result = false;
             }
-            // We may be missing counts if a single addition caused multiple
-            // evictions, or if updates can time out
-            if (m_NumberLookups - m_LostCount != totalCount) {
+            // Counts can legitimately be *lost* under concurrency: an item hit
+            // under the read lock may be evicted before its count is incremented
+            // under the write lock, a single insert can evict several items, and
+            // write updates can time out. They must never be *fabricated* though,
+            // so the surviving total may fall short of, but must not exceed, the
+            // number of accounted lookups. When updates cannot time out (the
+            // single-threaded cache) there is no concurrency and the relation is
+            // exact.
+            std::uint64_t accountedLookups{m_NumberLookups - m_LostCount};
+            bool countInvariantHolds{this->updatesCanTimeOut()
+                                         ? totalCount <= accountedLookups
+                                         : totalCount == accountedLookups};
+            if (countInvariantHolds == false) {
                 LOG_ERROR(<< "Count mismatch " << m_NumberLookups.load() << " (less "
                           << m_LostCount.load() << " lost) vs " << totalCount << ".");
                 result = false;

lib/core/unittest/CCompressedLfuCacheTest.cc

@@ -296,6 +296,56 @@ BOOST_AUTO_TEST_CASE(testConcurrentReadsAndWrites) {
     BOOST_TEST_REQUIRE(cache.checkInvariants());
 }
 
+BOOST_AUTO_TEST_CASE(testConcurrentCountInvariantWithTimeouts) {
+
+    // Aggressively drive lock timeouts by combining a zero maximum wait with
+    // heavy contention, then check the count bookkeeping invariant still holds.
+    //
+    // Regression test: a lookup whose read lock timed out used to skip counting
+    // the lookup, yet its fall-through miss path could still add a count (or
+    // record a lost update). That let the surviving counts exceed the number of
+    // accounted lookups, tripping checkInvariants() intermittently under load.
+
+    using TTaskVec = std::vector<std::function<void()>>;
+
+    // A zero maximum wait makes every contended lock acquisition time out, so
+    // both the read and write guards fail frequently under load.
+    TConcurrentStrStrCache cache{
+        32 * core::constants::BYTES_IN_KILOBYTES, std::chrono::milliseconds{0},
+        [](const TStrStrCache::TDictionary& dictionary,
+           const std::string& key) { return dictionary.word(key); }};
+
+    std::atomic<std::size_t> errorCount{0};
+
+    // Reuse a modest key space so most keys are cached, exercising the hit path
+    // whose count increment is where fabrication used to occur when the read
+    // lock had already timed out.
+    TTaskVec tasks;
+    for (std::size_t i = 0; i < 5000; ++i) {
+        std::string key{"a_long_key_" + std::to_string(i % 200)};
+        tasks.push_back([&cache, &errorCount, key] {
+            cache.lookup(key, [](std::string key_) { return key_; },
+                         [&](const std::string& value, bool) {
+                             if (value != key) {
+                                 ++errorCount;
+                             }
+                         });
+        });
+    }
+
+    {
+        core::CStaticThreadPool pool{8};
+        for (auto& task : tasks) {
+            pool.schedule(std::move(task));
+        }
+    }
+
+    BOOST_REQUIRE_EQUAL(0, errorCount.load());
+    LOG_DEBUG(<< "hit fraction = " << cache.hitFraction());
+
+    BOOST_TEST_REQUIRE(cache.checkInvariants());
+}
+
 BOOST_AUTO_TEST_CASE(testEvictionStrategyFrequency) {
 
     // Check that frequent items are retained.