mooncake-store/include/local_hot_cache.h

@@ -19,6 +19,16 @@
 
 namespace mooncake {
 
+/**
+ * @brief Token captured at async hot cache fill submission time.
+ * Invalidated when RemoveHotKey, BumpKeyGeneration, or Clear bumps
+ * generation/epoch.
+ */
+struct HotCachePutToken {
+    uint64_t cache_epoch = 0;
+    uint64_t key_generation = 0;
+};
+
 /**
  * @brief Memory block metadata for hot cache.
  */
@@ -61,6 +71,15 @@ class LocalHotCache {
      */
     bool PutHotKey(HotMemBlock* block);
 
+    /**
+     * @brief Insert a populated block only if its async fill token is still
+     * valid, checking the token and publishing atomically under one lock.
+     * If the token is stale (the key was removed/overwritten since the fill
+     * started), the block is returned to the pool instead of being published.
+     * @return true if the block was published, false if cancelled or on error.
+     */
+    bool PutHotKey(HotMemBlock* block, const HotCachePutToken& token);
+
     /**
      * @brief Check if the key exists in cache.
      * @param key Cache key: {request key}
@@ -91,6 +110,65 @@ class LocalHotCache {
      */
     bool TouchHotKey(const std::string& key);
 
+    /**
+     * @brief Remove a key from the hot cache immediately.
+     * Bumps the key generation so in-flight async fills are invalidated.
+     * @param key The key to remove from cache.
+     * @return true if a published cache entry was removed, false otherwise.
+     */
+    bool RemoveHotKey(const std::string& key);
+
+    /**
+     * @brief Remove a batch of keys from the hot cache under one lock.
+     * Bumps each key generation so in-flight async fills are invalidated.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveHotKeys(const std::vector<std::string>& keys);
+
+    /**
+     * @brief Remove cached keys matching a regex from the hot cache.
+     * Bumps key generation for matching published entries.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveHotKeysByRegex(const std::string& regex_pattern);
+
+    /**
+     * @brief Remove every published hot cache entry and invalidate async fills.
+     * @return number of published cache entries removed.
+     */
+    size_t RemoveAllHotKeys();
+
+    /**
+     * @brief Invalidate in-flight async fills for a key without evicting it.
+     */
+    void BumpKeyGeneration(const std::string& key);
+
+    /**
+     * @brief Invalidate in-flight async fills for multiple keys.
+     */
+    void BumpKeyGenerations(const std::vector<std::string>& keys);
+
+    /**
+     * @brief Invalidate all in-flight async fills without evicting entries.
+     */
+    void BumpCacheEpoch();
+
+    /**
+     * @brief Clear all hot cache entries and invalidate in-flight async fills.
+     */
+    void Clear();
+
+    /**
+     * @brief Capture the current put token for async hot cache fill validation.
+     */
+    HotCachePutToken AcquirePutToken(const std::string& key);
+
+    /**
+     * @brief Check whether an async put token is still valid.
+     */
+    bool IsPutTokenValid(const std::string& key,
+                         const HotCachePutToken& token) const;
+
     /**
      * @brief Get a free block for writing.
      * Detaches a block from the LRU tail (evicting if necessary) and returns
@@ -146,6 +224,11 @@ class LocalHotCache {
    private:
     // Drain deferred LRU touches: splice accessed blocks to front
     void drainDeferredTouches();
+    bool putHotKeyLocked(HotMemBlock* block);
+    bool removeHotKeyLocked(const std::string& key);
+    bool hasActiveBlockForKeyLocked(const std::string& key) const;
+    bool isPutTokenValidLocked(const std::string& key,
+                               const HotCachePutToken& token) const;
 
     size_t block_size_;  // Actual block size used by this cache
 
@@ -165,6 +248,9 @@ class LocalHotCache {
     // key -> iterator of lru_queue_
     std::unordered_map<std::string, std::list<HotMemBlock*>::iterator>
         key_to_lru_it_ GUARDED_BY(lru_mutex_);
+    std::unordered_map<std::string, uint64_t> key_generation_
+        GUARDED_BY(lru_mutex_);
+    std::atomic<uint64_t> cache_epoch_{0};
 };
 
 /**
@@ -175,13 +261,19 @@ struct HotCachePutTask {
     HotMemBlock* block;  // Pointer to the allocated block
     size_t size;
     std::shared_ptr<LocalHotCache> hot_cache;
+    HotCachePutToken token;
 
     // Default constructor for empty task
     HotCachePutTask() : block(nullptr), size(0), hot_cache(nullptr) {}
 
     HotCachePutTask(const std::string& k, const Slice& slice, HotMemBlock* blk,
-                    std::shared_ptr<LocalHotCache> cache)
-        : key(k), block(blk), size(slice.size), hot_cache(std::move(cache)) {
+                    std::shared_ptr<LocalHotCache> cache,
+                    HotCachePutToken put_token)
+        : key(k),
+          block(blk),
+          size(slice.size),
+          hot_cache(std::move(cache)),
+          token(put_token) {
         // No data copy here; memcpy is done by SubmitPutTask into block->addr.
     }
 };

mooncake-store/src/client_service.cpp

@@ -1490,6 +1490,10 @@ tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
         client_cfg.preferred_segment = local_hostname_;
     }
 
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     // Start put operation
     auto start_result = master_client_.PutStart(key, slice_lengths, client_cfg);
     if (!start_result) {
@@ -1596,6 +1600,10 @@ tl::expected<void, ErrorCode> Client::Upsert(const ObjectKey& key,
         client_cfg.preferred_segment = local_hostname_;
     }
 
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     // Start upsert operation
     auto start_result =
         master_client_.UpsertStart(key, slice_lengths, client_cfg);
@@ -1658,6 +1666,14 @@ tl::expected<void, ErrorCode> Client::Upsert(const ObjectKey& key,
         return tl::unexpected(err);
     }
 
+    // Success-side invalidation: a concurrent read between the pre-upsert
+    // RemoveHotKey() and UpsertEnd could have read the old value and submitted
+    // an async hot-cache fill with a still-valid token. Invalidate again now
+    // that the new value is committed so that stale fill cannot publish.
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     return {};
 }
 
@@ -1824,6 +1840,13 @@ void Client::StartBatchPut(std::vector<PutOperation>& ops,
     keys.reserve(ops.size());
     slice_lengths.reserve(ops.size());
 
+    if (hot_cache_) {
+        std::vector<std::string> hot_keys;
+        hot_keys.reserve(ops.size());
+        for (const auto& op : ops) hot_keys.emplace_back(op.key);
+        hot_cache_->RemoveHotKeys(hot_keys);
+    }
+
     for (const auto& op : ops) {
         keys.emplace_back(op.key);
 
@@ -1883,6 +1906,13 @@ void Client::StartBatchUpsert(std::vector<PutOperation>& ops,
     keys.reserve(ops.size());
     slice_lengths.reserve(ops.size());
 
+    if (hot_cache_) {
+        std::vector<std::string> hot_keys;
+        hot_keys.reserve(ops.size());
+        for (const auto& op : ops) hot_keys.emplace_back(op.key);
+        hot_cache_->RemoveHotKeys(hot_keys);
+    }
+
     for (const auto& op : ops) {
         keys.emplace_back(op.key);
 
@@ -2261,7 +2291,9 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
     }
 
     // Process successful operations
+    std::vector<std::string> finalized_keys;
     if (!successful_keys.empty()) {
+        finalized_keys.reserve(successful_keys.size());
         auto end_responses = master_client_.BatchUpsertEnd(successful_keys);
         if (end_responses.size() != successful_keys.size()) {
             LOG(ERROR) << "BatchUpsertEnd response size mismatch: expected "
@@ -2282,6 +2314,7 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
                                          "BatchUpsertEnd failed");
                 } else {
                     ops[op_idx].SetSuccess();
+                    finalized_keys.emplace_back(successful_keys[i]);
                     VLOG(1) << "Successfully completed upsert for key "
                             << successful_keys[i];
                 }
@@ -2289,6 +2322,15 @@ void Client::FinalizeBatchUpsert(std::vector<PutOperation>& ops) {
         }
     }
 
+    // Success-side invalidation for finalized upserts only: a concurrent read
+    // between StartBatchUpsert()'s pre-invalidation and BatchUpsertEnd could
+    // have read the old value and submitted an async hot-cache fill with a
+    // still-valid token. Invalidate again now that the new values are
+    // committed so those stale fills cannot publish.
+    if (hot_cache_ && !finalized_keys.empty()) {
+        hot_cache_->RemoveHotKeys(finalized_keys);
+    }
+
     // Process failed operations that need cleanup
     if (!failed_keys.empty()) {
         auto revoke_responses = master_client_.BatchUpsertRevoke(failed_keys);
@@ -2513,6 +2555,10 @@ std::vector<tl::expected<void, ErrorCode>> Client::BatchPut(
 }
 
 tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key, bool force) {
+    if (hot_cache_) {
+        hot_cache_->BumpKeyGeneration(key);
+    }
+
     auto result = master_client_.Remove(key, force);
     // if (storage_backend_) {
     //     storage_backend_->RemoveFile(key);
@@ -2520,11 +2566,20 @@ tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key, bool force) {
     if (!result) {
         return tl::unexpected(result.error());
     }
+
+    if (hot_cache_) {
+        hot_cache_->RemoveHotKey(key);
+    }
+
     return {};
 }
 
 tl::expected<long, ErrorCode> Client::RemoveByRegex(const ObjectKey& str,
                                                     bool force) {
+    if (hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+    }
+
     auto result = master_client_.RemoveByRegex(str, force);
     // if (storage_backend_) {
     //     storage_backend_->RemoveByRegex(str);
@@ -2532,20 +2587,50 @@ tl::expected<long, ErrorCode> Client::RemoveByRegex(const ObjectKey& str,
     if (!result) {
         return tl::unexpected(result.error());
     }
+    if (result.value() > 0 && hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+        hot_cache_->RemoveHotKeysByRegex(str);
+    }
     return result.value();
 }
 
 tl::expected<long, ErrorCode> Client::RemoveAll(bool force) {
+    if (hot_cache_) {
+        hot_cache_->BumpCacheEpoch();
+    }
+
     auto result = master_client_.RemoveAll(force);
     if (result && storage_backend_) {
         storage_backend_->RemoveAll();
     }
+    if (result && result.value() > 0 && hot_cache_) {
+        hot_cache_->RemoveAllHotKeys();
+    }
     return result;
 }
 
 std::vector<tl::expected<void, ErrorCode>> Client::BatchRemove(
     const std::vector<ObjectKey>& keys, bool force) {
-    return master_client_.BatchRemove(keys, force);
+    if (hot_cache_) {
+        hot_cache_->BumpKeyGenerations(keys);
+    }
+
+    auto results = master_client_.BatchRemove(keys, force);
+
+    if (hot_cache_) {
+        std::vector<std::string> removed_keys;
+        removed_keys.reserve(std::min(keys.size(), results.size()));
+        for (size_t i = 0; i < keys.size(); ++i) {
+            if (i < results.size() && results[i].has_value()) {
+                removed_keys.emplace_back(keys[i]);
+            }
+        }
+        if (!removed_keys.empty()) {
+            hot_cache_->RemoveHotKeys(removed_keys);
+        }
+    }
+
+    return results;
 }
 
 tl::expected<void, ErrorCode> Client::EvictDiskReplica(

mooncake-store/src/local_hot_cache.cpp

@@ -4,6 +4,7 @@
 #include <cstring>
 #include <cstdlib>
 #include <cstdint>
+#include <regex>
 #include <shared_mutex>
 #include <glog/logging.h>
 
@@ -75,7 +76,23 @@ LocalHotCache::~LocalHotCache() {
 
 bool LocalHotCache::PutHotKey(HotMemBlock* block) {
     std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    return putHotKeyLocked(block);
+}
 
+bool LocalHotCache::PutHotKey(HotMemBlock* block,
+                              const HotCachePutToken& token) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    // Validate the token and publish under the same lock so a concurrent
+    // Remove/Bump cannot slip in between the check and the publish.
+    if (block && !block->key_.empty() &&
+        !isPutTokenValidLocked(block->key_, token)) {
+        // Stale async fill: drop the data and return the block to the pool.
+        block->key_.clear();
+    }
+    return putHotKeyLocked(block);
+}
+
+bool LocalHotCache::putHotKeyLocked(HotMemBlock* block) {
     // Drain deferred LRU touches
     drainDeferredTouches();
 
@@ -92,7 +109,8 @@ bool LocalHotCache::PutHotKey(HotMemBlock* block) {
 
     // Race condition check: did someone else insert this key while we were
     // copying
-    if (key_to_lru_it_.find(key) != key_to_lru_it_.end()) {
+    if (key_to_lru_it_.find(key) != key_to_lru_it_.end() ||
+        hasActiveBlockForKeyLocked(key)) {
         // Lost race -> Return to lru tail as free block
         block->key_.clear();
         block->ref_count = 0;
@@ -134,17 +152,211 @@ HotMemBlock* LocalHotCache::GetHotKey(const std::string& key) {
 }
 
 void LocalHotCache::ReleaseHotKey(const std::string& key) {
-    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    auto it = key_to_lru_it_.find(key);
+    if (it != key_to_lru_it_.end()) {
+        HotMemBlock* block = *(it->second);
+        if (block && block->ref_count > 0) {
+            block->ref_count--;
+            return;
+        }
+    }
+
+    // The entry may have been removed while a reader still holds the block.
+    // Keep the key on active removed blocks so release can find and retire it.
+    for (auto& owned_block : blocks_) {
+        HotMemBlock* block = owned_block.get();
+        if (block && block->key_ == key && block->ref_count > 0) {
+            block->ref_count--;
+            if (block->ref_count == 0 &&
+                key_to_lru_it_.find(key) == key_to_lru_it_.end()) {
+                block->key_.clear();
+                block->accessed.store(false, std::memory_order_relaxed);
+            }
+            return;
+        }
+    }
+}
+
+bool LocalHotCache::hasActiveBlockForKeyLocked(const std::string& key) const {
+    for (const auto& owned_block : blocks_) {
+        const HotMemBlock* block = owned_block.get();
+        if (block && block->key_ == key && block->ref_count > 0) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool LocalHotCache::removeHotKeyLocked(const std::string& key) {
     auto it = key_to_lru_it_.find(key);
     if (it == key_to_lru_it_.end()) {
+        return false;
+    }
+
+    HotMemBlock* block = *(it->second);
+    lru_queue_.erase(it->second);
+    key_to_lru_it_.erase(it);
+
+    if (block->ref_count == 0) {
+        block->key_.clear();
+        block->accessed.store(false, std::memory_order_relaxed);
+    }
+    lru_queue_.push_back(block);
+
+    return true;
+}
+
+bool LocalHotCache::RemoveHotKey(const std::string& key) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+
+    key_generation_[key]++;
+    drainDeferredTouches();
+
+    const bool removed = removeHotKeyLocked(key);
+    if (removed) {
+        VLOG(2) << "Removed hot key: " << key << " from hot cache";
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveHotKeys(const std::vector<std::string>& keys) {
+    if (keys.empty()) {
+        return 0;
+    }
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    drainDeferredTouches();
+
+    size_t removed = 0;
+    for (const auto& key : keys) {
+        key_generation_[key]++;
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveHotKeysByRegex(const std::string& regex_pattern) {
+    std::regex pattern;
+    try {
+        pattern = std::regex(regex_pattern, std::regex::ECMAScript);
+    } catch (const std::regex_error& e) {
+        LOG(ERROR) << "RemoveHotKeysByRegex: invalid pattern: " << regex_pattern
+                   << ", error: " << e.what();
+        return 0;
+    }
+
+    std::vector<std::string> matching_keys;
+    {
+        std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+        matching_keys.reserve(key_to_lru_it_.size());
+        for (const auto& [key, _] : key_to_lru_it_) {
+            if (std::regex_search(key, pattern)) {
+                matching_keys.emplace_back(key);
+            }
+        }
+    }
+
+    if (matching_keys.empty()) {
+        return 0;
+    }
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    drainDeferredTouches();
+
+    size_t removed = 0;
+    for (const auto& key : matching_keys) {
+        key_generation_[key]++;
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    return removed;
+}
+
+size_t LocalHotCache::RemoveAllHotKeys() {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+
+    std::vector<std::string> keys;
+    keys.reserve(key_to_lru_it_.size());
+    for (const auto& [key, _] : key_to_lru_it_) {
+        keys.emplace_back(key);
+    }
+
+    size_t removed = 0;
+    for (const auto& key : keys) {
+        if (removeHotKeyLocked(key)) {
+            ++removed;
+        }
+    }
+    key_generation_.clear();
+    return removed;
+}
+
+void LocalHotCache::BumpKeyGeneration(const std::string& key) {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    key_generation_[key]++;
+}
+
+void LocalHotCache::BumpKeyGenerations(const std::vector<std::string>& keys) {
+    if (keys.empty()) {
         return;
     }
-    HotMemBlock* block = *(it->second);
-    if (block) {
-        block->ref_count--;
+
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    for (const auto& key : keys) {
+        key_generation_[key]++;
     }
 }
 
+void LocalHotCache::BumpCacheEpoch() {
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+}
+
+void LocalHotCache::Clear() {
+    std::unique_lock<std::shared_mutex> lk(lru_mutex_);
+    cache_epoch_.fetch_add(1, std::memory_order_relaxed);
+
+    std::vector<std::string> keys;
+    keys.reserve(key_to_lru_it_.size());
+    for (const auto& [key, _] : key_to_lru_it_) {
+        keys.emplace_back(key);
+    }
+
+    for (const auto& key : keys) {
+        removeHotKeyLocked(key);
+    }
+    key_generation_.clear();
+}
+
+HotCachePutToken LocalHotCache::AcquirePutToken(const std::string& key) {
+    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    HotCachePutToken token;
+    token.cache_epoch = cache_epoch_.load(std::memory_order_relaxed);
+    auto it = key_generation_.find(key);
+    token.key_generation = (it != key_generation_.end()) ? it->second : 0;
+    return token;
+}
+
+bool LocalHotCache::isPutTokenValidLocked(const std::string& key,
+                                          const HotCachePutToken& token) const {
+    if (token.cache_epoch != cache_epoch_.load(std::memory_order_relaxed)) {
+        return false;
+    }
+    auto it = key_generation_.find(key);
+    const uint64_t current_gen = (it != key_generation_.end()) ? it->second : 0;
+    return token.key_generation == current_gen;
+}
+
+bool LocalHotCache::IsPutTokenValid(const std::string& key,
+                                    const HotCachePutToken& token) const {
+    std::shared_lock<std::shared_mutex> lk(lru_mutex_);
+    return isPutTokenValidLocked(key, token);
+}
+
 bool LocalHotCache::TouchHotKey(const std::string& key) {
     std::shared_lock<std::shared_mutex> lk(lru_mutex_);
     auto it = key_to_lru_it_.find(key);
@@ -303,6 +515,7 @@ bool LocalHotCacheHandler::SubmitPutTask(const std::string& key,
     }
 
     // Try to get a free block (may evict from LRU tail)
+    HotCachePutToken token = hot_cache_->AcquirePutToken(key);
     HotMemBlock* block = hot_cache_->GetFreeBlock();
     if (!block) {
         LOG(ERROR) << "Hot cache is fully in-use, fail to get a free block: "
@@ -330,7 +543,7 @@ bool LocalHotCacheHandler::SubmitPutTask(const std::string& key,
     block->size = slice.size;
     block->key_ = key;  // Set key for insertion
 
-    HotCachePutTask task(key, slice, block, hot_cache_);
+    HotCachePutTask task(key, slice, block, hot_cache_, token);
 
     {
         std::lock_guard<std::mutex> lock(queue_mutex_);
@@ -373,11 +586,12 @@ void LocalHotCacheHandler::workerThread() {
         // Execute the task if we have one
         if (task.hot_cache && task.block) {
             try {
-                // Insert the pre-filled block into LRU
-                if (task.hot_cache->PutHotKey(task.block)) {
+                // Validate token and publish atomically: a Remove/Bump that
+                // races after the check can no longer resurrect a stale fill.
+                if (task.hot_cache->PutHotKey(task.block, task.token)) {
                     VLOG(2) << "Put task completed: " << task.key;
                 } else {
-                    VLOG(2) << "Put task skipped: " << task.key;
+                    VLOG(2) << "Put task skipped or cancelled: " << task.key;
                 }
             } catch (const std::exception& e) {
                 LOG(ERROR) << "Exception during async hot cache put for key "

mooncake-store/tests/client_local_hot_cache_test.cpp

@@ -476,11 +476,40 @@ TEST_F(LocalHotCacheTest, GetHotKeyProtectsBlockFromReuse) {
     HotMemBlock* block2 = cache.GetHotKey("key2");
     ASSERT_NE(block2, nullptr);
     VerifySliceData(block2, 1024, 'B');
+    cache.ReleaseHotKey("key2");
 
     // key1 should be evicted since we only have 1 block
     EXPECT_FALSE(cache.HasHotKey("key1"));
 }
 
+// Token check and publish must be atomic: a generation bump that races after
+// the token was captured must cancel the fill instead of resurrecting it.
+TEST_F(LocalHotCacheTest, PutHotKeyWithTokenRejectsStaleFill) {
+    const size_t cache_size = 16 * 1024 * 1024;  // 1 block
+    LocalHotCache cache(cache_size);
+
+    // Stale: token captured, then key generation bumped (as Remove would).
+    HotCachePutToken stale = cache.AcquirePutToken("k");
+    cache.BumpKeyGeneration("k");
+    HotMemBlock* b1 = cache.GetFreeBlock();
+    ASSERT_NE(b1, nullptr);
+    b1->key_ = "k";
+    b1->size = 1024;
+    EXPECT_FALSE(cache.PutHotKey(b1, stale));
+    EXPECT_FALSE(cache.HasHotKey("k"));
+    EXPECT_EQ(cache.GetCacheSize(), 1)
+        << "Stale fill must be returned to the pool, not published";
+
+    // Valid: token still current -> published.
+    HotCachePutToken fresh = cache.AcquirePutToken("k");
+    HotMemBlock* b2 = cache.GetFreeBlock();
+    ASSERT_NE(b2, nullptr);
+    b2->key_ = "k";
+    b2->size = 1024;
+    EXPECT_TRUE(cache.PutHotKey(b2, fresh));
+    EXPECT_TRUE(cache.HasHotKey("k"));
+}
+
 // Test LocalHotCacheHandler basic functionality
 TEST_F(LocalHotCacheTest, LocalHotCacheHandlerBasic) {
     const size_t cache_size = 32 * 1024 * 1024;  // 32MB = 2 blocks
@@ -579,8 +608,7 @@ TEST_F(LocalHotCacheTest, ConcurrentAccess) {
 
     // Each thread puts and gets keys
     for (int t = 0; t < num_threads; ++t) {
-        threads.emplace_back([&cache, t, keys_per_thread, &successful_puts,
-                              &successful_gets]() {
+        threads.emplace_back([&cache, t, &successful_puts, &successful_gets]() {
             for (int i = 0; i < keys_per_thread; ++i) {
                 std::string key =
                     "thread_" + std::to_string(t) + "_key_" + std::to_string(i);

mooncake-store/tests/dummy_client_get_buffer_test.cpp

@@ -9,6 +9,7 @@
 #include <chrono>
 #include <cstdlib>
 #include <cstring>
+#include <iostream>
 #include <memory>
 #include <string>
 #include <thread>
@@ -184,6 +185,16 @@ class DummyClientGetBufferTest : public ::testing::Test {
         std::this_thread::sleep_for(std::chrono::milliseconds(2000));
     }
 
+    // Trigger hot-cache admission via DummyClient fallback reads (CMS threshold
+    // default 2) but do NOT wait for async PutHotKey to publish.
+    void TriggerAdmissionWithoutWaiting(const std::string &key) {
+        for (int i = 0; i < 2; ++i) {
+            auto buf = dummy_client_->get_buffer(key);
+            ASSERT_NE(buf, nullptr)
+                << "get_buffer failed while triggering admission";
+        }
+    }
+
     mooncake::testing::InProcMaster master_;
     std::shared_ptr<RealClient> real_client_;
     std::shared_ptr<DummyClient> dummy_client_;
@@ -194,6 +205,156 @@ class DummyClientGetBufferTest : public ::testing::Test {
     std::optional<std::string> saved_hot_block_env_;
     std::optional<std::string> saved_hot_shm_env_;
 };
+// ---- Regression: a stable hot-cache entry must be invalidated by remove ----
+TEST_F(DummyClientGetBufferTest,
+       StableHotCacheEntryShouldBeInvalidatedByRemove) {
+    ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack";
+
+    const std::string key = "stable_hot_cache_remove_invalidation";
+
+    const char fill = static_cast<char>(7);
+    const std::string data(kPayloadSize, fill);
+
+    PutData(key, data);
+    WarmHotCache(key);
+
+    auto hot_buf_1 = dummy_client_->get_buffer(key);
+    ASSERT_NE(hot_buf_1, nullptr)
+        << "dummy get_buffer should succeed after WarmHotCache";
+    ASSERT_EQ(hot_buf_1->size(), data.size())
+        << "hot cache buffer size mismatch before remove";
+    ASSERT_TRUE(dummy_client_->is_hot_cache_ptr(hot_buf_1->ptr()))
+        << "buffer should be in hot cache shm before remove";
+
+    {
+        std::string got(static_cast<char *>(hot_buf_1->ptr()),
+                        hot_buf_1->size());
+        ASSERT_EQ(got.size(), data.size());
+        ASSERT_TRUE(std::equal(got.begin(), got.end(), data.begin()))
+            << "hot cache data mismatch before remove";
+    }
+
+    auto hot_buf_2 = dummy_client_->get_buffer(key);
+    ASSERT_NE(hot_buf_2, nullptr)
+        << "second dummy get_buffer should still succeed before remove";
+    ASSERT_TRUE(dummy_client_->is_hot_cache_ptr(hot_buf_2->ptr()))
+        << "second buffer should still be in hot cache shm before remove";
+    ASSERT_EQ(hot_buf_2->size(), data.size())
+        << "second hot cache buffer size mismatch before remove";
+
+    {
+        std::string got(static_cast<char *>(hot_buf_2->ptr()),
+                        hot_buf_2->size());
+        ASSERT_EQ(got.size(), data.size());
+        ASSERT_TRUE(std::equal(got.begin(), got.end(), data.begin()))
+            << "second hot cache data mismatch before remove";
+    }
+
+    int remove_rc = real_client_->remove(key, true);
+    ASSERT_EQ(remove_rc, 0) << "remove failed";
+
+    auto after_remove_buf = dummy_client_->get_buffer(key);
+    if (after_remove_buf != nullptr) {
+        const bool after_remove_is_hot =
+            dummy_client_->is_hot_cache_ptr(after_remove_buf->ptr());
+        std::cerr << "[StableHotCacheEntryShouldBeInvalidatedByRemove] "
+                  << "after remove get_buffer returned non-null; "
+                  << "size=" << after_remove_buf->size()
+                  << ", is_hot_cache_ptr=" << after_remove_is_hot << std::endl;
+        FAIL() << "remove(key, true) succeeded, but dummy get_buffer(key) "
+               << "still returned a non-null buffer";
+    }
+}
+
+// ---- Diagnosis: removed key must not be resurrected by async hot-cache fill
+// ----
+//
+// Compare against pre-fix behavior:
+//   mooncake-store/tests/scripts/compare_hot_cache_fix.sh all
+TEST_F(DummyClientGetBufferTest,
+       RemoveShouldNotBeResurrectedByAsyncHotCacheFill) {
+    ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack";
+
+    const std::string key = "async_fill_resurrect_removed_key";
+    constexpr int kRounds = 8;
+
+    for (int round = 0; round < kRounds; ++round) {
+        SCOPED_TRACE("round=" + std::to_string(round));
+
+        const char fill = static_cast<char>(1 + (round % 250));
+        const std::string data(kPayloadSize, fill);
+
+        PutData(key, data);
+
+        // Two fallback reads reach CMS admission threshold (default 2) and
+        // submit async fill, but we do not wait for worker PutHotKey.
+        TriggerAdmissionWithoutWaiting(key);
+
+        int remove_rc = real_client_->remove(key, true);
+        ASSERT_EQ(remove_rc, 0) << "remove failed at round=" << round;
+
+        auto immediately_after_remove = dummy_client_->get_buffer(key);
+        const bool immediate_non_null = (immediately_after_remove != nullptr);
+        bool immediate_is_hot = false;
+        size_t immediate_size = 0;
+        if (immediate_non_null) {
+            immediate_is_hot = dummy_client_->is_hot_cache_ptr(
+                immediately_after_remove->ptr());
+            immediate_size = immediately_after_remove->size();
+            std::cerr << "[round=" << round
+                      << "] immediately after remove: NON_NULL"
+                      << ", size=" << immediate_size
+                      << ", is_hot_cache_ptr=" << immediate_is_hot << std::endl;
+        } else {
+            std::cerr << "[round=" << round
+                      << "] immediately after remove: nullptr" << std::endl;
+        }
+
+        std::this_thread::sleep_for(std::chrono::milliseconds(2000));
+
+        auto delayed_after_remove = dummy_client_->get_buffer(key);
+        const bool delayed_non_null = (delayed_after_remove != nullptr);
+        bool delayed_is_hot = false;
+        size_t delayed_size = 0;
+        if (delayed_non_null) {
+            delayed_is_hot =
+                dummy_client_->is_hot_cache_ptr(delayed_after_remove->ptr());
+            delayed_size = delayed_after_remove->size();
+            std::cerr << "[round=" << round
+                      << "] delayed after remove: NON_NULL"
+                      << ", size=" << delayed_size
+                      << ", is_hot_cache_ptr=" << delayed_is_hot << std::endl;
+        } else {
+            std::cerr << "[round=" << round << "] delayed after remove: nullptr"
+                      << std::endl;
+        }
+
+        if (immediate_non_null || delayed_non_null) {
+            if (!immediate_non_null && delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was nullptr immediately after remove, "
+                       << "but became non-null after waiting. "
+                       << "This strongly suggests async hot-cache fill "
+                       << "resurrected a removed key. "
+                       << "delayed_is_hot_cache_ptr=" << delayed_is_hot
+                       << ", delayed_size=" << delayed_size;
+            }
+            if (immediate_non_null && delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was non-null immediately after remove "
+                       << "and remained non-null after waiting. "
+                       << "immediate_is_hot_cache_ptr=" << immediate_is_hot
+                       << ", delayed_is_hot_cache_ptr=" << delayed_is_hot;
+            }
+            if (immediate_non_null && !delayed_non_null) {
+                FAIL() << "FAILED at round=" << round
+                       << ": key was non-null immediately after remove, "
+                       << "but became nullptr after waiting. "
+                       << "immediate_is_hot_cache_ptr=" << immediate_is_hot;
+            }
+        }
+    }
+}
 
 // ---- Test: get_buffer via allocator fallback (no hot cache hit) ----
 TEST_F(DummyClientGetBufferTest, GetBuffer_AllocatorFallback) {
@@ -409,4 +570,4 @@ int main(int argc, char **argv) {
     ::testing::InitGoogleTest(&argc, argv);
     gflags::ParseCommandLineFlags(&argc, &argv, false);
     return RUN_ALL_TESTS();
-}
+}