[Bugfix] Roll back failed OffsetAllocator evictions

mooncake-store/include/storage_backend.h

@@ -1271,6 +1271,12 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface {
               fifo_seq(seq) {}
     };
 
+    // Keeps evicted metadata and allocation handles alive until the master
+    // accepts the replica-removal notification.
+    struct PendingEviction {
+        std::vector<std::pair<std::string, ObjectEntry>> objects;
+    };
+
     // Returns full path to data file: {storage_path_}/kv_cache.data
     std::string GetDataFilePath() const;
 
@@ -1360,10 +1366,19 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface {
     int64_t low_watermark_keys_ = 0;
 
     // Evict keys from the FIFO index until both byte and key-count watermarks
-    // are satisfied (or until the eviction cap is reached).
+    // are satisfied (or until the eviction cap is reached). Allocations remain
+    // pinned in out_pending until notification succeeds.
     void EvictToMakeRoom(int64_t required_bytes, size_t min_victims,
                          const std::unordered_set<std::string>& batch_keys,
-                         std::vector<std::string>& out_evicted);
+                         PendingEviction& out_pending);
+
+    // Restore prepared victims when the master rejects their removal.
+    void RestorePreparedEviction(PendingEviction&& pending);
+
+    // Notify the master, then release prepared allocations on success or
+    // restore their metadata on failure.
+    tl::expected<void, ErrorCode> NotifyAndCommitPreparedEviction(
+        const EvictionHandler& eviction_handler, PendingEviction& pending);
 
     // Test-only: Predicate to determine which keys should fail in BatchOffload.
     // Used for deterministic testing of partial success behavior.

mooncake-store/src/storage_backend.cpp

@@ -3631,7 +3631,7 @@ tl::expected<void, ErrorCode> OffsetAllocatorStorageBackend::Init() {
 void OffsetAllocatorStorageBackend::EvictToMakeRoom(
     int64_t required_bytes, size_t min_victims,
     const std::unordered_set<std::string>& batch_keys,
-    std::vector<std::string>& out_evicted) {
+    PendingEviction& out_pending) {
     if (cfg_.eviction_policy == OffsetEvictionPolicy::NONE) return;
 
     MutexLocker ev(&eviction_mutex_);
@@ -3687,19 +3687,76 @@ void OffsetAllocatorStorageBackend::EvictToMakeRoom(
             DCHECK_GE(total_size_.load(std::memory_order_relaxed),
                       it->second.total_size);
             DCHECK_GE(total_keys_.load(std::memory_order_relaxed), 1);
+            out_pending.objects.emplace_back(vkey, it->second);
             total_size_.fetch_sub(it->second.total_size,
                                   std::memory_order_relaxed);
             total_keys_.fetch_sub(1, std::memory_order_relaxed);
             shard.map.erase(it);
         }
         fifo_index_.erase(oldest);
-        out_evicted.push_back(std::move(vkey));
         ++n;
     }
 }
 
 //-----------------------------------------------------------------------------
 
+void OffsetAllocatorStorageBackend::RestorePreparedEviction(
+    PendingEviction&& pending) {
+    MutexLocker ev(&eviction_mutex_);
+
+    for (auto& [key, entry] : pending.objects) {
+        const uint32_t restored_size = entry.total_size;
+        const uint64_t restored_seq = entry.fifo_seq;
+        auto& shard = shards_[ShardForKey(key)];
+        SharedMutexLocker lk(&shard.mutex);
+
+        // BatchOffload's single-writer precondition guarantees that no writer
+        // can recreate this key or consume its FIFO sequence while the
+        // eviction notification is in flight. Readers do not mutate either
+        // index, so both entries must still be absent here.
+        const bool map_inserted =
+            shard.map.emplace(key, std::move(entry)).second;
+        CHECK(map_inserted) << "Failed to restore evicted key: " << key;
+        const bool fifo_inserted =
+            fifo_index_.emplace(restored_seq, key).second;
+        CHECK(fifo_inserted)
+            << "Failed to restore FIFO entry for evicted key: " << key;
+        total_size_.fetch_add(restored_size, std::memory_order_relaxed);
+        total_keys_.fetch_add(1, std::memory_order_relaxed);
+    }
+    pending.objects.clear();
+}
+
+tl::expected<void, ErrorCode>
+OffsetAllocatorStorageBackend::NotifyAndCommitPreparedEviction(
+    const EvictionHandler& eviction_handler, PendingEviction& pending) {
+    if (pending.objects.empty()) return {};
+    if (!eviction_handler) {
+        pending.objects.clear();
+        return {};
+    }
+
+    std::vector<std::string> evicted_keys;
+    evicted_keys.reserve(pending.objects.size());
+    for (const auto& [key, _] : pending.objects) {
+        evicted_keys.push_back(key);
+    }
+
+    auto notify_result = eviction_handler(evicted_keys);
+    if (!notify_result) {
+        const ErrorCode error = notify_result.error();
+        RestorePreparedEviction(std::move(pending));
+        return tl::make_unexpected(error);
+    }
+
+    // Dropping the final metadata references commits the eviction and makes
+    // unpinned extents available to the allocator.
+    pending.objects.clear();
+    return {};
+}
+
+//-----------------------------------------------------------------------------
+
 tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
     const std::unordered_map<std::string, std::vector<Slice>>& batch_object,
     std::function<ErrorCode(const std::vector<std::string>& keys,
@@ -3764,10 +3821,9 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
     keys.reserve(batch_object.size());
     metadatas.reserve(batch_object.size());
 
-    // Accumulated evicted keys across the per-key loop.
-    // Flushed to eviction_handler before each allocate() that may reuse
-    // freed space, and again after the loop for any leftover victims.
-    std::vector<std::string> evicted_keys;
+    // Prepared victims are held here until the master accepts their removal.
+    // Their allocation handles prevent reuse before notification succeeds.
+    PendingEviction pending_eviction;
 
     for (const auto& [key, slices] : batch_object) {
         if (slices.empty()) continue;
@@ -3817,7 +3873,7 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
                 // fallback_evict_batch victims even if bytes are low.
                 size_t min_v = over_keys ? cfg_.fallback_evict_batch : 0;
                 EvictToMakeRoom(static_cast<int64_t>(record_size), min_v,
-                                batch_keys, evicted_keys);
+                                batch_keys, pending_eviction);
             }
         }
 
