diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 8a33264acc..b1062958e8 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -200,6 +200,44 @@ struct BucketBackendConfig { static BucketBackendConfig FromEnvironment(); }; +enum class OffsetEvictionPolicy { + NONE, // No eviction + FIFO, // Evict oldest key first (by insertion order) + LRU, // Approximate LRU via cross-shard sampling (phase 2) +}; + +struct OffsetAllocatorBackendConfig { + OffsetEvictionPolicy eviction_policy = OffsetEvictionPolicy::NONE; + + // Watermark thresholds: eviction triggers when total_size_ exceeds high, + // drives down to low. 0 = auto-resolved in Init() from ratios. + int64_t high_watermark_bytes = 0; + int64_t low_watermark_bytes = 0; + double high_ratio = 0.90; + double low_ratio = 0.80; + + // Key-count watermarks (symmetric with byte watermarks). + // high triggers eviction, drives down to low. + int64_t high_watermark_keys = 0; + int64_t low_watermark_keys = 0; + double keys_high_ratio = 0.95; + double keys_low_ratio = 0.90; + + // Eviction caps + size_t max_evict_per_offload = 4096; + size_t fallback_evict_batch = 16; + + // Allocator node capacity override. + // 0 = auto-derived from capacity_ / kMinObjectSize (capped at RAM budget). + // Must be <= UINT32_MAX (OffsetAllocator::create takes uint32 + // max_capacity). + int64_t max_capacity_nodes = 0; + + bool Validate() const; + + static OffsetAllocatorBackendConfig FromEnvironment(); +}; + struct FileStorageConfig { // type of the storage backend StorageBackendType storage_backend_type = StorageBackendType::kBucket; @@ -992,7 +1030,8 @@ class BucketStorageBackend : public StorageBackendInterface { class OffsetAllocatorStorageBackend : public StorageBackendInterface { public: OffsetAllocatorStorageBackend( - const FileStorageConfig& file_storage_config_); + const FileStorageConfig& file_storage_config_, + const OffsetAllocatorBackendConfig& offset_backend_config = {}); /** * @brief Initializes the offset allocator storage backend. @@ -1062,6 +1101,15 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { test_failure_predicate_ = std::move(predicate); } + // Returns the number of keys skipped after fallback eviction + // could not make enough room (fragmentation, extents pinned by + // in-flight reads, or allocator node exhaustion). Monotonically + // increasing; useful for distinguishing "watermark working" from + // "thrashing but unable to free space". + int64_t GetEvictionSkips() const { + return eviction_skips_.load(std::memory_order_relaxed); + } + private: // On-disk record header: [u32 key_len][u32 value_len] (8 bytes total) struct RecordHeader { @@ -1132,12 +1180,19 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // Refcounted handle keeps physical extent alive during reads AllocationPtr allocation; + + // Monotonic insertion sequence number. Points back to the slot in + // fifo_index_ (seq -> key). Used during eviction to detect stale + // index entries (lazy-repair) and to remove old slots on overwrite. + uint64_t fifo_seq = 0; + ObjectEntry(uint64_t off, uint32_t total, uint32_t val, - AllocationPtr alloc_ptr) + AllocationPtr alloc_ptr, uint64_t seq = 0) : offset(off), total_size(total), value_size(val), - allocation(std::move(alloc_ptr)) {} + allocation(std::move(alloc_ptr)), + fifo_seq(seq) {} }; // Returns full path to data file: {storage_path_}/kv_cache.data @@ -1201,6 +1256,39 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // counting) std::atomic total_keys_{0}; + // ===== Eviction-related members ===== + OffsetAllocatorBackendConfig cfg_; + + // Counter for keys skipped due to fallback eviction exhaustion. + // See GetEvictionSkips() for the public accessor. + std::atomic eviction_skips_{0}; + + // Mutex protecting fifo_index_ and insert_seq_. Must be acquired BEFORE + // any shard mutex (shards_[i].mutex) when both are held. + mutable Mutex eviction_mutex_; + + // Global FIFO index: insertion sequence number -> key. + // begin() = oldest key, the default eviction victim. + // Entries allowed to be stale; lazy-repair at eviction time. + std::map fifo_index_; + + // Monotonic sequence number source for fifo_index_. + std::atomic insert_seq_{0}; + + // Resolved watermark thresholds (bytes), computed in Init(). + int64_t high_watermark_bytes_ = 0; + int64_t low_watermark_bytes_ = 0; + + // Resolved watermark thresholds (key count), computed in Init(). + int64_t high_watermark_keys_ = 0; + 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). + void EvictToMakeRoom(int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + std::vector& out_evicted); + // Test-only: Predicate to determine which keys should fail in BatchOffload. // Used for deterministic testing of partial success behavior. std::function test_failure_predicate_; diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 3612624a3f..2a1d7cc030 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -86,6 +87,93 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { return config; } +bool OffsetAllocatorBackendConfig::Validate() const { + if (high_ratio <= 0.0 || high_ratio > 1.0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: high_ratio must be in (0,1]"; + return false; + } + if (low_ratio <= 0.0 || low_ratio >= high_ratio) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: low_ratio must be in (0, " + "high_ratio)"; + return false; + } + if (keys_high_ratio <= 0.0 || keys_high_ratio > 1.0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: keys_high_ratio must be in (0,1]"; + return false; + } + if (keys_low_ratio <= 0.0 || keys_low_ratio >= keys_high_ratio) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: keys_low_ratio must be in " + "(0, keys_high_ratio)"; + return false; + } + if (max_evict_per_offload == 0) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: max_evict_per_offload " + "must be > 0"; + return false; + } + if (fallback_evict_batch == 0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: fallback_evict_batch must be > 0"; + return false; + } + if (max_capacity_nodes < 0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: max_capacity_nodes must be >= 0"; + return false; + } + return true; +} + +static std::optional GetEnvDouble(const char* name) { + const char* env = std::getenv(name); + if (!env || env[0] == '\0') return std::nullopt; + try { + return std::stod(env); + } catch (...) { + return std::nullopt; + } +} + +OffsetAllocatorBackendConfig OffsetAllocatorBackendConfig::FromEnvironment() { + OffsetAllocatorBackendConfig cfg; + + const char* pol = std::getenv("MOONCAKE_OFFSET_EVICTION_POLICY"); + if (pol) { + std::string s(pol); + if (s == "fifo" || s == "FIFO" || s == "Fifo") { + cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + } + // NONE is default; LRU reserved for phase 2 + } + + if (auto v = GetEnvDouble("MOONCAKE_OFFSET_HIGH_RATIO")) + cfg.high_ratio = *v; + if (auto v = GetEnvDouble("MOONCAKE_OFFSET_LOW_RATIO")) cfg.low_ratio = *v; + // Both byte and key watermarks derive from the same ratio pair. + cfg.keys_high_ratio = cfg.high_ratio; + cfg.keys_low_ratio = cfg.low_ratio; + + cfg.max_capacity_nodes = GetEnvOr( + "MOONCAKE_OFFSET_MAX_CAPACITY_NODES", cfg.max_capacity_nodes); + + // Read eviction cap as int64_t to guard against negative env values + // which would wrap to SIZE_MAX with GetEnvOr. + auto max_evict_raw = + GetEnvOr("MOONCAKE_OFFSET_MAX_EVICT_PER_OFFLOAD", + static_cast(cfg.max_evict_per_offload)); + if (max_evict_raw > 0) { + cfg.max_evict_per_offload = static_cast(max_evict_raw); + } else if (max_evict_raw <= 0) { + LOG(WARNING) << "MOONCAKE_OFFSET_MAX_EVICT_PER_OFFLOAD=" + << max_evict_raw << " is non-positive; using default " + << cfg.max_evict_per_offload; + } + + return cfg; +} + StorageBackendInterface::StorageBackendInterface( const FileStorageConfig& config) : file_storage_config_(config) {} @@ -2716,9 +2804,11 @@ BucketStorageBackend::GetFileInstance() const { // ============================================================================ OffsetAllocatorStorageBackend::OffsetAllocatorStorageBackend( - const FileStorageConfig& file_storage_config_) + const FileStorageConfig& file_storage_config_, + const OffsetAllocatorBackendConfig& offset_backend_config) : StorageBackendInterface(file_storage_config_), - storage_path_(file_storage_config_.storage_filepath) { + storage_path_(file_storage_config_.storage_filepath), + cfg_(offset_backend_config) { capacity_ = file_storage_config_.total_size_limit; } @@ -2812,16 +2902,128 @@ tl::expected OffsetAllocatorStorageBackend::Init() { fd_guard.release()); } - // Create allocator with base=0, size=capacity - allocator_ = offset_allocator::OffsetAllocator::create(0, capacity_); + // Resolve watermark thresholds from config + high_watermark_bytes_ = + cfg_.high_watermark_bytes > 0 + ? cfg_.high_watermark_bytes + : static_cast(capacity_ * cfg_.high_ratio); + low_watermark_bytes_ = + cfg_.low_watermark_bytes > 0 + ? cfg_.low_watermark_bytes + : static_cast(capacity_ * cfg_.low_ratio); + high_watermark_keys_ = + cfg_.high_watermark_keys > 0 + ? cfg_.high_watermark_keys + : static_cast(file_storage_config_.total_keys_limit * + cfg_.keys_high_ratio); + low_watermark_keys_ = + cfg_.low_watermark_keys > 0 + ? cfg_.low_watermark_keys + : static_cast(file_storage_config_.total_keys_limit * + cfg_.keys_low_ratio); + + // Auto-nudge ratio-derived low watermarks when integer truncation + // collapses them to the same value as high (e.g. limit=5, ratio + // 0.95->4, ratio 0.90->4 => low==high==4). Only applies to + // auto-derived values; explicit config values are validated strictly. + if (cfg_.low_watermark_bytes == 0 && high_watermark_bytes_ > 0 && + low_watermark_bytes_ >= high_watermark_bytes_) { + low_watermark_bytes_ = + std::max(1, high_watermark_bytes_ - 1); + } + if (cfg_.low_watermark_keys == 0 && high_watermark_keys_ > 0 && + low_watermark_keys_ >= high_watermark_keys_) { + low_watermark_keys_ = + std::max(1, high_watermark_keys_ - 1); + } + + // Validate watermarks + if (low_watermark_bytes_ >= high_watermark_bytes_) { + LOG(ERROR) << "Invalid watermark: low_bytes=" + << low_watermark_bytes_ + << " >= high_bytes=" << high_watermark_bytes_; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (low_watermark_keys_ >= high_watermark_keys_) { + LOG(ERROR) << "Invalid watermark: low_keys=" << low_watermark_keys_ + << " >= high_keys=" << high_watermark_keys_; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + // Clamp watermarks to not exceed capacity / total_keys_limit + if (high_watermark_bytes_ > static_cast(capacity_)) { + LOG(WARNING) << "high_watermark_bytes clamped from " + << high_watermark_bytes_ + << " to capacity=" << capacity_; + high_watermark_bytes_ = static_cast(capacity_); + low_watermark_bytes_ = + std::min(low_watermark_bytes_, high_watermark_bytes_ - 1); + } + if (high_watermark_keys_ > file_storage_config_.total_keys_limit) { + LOG(WARNING) << "high_watermark_keys clamped from " + << high_watermark_keys_ << " to total_keys_limit=" + << file_storage_config_.total_keys_limit; + high_watermark_keys_ = file_storage_config_.total_keys_limit; + low_watermark_keys_ = + std::min(low_watermark_keys_, high_watermark_keys_ - 1); + } + + // Guard against zero low-watermark on very small capacity + if (low_watermark_bytes_ <= 0 && high_watermark_bytes_ > 0) { + low_watermark_bytes_ = + std::max(1, high_watermark_bytes_ / 2); + } + if (low_watermark_keys_ <= 0 && high_watermark_keys_ > 0) { + low_watermark_keys_ = + std::max(1, high_watermark_keys_ / 2); + } + + // Create allocator with tuned node capacity + constexpr int64_t kMinObjectSize = 256; + constexpr int64_t kMaxNodeRamBytes = + 512LL * 1024 * 1024; // 512MB node RAM budget + constexpr uint32_t kRamBasedMaxNodes = + static_cast(kMaxNodeRamBytes / 56); + constexpr uint32_t kAbsoluteMaxNodes = + std::min(kRamBasedMaxNodes, 32U << 20); + + uint32_t max_nodes = (1U << 20); // default 1M nodes + if (cfg_.max_capacity_nodes > 0) { + if (cfg_.max_capacity_nodes > kAbsoluteMaxNodes) { + LOG(WARNING) + << "max_capacity_nodes " << cfg_.max_capacity_nodes + << " exceeds RAM budget; clamped to " << kAbsoluteMaxNodes; + max_nodes = kAbsoluteMaxNodes; + } else { + max_nodes = static_cast(cfg_.max_capacity_nodes); + } + } else { + int64_t auto_nodes = std::max( + 1LL << 20, std::min(capacity_ / kMinObjectSize, + kAbsoluteMaxNodes)); + max_nodes = static_cast(auto_nodes); + } + uint32_t init_nodes = std::min(128U * 1024, max_nodes); + allocator_ = offset_allocator::OffsetAllocator::create( + 0, capacity_, init_nodes, max_nodes); if (!allocator_) { LOG(ERROR) << "Failed to create OffsetAllocator"; return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } + // Initialize eviction index + { + MutexLocker ev(&eviction_mutex_); + fifo_index_.clear(); + insert_seq_.store(0, std::memory_order_relaxed); + } + initialized_.store(true, std::memory_order_release); LOG(INFO) << "OffsetAllocatorStorageBackend initialized, capacity: " - << capacity_ << " bytes, data file: " << data_file_path_; + << capacity_ + << " bytes, high_watermark: " << high_watermark_bytes_ + << " bytes, " << high_watermark_keys_ << " keys" + << ", data file: " << data_file_path_; } catch (const std::exception& e) { LOG(ERROR) << "OffsetAllocatorStorageBackend initialize error: " << e.what(); @@ -2831,6 +3033,80 @@ tl::expected OffsetAllocatorStorageBackend::Init() { return {}; } +//----------------------------------------------------------------------------- +// EvictToMakeRoom +//----------------------------------------------------------------------------- + +void OffsetAllocatorStorageBackend::EvictToMakeRoom( + int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + std::vector& out_evicted) { + if (cfg_.eviction_policy == OffsetEvictionPolicy::NONE) return; + + MutexLocker ev(&eviction_mutex_); + size_t n = 0; + + while (n < cfg_.max_evict_per_offload) { + int64_t cur_size = total_size_.load(std::memory_order_relaxed); + int64_t cur_keys = total_keys_.load(std::memory_order_relaxed); + bool below_bytes = (cur_size + required_bytes <= low_watermark_bytes_); + bool below_keys = (cur_keys <= low_watermark_keys_); + // Stop when both byte and key-count watermarks are satisfied + // and we have met the minimum victim count. + if (below_bytes && below_keys && n >= min_victims) break; + + if (fifo_index_.empty()) break; + + auto oldest = fifo_index_.begin(); + uint64_t vseq = oldest->first; + std::string vkey = oldest->second; + + // Skip batch_keys prefix — keys being written in this batch. + // Do NOT erase their FIFO slots (they may fail allocate() and + // need the slot to remain in the index for future eviction). + // Worst-case comparison cost: O(|batch_keys_prefix|), bounded. + while (oldest != fifo_index_.end() && + batch_keys.count(oldest->second)) { + ++oldest; + } + if (oldest == fifo_index_.end()) break; // all are batch_keys + vkey = oldest->second; + vseq = oldest->first; + + size_t shard_idx = ShardForKey(vkey); + auto& shard = shards_[shard_idx]; + { + SharedMutexLocker lk(&shard.mutex); + auto it = shard.map.find(vkey); + if (it == shard.map.end() || it->second.fifo_seq != vseq) { + // Orphan slot in fifo_index_: the key is no longer in + // shard.map, or its fifo_seq was replaced by a newer + // overwrite. Overwrites are cleaned up in BatchOffload + // Step-4 (fifo_index_.erase(old_seq) under the lock), + // so today this branch is only reachable if a future + // per-key delete path neglects to also erase from + // fifo_index_. Keep the lazy-repair as a defense. + fifo_index_.erase(oldest); + ++n; // counted toward scan budget + continue; + } + // Defensive assertions against double-evict underflow. + // Precondition: single heartbeat_thread_ serialises offload; + // if concurrent offload is added, these must be re-evaluated. + DCHECK_GE(total_size_.load(std::memory_order_relaxed), + it->second.total_size); + DCHECK_GE(total_keys_.load(std::memory_order_relaxed), 1); + 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; + } +} + //----------------------------------------------------------------------------- tl::expected OffsetAllocatorStorageBackend::BatchOffload( @@ -2838,8 +3114,21 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( std::function& keys, std::vector& metadatas)> complete_handler, - std::function& /*evicted_keys*/)> - /*eviction_handler*/) { + std::function& evicted_keys)> + eviction_handler) { + // ================================================================ + // SINGLE-WRITER PRECONDITION + // + // BatchOffload, EvictToMakeRoom, and the watermark accounting on + // total_size_ / total_keys_ assume only ONE thread calls + // BatchOffload at a time (currently guaranteed by FileStorage's + // single heartbeat_thread_). The atomics make individual loads + // and stores atomic, but the read-modify-write sequences (check + // watermark → evict → update counters) are NOT atomic across + // threads. If concurrent offload is added, the DCHECK_GE guards + // in EvictToMakeRoom must also be re-evaluated. + // ================================================================ + if (!initialized_.load(std::memory_order_acquire)) { LOG(ERROR) << "Storage backend is not initialized. Call Init() before use."; @@ -2858,33 +3147,47 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); } + const bool eviction_on = + (cfg_.eviction_policy != OffsetEvictionPolicy::NONE) && + (eviction_handler != nullptr); + + // Warn if eviction policy is set but caller omitted the handler. + // IsEnableOffloading() will still return true in this mode, but + // eviction is effectively disabled (degraded to NONE behavior). + if (cfg_.eviction_policy != OffsetEvictionPolicy::NONE && + eviction_handler == nullptr) { + LOG_FIRST_N(WARNING, 1) + << "Eviction policy is " << static_cast(cfg_.eviction_policy) + << " but eviction_handler is null; eviction is disabled. " + "IsEnableOffloading() will still return true."; + } + + // Build the set of keys being offloaded in this batch so that + // EvictToMakeRoom does not evict them (they aren't committed yet). + std::unordered_set batch_keys; + if (eviction_on) { + for (const auto& [k, _] : batch_object) batch_keys.insert(k); + } + std::vector keys; std::vector metadatas; keys.reserve(batch_object.size()); metadatas.reserve(batch_object.size()); - // Process each object in the batch; continue on individual failures to - // support partial success + // 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; + for (const auto& [key, slices] : batch_object) { - if (slices.empty()) { - // Skip empty slices (empty values are allowed but not stored) - continue; - } + if (slices.empty()) continue; - // Test-only: Check if this key should fail (deterministic failure - // injection) if (test_failure_predicate_ && test_failure_predicate_(key)) { LOG(INFO) << "[TEST] Injecting failure for key: " << key << " (test failure predicate)"; - continue; // Simulate allocation/write failure + continue; } - // Calculate total value size. RecordHeader stores value_len as a - // uint32_t (RecordHeader::SIZE is 8 bytes), so accumulating directly - // into a uint32_t would silently overflow for an object larger than - // 4 GiB: the record would be under-allocated and then its full slices - // written past the allocation. Sum in 64 bits and reject oversized - // objects instead of corrupting the storage arena. uint64_t total_value_size = 0; for (const auto& slice : slices) { total_value_size += slice.size; @@ -2895,48 +3198,116 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( "in 4 GiB) for key: " << key << ", size: " << total_value_size << " - skipping this key"; - continue; // partial-success model: keep processing other keys + continue; } uint32_t value_size = static_cast(total_value_size); - // Prepare record header RecordHeader header{.key_len = static_cast(key.size()), .value_len = value_size}; - - // Use size_t for record_size to handle large objects (up to 4GB per - // RecordHeader) size_t record_size = RecordHeader::SIZE + header.key_len + header.value_len; - // Step 1: Allocate space (allocator is thread-safe, ensures unique - // offsets) No locks held during allocation + // Guard against record_size exceeding what the on-disk format + // can represent (ObjectEntry::total_size is uint32_t). + if (record_size > UINT32_MAX) { + LOG(ERROR) << "Record too large for key: " << key + << ", record_size=" << record_size; + continue; + } + + // ---- (A) Proactive eviction (watermark-driven) ---- + if (eviction_on) { + int64_t cur_size = total_size_.load(std::memory_order_relaxed); + int64_t cur_keys = total_keys_.load(std::memory_order_relaxed); + bool over_bytes = (cur_size + static_cast(record_size) > + high_watermark_bytes_); + bool over_keys = (cur_keys > high_watermark_keys_); + if (over_bytes || over_keys) { + // When triggered by key-count overflow, force at least + // 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); + } + } + + // ---- (B) Notify master of evicted keys BEFORE allocating ---- + // Layer-1 (byte safety): BatchLoad pins extents via shared_ptr, + // 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()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); + } + + // ---- (C) Allocate ---- auto allocation = allocator_->allocate(record_size); + + // ---- (D) Fallback eviction (nullopt retry loop) ---- + if (!allocation.has_value() && eviction_on) { + uint64_t prev_largest = + allocator_->get_metrics().largest_free_region_; + size_t fallback_total_evicted = 0; + const size_t kMaxFallbackEvicted = cfg_.max_evict_per_offload; + + 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; + fallback_total_evicted += evicted_this_turn; + + // Notify master of fallback victims before retrying. + if (eviction_handler && !evicted_keys.empty()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); + } + + uint64_t now_largest = + allocator_->get_metrics().largest_free_region_; + if (evicted_this_turn == 0) break; // no victims at all + // Stop if the largest free region did not grow at all. + // Using `prev_largest` (rather than `prev_largest + + // record_size / 2`) allows gradual coalescence when + // many small victims must be evicted for one large + // allocation. The `fallback_total_evicted` cap still + // bounds total eviction per key. + if (now_largest <= prev_largest) break; + prev_largest = now_largest; + allocation = allocator_->allocate(record_size); + } + } + + // ---- Handle allocation failure ---- if (!allocation.has_value()) { - LOG(ERROR) << "Failed to allocate " << record_size - << " bytes for key: " << key - << " - stopping processing for this batch"; - break; // Stop processing other keys as space is likely exhausted + if (eviction_on) { + eviction_skips_.fetch_add(1, std::memory_order_relaxed); + LOG(WARNING) << "Skipping key after eviction attempts: " << key; + continue; // eviction enabled: try next key + } else { + LOG(ERROR) << "Failed to allocate " << record_size + << " bytes for key: " << key + << " - stopping processing for this batch"; + break; // eviction disabled: preserve old break semantics + } } uint64_t offset = allocation->address(); - // Step 2: Write data to disk (no metadata locks held during I/O) + // ---- (E) Disk write (unchanged from original) ---- std::vector iovs; - iovs.reserve(2 + slices.size()); - - // Header + iovs.reserve(2 + 1 + slices.size()); iovs.push_back( {const_cast(reinterpret_cast(&header.key_len)), sizeof(header.key_len)}); iovs.push_back({const_cast( reinterpret_cast(&header.value_len)), sizeof(header.value_len)}); - - // Key iovs.push_back({const_cast(key.data()), static_cast(header.key_len)}); - - // Value slices for (const auto& slice : slices) { iovs.push_back({slice.ptr, slice.size}); } @@ -2945,78 +3316,81 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( data_file_->vector_write(iovs.data(), iovs.size(), offset); if (!write_result) { LOG(ERROR) << "Failed to write record for key: " << key - << ", error: " << write_result.error() - << " - continuing with remaining keys"; - // Allocation handle is still local (not yet stored in the metadata - // map) and will be freed automatically when going out of scope. - continue; // Continue processing other keys + << ", error: " << write_result.error(); + continue; } - - // Handle the case where the data was written partially. - size_t written = write_result.value(); - if (written != record_size) { + if (write_result.value() != record_size) { LOG(ERROR) << "Write size mismatch for key: " << key - << ", expected: " << record_size << ", got: " << written - << " - continuing with remaining keys"; - continue; // Continue processing other keys + << ", expected: " << record_size + << ", got: " << write_result.value(); + continue; } - // Step 3: Wrap allocation in refcounted handle - auto allocation_ptr = std::make_shared( - std::move(allocation.value())); - - // Step 4: Update metadata map under exclusive shard lock - // Lock only the shard for this key (other shards can proceed in - // parallel) + // ---- (F) Metadata update with FIFO index maintenance ---- { + auto allocation_ptr = std::make_shared( + std::move(allocation.value())); size_t shard_idx = ShardForKey(key); auto& shard = shards_[shard_idx]; - SharedMutexLocker lock(&shard.mutex); - // Check if key exists to update size accounting + // Lock order: eviction_mutex_ -> shard.mutex. + // Both insert and evict paths obey this order, preventing + // the overwrite-vs-evict race on fifo_index_. + std::optional ev_lock; + if (eviction_on) ev_lock.emplace(&eviction_mutex_); + SharedMutexLocker shard_lock(&shard.mutex); + auto it = shard.map.find(key); int64_t size_delta = static_cast(record_size); bool is_new_key = (it == shard.map.end()); - - if (!is_new_key) { - // Overwrite: subtract old size + uint64_t seq = 0; + + if (eviction_on) { + seq = insert_seq_.fetch_add(1, std::memory_order_relaxed); + if (!is_new_key) { + // Overwrite: drop old size and remove old FIFO slot. + size_delta -= static_cast(it->second.total_size); + fifo_index_.erase(it->second.fifo_seq); + } + } else if (!is_new_key) { size_delta -= static_cast(it->second.total_size); - // Old AllocationPtr will be dropped, refcount decremented - // Physical extent freed when last reader releases it } - // Update map (insert_or_assign handles both insert and overwrite) shard.map.insert_or_assign( - key, ObjectEntry(offset, record_size, value_size, - std::move(allocation_ptr))); + key, ObjectEntry(offset, static_cast(record_size), + value_size, std::move(allocation_ptr), seq)); - // Update total size atomically (lock-free, separate from map - // updates) - total_size_.fetch_add(size_delta, std::memory_order_relaxed); + if (eviction_on) fifo_index_.emplace(seq, key); - // Update total keys only if inserting a new key + total_size_.fetch_add(size_delta, std::memory_order_relaxed); if (is_new_key) { total_keys_.fetch_add(1, std::memory_order_relaxed); } } keys.push_back(key); - metadatas.push_back(StorageObjectMetadata{ - 0, // bucket_id not used for this backend - static_cast(offset), static_cast(header.key_len), - static_cast(value_size), ""}); + metadatas.push_back( + StorageObjectMetadata{0, static_cast(offset), + static_cast(header.key_len), + static_cast(value_size), ""}); + } + + // ---- 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()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); } - // Invoke complete handler only if we have successful keys to report if (complete_handler != nullptr && !keys.empty()) { auto error_code = complete_handler(keys, metadatas); if (error_code != ErrorCode::OK) { - LOG(ERROR) - << "Complete handler failed: " << error_code << " - " - << keys.size() - << " keys were successfully written to disk but master was not " - "notified. " - << "Master will learn about them via ScanMeta on next restart."; + LOG(ERROR) << "Complete handler failed: " << error_code << " - " + << keys.size() + << " keys were successfully written to disk but master " + "was not notified. " + << "Master will learn about them via ScanMeta on next " + "restart."; return tl::make_unexpected(error_code); } } @@ -3179,15 +3553,17 @@ OffsetAllocatorStorageBackend::IsEnableOffloading() { return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - // TODO: See if free space check is needed here. - // Check quota limits only (atomic counters, completely lock-free!) + // When eviction is enabled, BatchOffload's EvictToMakeRoom is + // responsible for making room — do not block offload here. + if (cfg_.eviction_policy != OffsetEvictionPolicy::NONE) { + return true; + } + + // Eviction disabled: keep the original quota-check behavior. bool within_size_limit = total_size_.load(std::memory_order_relaxed) < file_storage_config_.total_size_limit; - - // Check keys limit (atomic counter maintained during BatchOffload) bool within_keys_limit = total_keys_.load(std::memory_order_relaxed) < file_storage_config_.total_keys_limit; - return within_size_limit && within_keys_limit; } @@ -3301,7 +3677,14 @@ CreateStorageBackend(const FileStorageConfig& config) { config, file_per_key_backend_config); } case StorageBackendType::kOffsetAllocator: { - return std::make_shared(config); + auto offset_backend_config = + OffsetAllocatorBackendConfig::FromEnvironment(); + if (!offset_backend_config.Validate()) { + throw std::invalid_argument( + "Invalid OffsetAllocatorBackendConfig"); + } + return std::make_shared( + config, offset_backend_config); } case StorageBackendType::kDistributed: { auto distributed_config = diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 1c66f376e4..1a6722816f 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -2711,4 +2711,413 @@ TEST_F(StorageBackendTest, AdaptorBatchOffload_EvictionHandlerCalled) { //----------------------------------------------------------------------------- +//----------------------------------------------------------------------------- +// OffsetAllocatorStorageBackend Eviction Tests +//----------------------------------------------------------------------------- + +// Helper: build a BatchOffload request for a single key/value pair. +std::unordered_map> MakeSingleKeyBatch( + const std::string& key, const std::string& value, + std::vector>& buffers) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + buffers.push_back(std::move(buf)); + std::unordered_map> batch; + batch.emplace( + key, std::vector{Slice{buffers.back().get(), value.size()}}); + return batch; +} + +TEST_F(StorageBackendTest, OffsetAllocatorStorageBackend_Eviction_FifoOrder) { + // Verify that when watermark-triggered eviction fires, the oldest + // key (by insertion order) is evicted first. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 4 * 1024; // 4KB — very small arena + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Write A, B, C one by one. The arena is 4 KB; each 1 KB record + // (8 + key + value) ≈ 1040 bytes. 4 records should trigger eviction. + std::string data(1000, 'x'); + std::vector> buffers; + + for (const auto& key : {"key_a", "key_b", "key_c"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + ASSERT_TRUE(res.has_value()) << "key=" << key; + } + + // The fourth write should push total_size_ over high_watermark_bytes_ + // and evict key_a (the oldest). + auto batch = MakeSingleKeyBatch("key_d", data, buffers); + [[maybe_unused]] auto res = + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + ASSERT_TRUE(res.has_value()); + + ASSERT_FALSE(evicted_keys.empty()) + << "Should have evicted at least one key"; + EXPECT_EQ(evicted_keys[0], "key_a") + << "FIFO eviction must evict the oldest key first"; + EXPECT_FALSE(storage_backend.IsExist("key_a").value_or(true)) + << "key_a should no longer exist after eviction"; + EXPECT_TRUE(storage_backend.IsExist("key_d").value_or(false)) + << "key_d should exist"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_NoEvictionWhenNONE) { + // Under default NONE policy, the allocate-fail path still breaks. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 4 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorStorageBackend storage_backend(config); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + std::string data(1500, 'x'); + std::vector> buffers; + + int offloaded = 0; + for (const auto& key : {"key_a", "key_b", "key_c", "key_d"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + if (res.has_value()) + ++offloaded; + else + break; // allocation failure should break + } + + EXPECT_GT(offloaded, 0); + EXPECT_LT(offloaded, 5) << "NONE policy should break on allocation failure"; + EXPECT_TRUE(evicted_keys.empty()) << "NONE policy should never evict"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_MasterNotifiedBeforeReuse) { + // Verify that eviction_handler is called BEFORE allocate() for the + // key whose eviction made room, i.e. the notify-before-reuse contract. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 4 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + bool handler_called_before_allocation = false; + bool allocate_happened = false; + + auto eviction_handler = + [&handler_called_before_allocation, + &allocate_happened](const std::vector& keys) { + if (!keys.empty() && !allocate_happened) { + handler_called_before_allocation = true; + } + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // The test doesn't have direct instrumentation for "allocate() just + // happened". But the design guarantees that eviction_handler is called + // at (B) before (C) in BatchOffload. We verify indirectly: + // after writing enough to trigger eviction, the handler must have been + // invoked with at least one key AND that key no longer exists. + std::string data(1000, 'x'); + std::vector> buffers; + + for (const auto& key : {"key_a", "key_b", "key_c"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + } + + std::vector captured_evicted; + auto capture_handler = + [&captured_evicted](const std::vector& keys) { + for (const auto& k : keys) captured_evicted.push_back(k); + }; + + auto batch = MakeSingleKeyBatch("key_d", data, buffers); + storage_backend.BatchOffload(batch, complete_handler, capture_handler); + + EXPECT_FALSE(captured_evicted.empty()) + << "Should have evicted at least one key"; + for (const auto& ek : captured_evicted) { + EXPECT_FALSE(storage_backend.IsExist(ek).value_or(true)) + << "Evicted key " << ek + << " should not exist (erased before reuse)"; + } + 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. + // The evicted keys must still be flushed to the handler before return. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + // Arena large enough for several keys but small enough to trigger + // eviction near capacity. + config.total_size_limit = 8 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector all_evicted; + auto eviction_handler = + [&all_evicted](const std::vector& keys) { + for (const auto& k : keys) all_evicted.push_back(k); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + std::string data(1500, 'x'); // each record ≈ 1516 bytes + std::vector> buffers; + + // Write enough to fill the arena and trigger eviction. + for (const auto& key : {"key_a", "key_b", "key_c", "key_d", "key_e"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + // Some may succeed, some may fail — we just care that evicted + // keys are eventually reported. + } + + // If any eviction happened, the evicted keys should be reported. + // We don't assert non-empty because capacity calculations can vary; + // we just assert that if keys were evicted, they're no longer present. + if (!all_evicted.empty()) { + for (const auto& ek : all_evicted) { + EXPECT_FALSE(storage_backend.IsExist(ek).value_or(true)) + << "Evicted key " << ek << " should not exist"; + } + } +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_KeyCountTrigger) { + // Verify that eviction fires when total_keys_ exceeds the key-count + // high watermark, even when bytes are well below the byte watermark. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 1 * 1024 * 1024; // 1 MB — plenty of bytes + config.total_keys_limit = 10; // only 10 keys allowed + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Each key is tiny (10 bytes), so bytes stay low. + std::string data(10, 'x'); + std::vector> buffers; + + for (int i = 0; i < 15; ++i) { + std::string key = "tiny_key_" + std::to_string(i); + auto batch = MakeSingleKeyBatch(key, data, buffers); + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + } + + // Key-count watermark should have triggered eviction. + EXPECT_FALSE(evicted_keys.empty()) + << "Key-count overflow should trigger eviction even with low bytes"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_ConcurrentReadSafety) { + // Validate the load-bearing safety property: AllocationPtr refcount + // must keep an evicted key's physical extent alive (pinned) until + // all in-flight BatchLoad calls release their shared_ptr copies. + // + // Mechanism being tested: + // 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. + // + // This test interleaves reads and eviction-triggering writes; + // any data corruption means the allocator re-issued a still-read + // offset, which would be a violation of the refcount contract. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = + 8 * 1024; // 8 KB — tight enough that 80 keys trigger eviction + config.total_keys_limit = 500; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Pre-populate with distinct-per-key data so a re-issued offset + // is detectable: if key_i's extent is handed to a new key and the + // reader still reads from it, read_buf[0] won't match 'A' + i%26. + std::vector> buffers; + const int kNumKeys = 30; + for (int i = 0; i < kNumKeys; ++i) { + std::string key = "ckey_" + std::to_string(i); + std::string val(100, static_cast('A' + (i % 26))); + auto batch = MakeSingleKeyBatch(key, val, buffers); + storage_backend.BatchOffload(batch, complete_handler); + } + // Confirm pre-populated keys exist before the stress phase. + for (int i = 0; i < std::min(kNumKeys, 5); ++i) { + EXPECT_TRUE(storage_backend.IsExist("ckey_" + std::to_string(i)) + .value_or(false)); + } + + // Eviction handler that tracks victim keys so we can assert post-hoc. + std::vector all_evicted; + std::mutex evict_mtx; + auto eviction_handler = [&all_evicted, + &evict_mtx](const std::vector& keys) { + std::lock_guard lk(evict_mtx); + for (const auto& k : keys) all_evicted.push_back(k); + }; + + std::atomic stop{false}; + std::atomic read_errors{0}; + std::atomic read_success{0}; + std::atomic read_not_found{0}; + + std::thread reader([&]() { + while (!stop) { + for (int i = 0; i < kNumKeys; ++i) { + std::string key = "ckey_" + std::to_string(i); + auto read_buf = std::make_unique(100); + std::unordered_map load; + load.emplace(key, Slice{read_buf.get(), 100}); + auto res = storage_backend.BatchLoad(load); + if (res.has_value()) { + read_success++; + if (read_buf[0] != static_cast('A' + (i % 26))) { + read_errors++; + } + } else { + read_not_found++; // expected after eviction + } + } + } + }); + + std::thread writer([&]() { + for (int i = kNumKeys; i < kNumKeys + 50 && !stop; ++i) { + std::string key = "newkey_" + std::to_string(i); + std::string val(100, 'Z'); + std::vector> wbufs; + auto batch = MakeSingleKeyBatch(key, val, wbufs); + storage_backend.BatchOffload(batch, complete_handler, + eviction_handler); + } + stop = true; + }); + + writer.join(); + reader.join(); + + // Core safety assertion: zero data corruption across all reads. + EXPECT_EQ(read_errors.load(), 0) + << "Refcount must prevent allocator from re-issuing in-use extents"; + + // Sanity: some reads succeeded and some keys were evicted. + EXPECT_GT(read_success.load(), 0); + { + std::lock_guard lk(evict_mtx); + EXPECT_FALSE(all_evicted.empty()) + << "Eviction must have occurred during concurrent stress"; + } + + // Post-condition: at least one evicted key is no longer in the map. + { + std::lock_guard lk(evict_mtx); + bool any_gone = false; + for (const auto& ek : all_evicted) { + if (!storage_backend.IsExist(ek).value_or(true)) { + any_gone = true; + break; + } + } + EXPECT_TRUE(any_gone) << "Evicted keys must be removed from shard.map"; + } +} + } // namespace mooncake::test