From 9cc3dbe2eecf51ca6af02ef0c3534250f26abac9 Mon Sep 17 00:00:00 2001 From: eric_song Date: Thu, 16 Jul 2026 23:13:20 +0800 Subject: [PATCH 1/2] fix(store): roll back failed offset evictions --- mooncake-store/include/storage_backend.h | 19 +++- mooncake-store/src/storage_backend.cpp | 85 ++++++++++++---- mooncake-store/tests/storage_backend_test.cpp | 99 +++++++++++++++++-- 3 files changed, 174 insertions(+), 29 deletions(-) diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 12ff19e05f..1ccef58e64 100644 --- a/mooncake-store/include/storage_backend.h +++ b/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> 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& batch_keys, - std::vector& 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 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. diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index ab54fc69f2..beba063523 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -3631,7 +3631,7 @@ tl::expected OffsetAllocatorStorageBackend::Init() { void OffsetAllocatorStorageBackend::EvictToMakeRoom( int64_t required_bytes, size_t min_victims, const std::unordered_set& batch_keys, - std::vector& out_evicted) { + PendingEviction& out_pending) { if (cfg_.eviction_policy == OffsetEvictionPolicy::NONE) return; MutexLocker ev(&eviction_mutex_); @@ -3687,19 +3687,68 @@ 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); + + 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 +OffsetAllocatorStorageBackend::NotifyAndCommitPreparedEviction( + const EvictionHandler& eviction_handler, PendingEviction& pending) { + if (pending.objects.empty()) return {}; + + std::vector 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 OffsetAllocatorStorageBackend::BatchOffload( const std::unordered_map>& batch_object, std::function& keys, @@ -3764,10 +3813,9 @@ tl::expected 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 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 +3865,7 @@ tl::expected 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(record_size), min_v, - batch_keys, evicted_keys); + batch_keys, pending_eviction); } } @@ -3826,12 +3874,12 @@ tl::expected 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 +3894,19 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( while (!allocation.has_value() && fallback_total_evicted < kMaxFallbackEvicted) { - size_t before = evicted_keys.size(); EvictToMakeRoom(static_cast(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 +4020,12 @@ tl::expected 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()) { diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 83123e0302..b87272358a 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/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::vector&) { + return ErrorCode::OK; + }; + auto successful_eviction_handler = + [](const std::vector&) -> tl::expected { + return {}; + }; + + std::string data(1000, 'x'); + std::vector> 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 failed_evictions; + auto failing_eviction_handler = + [&failed_evictions](const std::vector& keys) + -> tl::expected { + 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 output(data.size()); + std::unordered_map 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 retry_evictions; + auto retry_eviction_handler = + [&retry_evictions](const std::vector& keys) + -> tl::expected { + 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 From a31db231d5faca85f1b55fc8ec2a0a739abc3ecc Mon Sep 17 00:00:00 2001 From: eric_song Date: Fri, 17 Jul 2026 09:46:47 +0800 Subject: [PATCH 2/2] fix(store): guard eviction rollback assumptions --- mooncake-store/src/storage_backend.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index beba063523..3b666fade2 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -3710,6 +3710,10 @@ void OffsetAllocatorStorageBackend::RestorePreparedEviction( 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; @@ -3727,6 +3731,10 @@ tl::expected OffsetAllocatorStorageBackend::NotifyAndCommitPreparedEviction( const EvictionHandler& eviction_handler, PendingEviction& pending) { if (pending.objects.empty()) return {}; + if (!eviction_handler) { + pending.objects.clear(); + return {}; + } std::vector evicted_keys; evicted_keys.reserve(pending.objects.size());