@@ -3826,12 +3882,12 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
         // so the allocator cannot re-issue a still-read offset.  Layer-2
         // (master metadata): the master must be told the key's local-disk
         // replica is gone before we reuse its space for a new key.
-        if (eviction_on && eviction_handler && !evicted_keys.empty()) {
-            auto notify_result = eviction_handler(evicted_keys);
+        if (eviction_on && !pending_eviction.objects.empty()) {
+            auto notify_result = NotifyAndCommitPreparedEviction(
+                eviction_handler, pending_eviction);
             if (!notify_result) {
                 return tl::make_unexpected(notify_result.error());
             }
-            evicted_keys.clear();
         }
 
         // ---- (C) Allocate ----
@@ -3846,20 +3902,19 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
 
             while (!allocation.has_value() &&
                    fallback_total_evicted < kMaxFallbackEvicted) {
-                size_t before = evicted_keys.size();
                 EvictToMakeRoom(static_cast<int64_t>(record_size),
                                 cfg_.fallback_evict_batch, batch_keys,
-                                evicted_keys);
-                size_t evicted_this_turn = evicted_keys.size() - before;
+                                pending_eviction);
+                size_t evicted_this_turn = pending_eviction.objects.size();
                 fallback_total_evicted += evicted_this_turn;
 
                 // Notify master of fallback victims before retrying.
-                if (eviction_handler && !evicted_keys.empty()) {
-                    auto notify_result = eviction_handler(evicted_keys);
+                if (!pending_eviction.objects.empty()) {
+                    auto notify_result = NotifyAndCommitPreparedEviction(
+                        eviction_handler, pending_eviction);
                     if (!notify_result) {
                         return tl::make_unexpected(notify_result.error());
                     }
-                    evicted_keys.clear();
                 }
 
                 uint64_t now_largest =
@@ -3973,12 +4028,12 @@ tl::expected<int64_t, ErrorCode> OffsetAllocatorStorageBackend::BatchOffload(
 
     // ---- Post-loop flush: notify master of any evicted keys that
     // were accumulated by the last (possibly allocate-failing) key.
-    if (eviction_on && eviction_handler && !evicted_keys.empty()) {
-        auto notify_result = eviction_handler(evicted_keys);
+    if (eviction_on && !pending_eviction.objects.empty()) {
+        auto notify_result =
+            NotifyAndCommitPreparedEviction(eviction_handler, pending_eviction);
         if (!notify_result) {
             return tl::make_unexpected(notify_result.error());
         }
-        evicted_keys.clear();
     }
 
     if (complete_handler != nullptr && !keys.empty()) {

mooncake-store/tests/storage_backend_test.cpp

@@ -3785,6 +3785,90 @@ TEST_F(StorageBackendTest,
 
 //-----------------------------------------------------------------------------
 
+TEST_F(StorageBackendTest,
+       OffsetAllocatorStorageBackend_Eviction_NotificationFailureRollback) {
+    FileStorageConfig config;
+    config.storage_filepath = data_path;
+    config.storage_backend_type = StorageBackendType::kOffsetAllocator;
+    config.total_size_limit = 8 * 1024;
+    config.total_keys_limit = 100;
+
+    OffsetAllocatorBackendConfig evict_cfg;
+    evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO;
+    evict_cfg.high_watermark_bytes = 3500;
+    evict_cfg.low_watermark_bytes = 2100;
+
+    OffsetAllocatorStorageBackend storage_backend(config, evict_cfg);
+    ASSERT_TRUE(storage_backend.Init());
+
+    auto complete_handler = [](const std::vector<std::string>&,
+                               std::vector<StorageObjectMetadata>&) {
+        return ErrorCode::OK;
+    };
+    auto successful_eviction_handler =
+        [](const std::vector<std::string>&) -> tl::expected<void, ErrorCode> {
+        return {};
+    };
+
+    std::string data(1000, 'x');
+    std::vector<std::unique_ptr<char[]>> buffers;
+    for (const auto& key : {"key_a", "key_b", "key_c"}) {
+        auto batch = MakeSingleKeyBatch(key, data, buffers);
+        auto result = storage_backend.BatchOffload(batch, complete_handler,
+                                                   successful_eviction_handler);
+        ASSERT_TRUE(result.has_value()) << "key=" << key;
+    }
+
+    std::vector<std::string> failed_evictions;
+    auto failing_eviction_handler =
+        [&failed_evictions](const std::vector<std::string>& keys)
+        -> tl::expected<void, ErrorCode> {
+        failed_evictions = keys;
+        return tl::make_unexpected(ErrorCode::INTERNAL_ERROR);
+    };
+
+    auto batch = MakeSingleKeyBatch("key_d", data, buffers);
+    auto failed_result = storage_backend.BatchOffload(batch, complete_handler,
+                                                      failing_eviction_handler);
+    ASSERT_FALSE(failed_result.has_value());
+    EXPECT_EQ(failed_result.error(), ErrorCode::INTERNAL_ERROR);
+    ASSERT_FALSE(failed_evictions.empty());
+    EXPECT_EQ(failed_evictions.front(), "key_a");
+
+    for (const auto& key : failed_evictions) {
+        auto exists = storage_backend.IsExist(key);
+        ASSERT_TRUE(exists.has_value());
+        EXPECT_TRUE(exists.value()) << "key=" << key;
+
+        std::vector<char> output(data.size());
+        std::unordered_map<std::string, Slice> load_batch;
+        load_batch.emplace(key, Slice{output.data(), output.size()});
+        auto load_result = storage_backend.BatchLoad(load_batch);
+        ASSERT_TRUE(load_result.has_value()) << "key=" << key;
+        EXPECT_EQ(std::string(output.begin(), output.end()), data);
+    }
+    EXPECT_FALSE(storage_backend.IsExist("key_d").value_or(true));
+
+    std::vector<std::string> retry_evictions;
+    auto retry_eviction_handler =
+        [&retry_evictions](const std::vector<std::string>& keys)
+        -> tl::expected<void, ErrorCode> {
+        retry_evictions = keys;
+        return {};
+    };
+    auto retry_result = storage_backend.BatchOffload(batch, complete_handler,
+                                                     retry_eviction_handler);
+    ASSERT_TRUE(retry_result.has_value());
+    EXPECT_EQ(retry_evictions, failed_evictions);
+    for (const auto& key : retry_evictions) {
+        EXPECT_FALSE(storage_backend.IsExist(key).value_or(true))
+            << "key=" << key;
+    }
+    EXPECT_TRUE(storage_backend.IsExist("key_d").value_or(false));
+}
+
+//-----------------------------------------------------------------------------
+
 TEST_F(StorageBackendTest,
        OffsetAllocatorStorageBackend_Eviction_PostLoopFlush) {
     // The last key in a batch may trigger eviction but fail allocate.
@@ -3893,14 +3977,13 @@ TEST_F(StorageBackendTest,
     //   1. BatchLoad copies entry.allocation (shared_ptr) into its
     //      ReadPlan, incrementing the refcount.  (storage_backend.cpp
     //      ~line 3375: "entry.allocation" copy in ReadPlan)
-    //   2. EvictToMakeRoom erases the key from shard.map, decrementing
-    //      the map's shared_ptr.  If no reader holds a copy, the
-    //      RefCountedAllocationHandle destructor calls freeAllocation
-    //      and the extent returns to the allocator.
-    //   3. While a reader holds its shared_ptr copy (refcount >= 1),
-    //      freeAllocation does NOT fire → the extent is still marked
-    //      "used" in the allocator → allocate() cannot re-issue that
-    //      offset.  The reader always sees the original bytes.
+    //   2. EvictToMakeRoom retains the allocation shared_ptr in pending
+    //      eviction state before erasing the key from shard.map. The pending
+    //      reference is released only after the master accepts the removal.
+    //   3. While either pending eviction or a reader holds a shared_ptr,
+    //      freeAllocation does NOT fire → the extent is still marked "used"
+    //      in the allocator → allocate() cannot re-issue that offset. The
+    //      reader always sees the original bytes.
     //
     // This test interleaves reads and eviction-triggering writes;
     // any data corruption means the allocator re-issued a still-read