From 800ad5eacb031901cf2c35f601660a5548d7406c Mon Sep 17 00:00:00 2001 From: dusterbloom <32869278+dusterbloom@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:48:46 +0200 Subject: [PATCH] fix(prefix-cache): reuse slots left by aborted snapshots --- .github/workflows/ci.yml | 5 +- server/CMakeLists.txt | 15 ++ server/src/server/prefix_cache.cpp | 184 +++++----------- server/src/server/prefix_cache.h | 32 +-- server/src/server/prefix_cache_state.h | 281 ++++++++++++++++++++++++ server/test/test_prefix_cache_state.cpp | 201 +++++++++++++++++ 6 files changed, 562 insertions(+), 156 deletions(-) create mode 100644 server/src/server/prefix_cache_state.h create mode 100644 server/test/test_prefix_cache_state.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3cfe3e446..d2df963ab 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,12 +79,15 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build --target \ test_dflash test_generate test_flash_attn_sparse test_server_unit \ + test_prefix_cache_state \ test_deepseek4_unit -j$(nproc) - name: Run C++ server unit tests run: | cd server/build - ctest --output-on-failure -R "server_unit|deepseek4_unit" --no-tests=error + ctest --output-on-failure \ + -R "server_unit|prefix_cache_state|deepseek4_unit" \ + --no-tests=error - name: Populate venv with cu128 torch + setuptools # First pass: install the workspace's default deps. dflash declares diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 0da29ea33..5bb2de90e 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1163,6 +1163,18 @@ if(DFLASH27B_TESTS) list(APPEND _raw_unit_test_targets test_server_unit) endif() + # Dependency-free production state core used by native tests and ESBMC. + # Keep this target free of tokenizer, ggml, and GPU libraries so formal + # transition checks remain fast and reproducible on hosted CPU runners. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_prefix_cache_state.cpp") + add_executable( + test_prefix_cache_state + test/test_prefix_cache_state.cpp) + target_include_directories(test_prefix_cache_state PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + add_test(NAME prefix_cache_state COMMAND test_prefix_cache_state) + endif() + # Feature/architecture gate tests. check_feature_compatibility(), # collect_feature_warnings() and the capability table are pure functions, # so this target deliberately compiles only feature_gate.cpp and @@ -1236,6 +1248,9 @@ if(DFLASH27B_TESTS) if(TARGET test_feature_gate) list(APPEND _check_deps test_feature_gate) endif() + if(TARGET test_prefix_cache_state) + list(APPEND _check_deps test_prefix_cache_state) + endif() if(_check_deps) add_custom_target(check COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index 4ab20c590..0768363ab 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -140,38 +140,6 @@ PrefixHash hash_prefix(const int32_t * ids, int count) { return h; } -// ─── Prefix-aware eviction ────────────────────────────────────────────── - -static bool is_strict_prefix(const std::vector & a, - const std::vector & b) { - // True iff `a` is a strict (shorter) prefix of `b`. - if (a.size() >= b.size()) return false; - return std::equal(a.begin(), a.end(), b.begin()); -} - -int select_inline_evict_victim(const std::vector *> & ids_lru) { - const int n = (int)ids_lru.size(); - if (n <= 0) return 0; - // Oldest-first scan: evict the first entry that is not a strict prefix of any - // other entry (a leaf). Shared ancestor prefixes are thereby kept resident. - for (int i = 0; i < n; i++) { - bool is_ancestor = false; - for (int j = 0; j < n; j++) { - if (j == i) continue; - if (is_strict_prefix(*ids_lru[i], *ids_lru[j])) { is_ancestor = true; break; } - } - if (!is_ancestor) return i; // oldest leaf - } - return 0; // unreachable (the longest entry is always a leaf); pure-LRU fallback -} - -int select_inline_evict_victim(const std::vector> & ids_lru) { - std::vector *> ptrs; - ptrs.reserve(ids_lru.size()); - for (const auto & v : ids_lru) ptrs.push_back(&v); - return select_inline_evict_victim(ptrs); -} - int select_inline_snapshot_boundary(const std::vector & boundaries, int restored_prefix_len) { if (boundaries.empty()) return 0; @@ -184,7 +152,8 @@ int select_inline_snapshot_boundary(const std::vector & boundaries, // ─── PrefixCache ──────────────────────────────────────────────────────── PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) - : cap_(std::min(cap, MAX_SLOTS)) + : cap_(std::min(cap, MAX_SLOTS)), + inline_state_(std::max(0, std::min(cap, MAX_SLOTS))) { if (cap_ <= 0) { disabled_ = true; @@ -203,18 +172,9 @@ PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer) // ── LRU helpers ───────────────────────────────────────────────────────── -int PrefixCache::find_entry(const PrefixHash & h) const { - for (int i = 0; i < (int)entries_.size(); i++) { - if (entries_[i].hash == h) return i; - } - return -1; -} - -void PrefixCache::move_to_end(int idx) { - if (idx < 0 || idx >= (int)entries_.size()) return; - auto e = std::move(entries_[idx]); - entries_.erase(entries_.begin() + idx); - entries_.push_back(std::move(e)); +void PrefixCache::sync_inline_size() { + entries_size_count_.store( + inline_state_.size(), std::memory_order_relaxed); } int PrefixCache::find_full_entry(const PrefixHash & h) const { @@ -241,26 +201,21 @@ std::pair PrefixCache::lookup(const std::vector & prompt_ids) for (int cut : boundaries) { auto key = hash_prefix(prompt_ids.data(), cut); - int idx = find_entry(key); - if (idx >= 0) { - const int committed = (int)entries_[idx].ids.size(); - if (committed != cut) { - // Slot was refreshed in-place at a deeper boundary; a shallow - // hash→slot entry would restore the wrong cur_pos. - std::fprintf(stderr, - "[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n", - entries_[idx].slot, cut, committed); - entries_.erase(entries_.begin() + idx); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - continue; - } - if (cut > best_len) { - best_slot = entries_[idx].slot; - best_len = cut; - } - move_to_end(idx); + const auto result = inline_state_.lookup_candidate(key, cut); + if (result.stale_removed) { + // Slot was refreshed in-place at a deeper boundary; a shallow + // hash→slot entry would restore the wrong cur_pos. + std::fprintf(stderr, + "[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n", + result.stale_slot, cut, result.stale_committed_len); + continue; + } + if (result.slot >= 0 && cut > best_len) { + best_slot = result.slot; + best_len = result.prefix_len; } } + sync_inline_size(); if (best_slot >= 0) { lifetime_hits_.fetch_add(1, std::memory_order_relaxed); @@ -281,68 +236,44 @@ std::pair PrefixCache::prepare_inline_snap( if (target_cut <= 0) return {-1, 0}; auto key = hash_prefix(prompt_ids.data(), target_cut); - if (find_entry(key) >= 0) return {-1, 0}; // already cached - - int slot; - if ((int)entries_.size() >= cap_) { - // At capacity — reserve a slot without evicting yet. Prefix-aware: prefer - // the oldest leaf so shared ancestor prefixes (reused by later branches) - // stay resident. entries_ is already in LRU order (front = oldest). - std::vector *> ids_lru; - ids_lru.reserve(entries_.size()); - for (const auto & e : entries_) ids_lru.push_back(&e.ids); - int victim = select_inline_evict_victim(ids_lru); - pending_evict_key_ = entries_[victim].hash; - has_pending_evict_ = true; - slot = entries_[victim].slot; - if (victim != 0) { - std::fprintf(stderr, - "[pc] prefix-aware evict: victim idx=%d (len=%zu) kept oldest " - "ancestor (len=%zu)\n", - victim, entries_[victim].ids.size(), entries_.front().ids.size()); - } - } else { - slot = next_slot_; - next_slot_ = (next_slot_ + 1) % cap_; - has_pending_evict_ = false; + const auto reservation = inline_state_.prepare(key, target_cut); + if (reservation.slot < 0) return {-1, 0}; + if (reservation.victim_index > 0) { + std::fprintf(stderr, + "[pc] prefix-aware evict: victim idx=%d (len=%d) kept oldest " + "ancestor (len=%d)\n", + reservation.victim_index, reservation.victim_len, + reservation.oldest_len); } - return {slot, target_cut}; + return {reservation.slot, reservation.target_cut}; } void PrefixCache::confirm_inline_snap(int slot, int target_cut, const std::vector & prompt_ids) { if (disabled_) return; - - // Evict the reserved entry (if any). - if (has_pending_evict_) { - int idx = find_entry(pending_evict_key_); - if (idx >= 0) { - entries_.erase(entries_.begin() + idx); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } - has_pending_evict_ = false; + if (slot < 0 || slot >= cap_ || target_cut <= 0 || + target_cut > (int)prompt_ids.size()) { + std::fprintf(stderr, + "[pc] rejected inline-snap slot=%d prefix_len=%d prompt_len=%zu\n", + slot, target_cut, prompt_ids.size()); + return; } - // The new snapshot replaces whatever this slot previously held. Drop any - // other entries still pointing at the slot: their hashes describe a - // different (or shorter) token stream than the new snapshot, and a later - // restore through them would attach mismatched KV. Stale entries arise - // when an aborted snap burns a round-robin next_slot_ step and a later - // confirm wraps onto a slot with a live entry (PR #370 repro). - for (int i = (int)entries_.size() - 1; i >= 0; --i) { - if (entries_[(size_t)i].slot == slot) { - std::fprintf(stderr, - "[pc] dropping stale entry for reused slot=%d\n", slot); - entries_.erase(entries_.begin() + i); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } + const auto key = hash_prefix(prompt_ids.data(), target_cut); + const auto result = + inline_state_.confirm(slot, key, target_cut, prompt_ids); + if (!result.accepted) { + std::fprintf(stderr, + "[pc] rejected inline-snap slot=%d prefix_len=%d prompt_len=%zu\n", + slot, target_cut, prompt_ids.size()); + return; } - - auto key = hash_prefix(prompt_ids.data(), target_cut); - std::vector ids(prompt_ids.begin(), prompt_ids.begin() + target_cut); - entries_.push_back({key, slot, std::move(ids)}); - entries_size_count_.fetch_add(1, std::memory_order_relaxed); + for (int i = 0; i < result.stale_slot_entries_removed; ++i) { + std::fprintf(stderr, + "[pc] dropping stale entry for reused slot=%d\n", slot); + } + sync_inline_size(); std::fprintf(stderr, "[pc] inline-snap committed slot=%d prefix_len=%d\n", slot, target_cut); } @@ -353,31 +284,20 @@ void PrefixCache::abort_inline_snap(int slot) { // metadata still pointing at it is therefore invalid, whether the slot was // selected through the explicit eviction path or through a round-robin // hole left by an earlier aborted reservation. - for (int i = (int)entries_.size() - 1; i >= 0; --i) { - if (entries_[(size_t)i].slot == slot) { - entries_.erase(entries_.begin() + i); - entries_size_count_.fetch_sub(1, std::memory_order_relaxed); - } - } - has_pending_evict_ = false; + inline_state_.abort(slot); + sync_inline_size(); } void PrefixCache::cancel_inline_snap(int slot) { if (disabled_) return; - if (has_pending_evict_) { - const int idx = find_entry(pending_evict_key_); - if (idx >= 0 && entries_[idx].slot != slot) return; - } - has_pending_evict_ = false; + inline_state_.cancel(slot); } void PrefixCache::mark_all_cleared() { if (disabled_) return; - int n = (int)entries_.size(); - entries_.clear(); - entries_size_count_.store(0, std::memory_order_relaxed); - next_slot_ = 0; - has_pending_evict_ = false; + const int n = inline_state_.size(); + inline_state_.clear(); + sync_inline_size(); std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n); } diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 2a0515749..dd0d7b1de 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -12,9 +12,9 @@ #pragma once +#include "prefix_cache_state.h" #include "tokenizer.h" -#include #include #include #include @@ -42,7 +42,6 @@ std::vector find_all_boundaries(const std::vector & ids, const ChatMarkers & markers); // SHA-1 hash of a prefix (truncated to 16 bytes). -using PrefixHash = std::array; PrefixHash hash_prefix(const int32_t * ids, int count); // Prefix-aware inline eviction policy. Given the cached prefixes in LRU order @@ -50,12 +49,8 @@ PrefixHash hash_prefix(const int32_t * ids, int count); // whose ids are NOT a strict prefix of any other entry's ids (a "leaf"). Keeping // shared ancestor prefixes resident avoids re-prefilling them for later branches. // Returns 0 (pure-LRU fallback) when ids_lru is empty or, impossibly, no leaf -// is found. Pure and model-free so it can be unit-tested without a PrefixCache. -// The pointer overload is the core (the caller passes pointers into its own -// entries so no token vectors are copied); the value overload is a convenience -// wrapper for tests. -int select_inline_evict_victim(const std::vector *> & ids_lru); -int select_inline_evict_victim(const std::vector> & ids_lru); +// is found. The implementation lives in prefix_cache_state.h so the exact +// production policy can also be model-checked without server dependencies. // Pick the inline snapshot boundary for a request. We cache the boundary before // the current user turn (second-to-last marker) and only when it advances past @@ -163,17 +158,9 @@ class PrefixCache { int cap_ = 0; ChatMarkers markers_; - // LRU for inline prefix cache: ordered map of hash → slot. - // We use a vector to maintain insertion order (front = oldest). - struct LruEntry { - PrefixHash hash; - int slot; - std::vector ids; // prefix tokens [0, target_cut) for prefix-aware eviction - }; - std::vector entries_; - int next_slot_ = 0; - PrefixHash pending_evict_key_{}; - bool has_pending_evict_ = false; + // Boundary detection and hashing live in PrefixCache; all inline-cache + // transitions live in this dependency-free core shared with ESBMC. + InlinePrefixCacheState inline_state_; // Full-cache state bool full_disabled_ = true; @@ -194,19 +181,18 @@ class PrefixCache { std::atomic lifetime_hits_{0}; // inline cache hits std::atomic full_lifetime_hits_{0}; // full-compress cache hits std::atomic full_disk_bytes_{0}; // best-effort snapshot of disk usage - // Atomic mirrors of `entries_.size()` and `full_entries_.size()`. + // Atomic mirrors of `inline_state_.size()` and `full_entries_.size()`. // The vectors themselves are mutated only on the daemon thread // under the daemon's serialised request loop, but `/props` reads // happen from the client thread — calling `.size()` there is a // data race per the C++ memory model. Bump these alongside every // push_back / erase / clear so the public introspection counters // stay well-defined. (Codex r1 P2 follow-up.) - std::atomic entries_size_count_{0}; // mirrors entries_.size() + std::atomic entries_size_count_{0}; // mirrors inline_state_.size() std::atomic full_entries_size_count_{0}; // mirrors full_entries_.size() // Helpers - int find_entry(const PrefixHash & h) const; - void move_to_end(int idx); + void sync_inline_size(); int find_full_entry(const PrefixHash & h) const; void move_full_to_end(int idx); }; diff --git a/server/src/server/prefix_cache_state.h b/server/src/server/prefix_cache_state.h new file mode 100644 index 000000000..e14127565 --- /dev/null +++ b/server/src/server/prefix_cache_state.h @@ -0,0 +1,281 @@ +// Verification-friendly state core for the inline prefix cache. +// +// This header deliberately contains no tokenizer, hashing, ggml, CUDA, or +// server dependencies. PrefixCache performs boundary detection and key +// derivation, then delegates its state transitions here. The same production +// transition code is therefore usable by native unit tests and ESBMC harnesses. + +#pragma once + +#include +#include +#include +#include +#include + +namespace dflash::common { + +using PrefixHash = std::array; + +// Keep equality explicit instead of delegating to std::array's loop. Besides +// being cheap for a fixed 128-bit key, this gives model checkers a finite, +// fully unrolled comparison while remaining the production implementation. +inline bool prefix_hash_equal( + const PrefixHash & lhs, const PrefixHash & rhs) { + return lhs[0] == rhs[0] && lhs[1] == rhs[1] && + lhs[2] == rhs[2] && lhs[3] == rhs[3] && + lhs[4] == rhs[4] && lhs[5] == rhs[5] && + lhs[6] == rhs[6] && lhs[7] == rhs[7] && + lhs[8] == rhs[8] && lhs[9] == rhs[9] && + lhs[10] == rhs[10] && lhs[11] == rhs[11] && + lhs[12] == rhs[12] && lhs[13] == rhs[13] && + lhs[14] == rhs[14] && lhs[15] == rhs[15]; +} + +namespace prefix_cache_detail { + +inline bool is_strict_prefix(const std::vector & a, + const std::vector & b) { + if (a.size() >= b.size()) return false; + return std::equal(a.begin(), a.end(), b.begin()); +} + +} // namespace prefix_cache_detail + +// Prefix-aware inline eviction policy. Inputs are in LRU order +// (index 0 = oldest). Prefer the oldest leaf so shared ancestors remain hot. +inline int select_inline_evict_victim( + const std::vector *> & ids_lru) { + const int n = (int)ids_lru.size(); + if (n <= 0) return 0; + for (int i = 0; i < n; ++i) { + bool is_ancestor = false; + for (int j = 0; j < n; ++j) { + if (j == i) continue; + if (prefix_cache_detail::is_strict_prefix( + *ids_lru[(size_t)i], *ids_lru[(size_t)j])) { + is_ancestor = true; + break; + } + } + if (!is_ancestor) return i; + } + return 0; +} + +inline int select_inline_evict_victim( + const std::vector> & ids_lru) { + std::vector *> ptrs; + ptrs.reserve(ids_lru.size()); + for (const auto & ids : ids_lru) ptrs.push_back(&ids); + return select_inline_evict_victim(ptrs); +} + +// Select a slot below the production PrefixCache limit of 64 slots. Keeping +// this allocation decision scalar makes the behavior independently +// model-checkable; InlinePrefixCacheState remains responsible for deriving the +// occupancy mask from its committed entries. +inline int select_inline_free_slot( + int next_slot, int capacity, uint64_t occupied_slots) { + if (capacity <= 0 || next_slot < 0 || next_slot >= capacity) { + return -1; + } + // InlinePrefixCacheState is independently usable, while the production + // PrefixCache clamps capacity to 64. Preserve the legacy round-robin + // behavior for out-of-contract standalone capacities that do not fit the + // occupancy mask. + if (capacity > 64) return next_slot; + + for (int offset = 0; offset < capacity; ++offset) { + const int candidate = (next_slot + offset) % capacity; + const uint64_t candidate_bit = uint64_t{1} << candidate; + if ((occupied_slots & candidate_bit) == 0) { + return candidate; + } + } + return -1; +} + +class InlinePrefixCacheState { +public: + struct Entry { + PrefixHash hash{}; + int slot = -1; + std::vector ids; + }; + + struct LookupResult { + int slot = -1; + int prefix_len = 0; + bool stale_removed = false; + int stale_slot = -1; + int stale_committed_len = 0; + }; + + struct PrepareResult { + int slot = -1; + int target_cut = 0; + int victim_index = -1; + int victim_len = 0; + int oldest_len = 0; + }; + + struct ConfirmResult { + bool accepted = false; + int pending_removed = 0; + int stale_slot_entries_removed = 0; + }; + + explicit InlinePrefixCacheState(int capacity = 0) + : capacity_(std::max(0, capacity)) {} + + int capacity() const { return capacity_; } + int size() const { return (int)entries_.size(); } + int next_slot() const { return next_slot_; } + bool has_pending_eviction() const { return has_pending_evict_; } + const PrefixHash & pending_eviction_key() const { + return pending_evict_key_; + } + const std::vector & entries() const { return entries_; } + + int find(const PrefixHash & hash) const { + for (int i = 0; i < (int)entries_.size(); ++i) { + if (prefix_hash_equal(entries_[(size_t)i].hash, hash)) return i; + } + return -1; + } + + bool contains(const PrefixHash & hash) const { return find(hash) >= 0; } + + LookupResult lookup_candidate(const PrefixHash & hash, int cut) { + LookupResult result; + const int idx = find(hash); + if (idx < 0) return result; + + const int committed = (int)entries_[(size_t)idx].ids.size(); + if (committed != cut) { + result.stale_removed = true; + result.stale_slot = entries_[(size_t)idx].slot; + result.stale_committed_len = committed; + entries_.erase(entries_.begin() + idx); + return result; + } + + result.slot = entries_[(size_t)idx].slot; + result.prefix_len = cut; + move_to_end(idx); + return result; + } + + PrepareResult prepare(const PrefixHash & hash, int target_cut) { + PrepareResult result; + if (capacity_ <= 0 || target_cut <= 0 || contains(hash)) return result; + + result.target_cut = target_cut; + if ((int)entries_.size() >= capacity_) { + std::vector *> ids_lru; + ids_lru.reserve(entries_.size()); + for (const auto & entry : entries_) ids_lru.push_back(&entry.ids); + + const int victim = select_inline_evict_victim(ids_lru); + pending_evict_key_ = entries_[(size_t)victim].hash; + has_pending_evict_ = true; + result.slot = entries_[(size_t)victim].slot; + result.victim_index = victim; + result.victim_len = + (int)entries_[(size_t)victim].ids.size(); + result.oldest_len = (int)entries_.front().ids.size(); + } else { + uint64_t occupied_slots = 0; + if (capacity_ <= 64) { + for (const auto & entry : entries_) { + if (entry.slot >= 0 && entry.slot < 64) { + occupied_slots |= uint64_t{1} << entry.slot; + } + } + } + result.slot = select_inline_free_slot( + next_slot_, capacity_, occupied_slots); + next_slot_ = (result.slot + 1) % capacity_; + has_pending_evict_ = false; + } + return result; + } + + ConfirmResult confirm(int slot, const PrefixHash & hash, int target_cut, + const std::vector & prompt_ids) { + ConfirmResult result; + if (slot < 0 || slot >= capacity_ || target_cut <= 0 || + target_cut > (int)prompt_ids.size()) { + return result; + } + + if (has_pending_evict_) { + const int idx = find(pending_evict_key_); + if (idx >= 0) { + entries_.erase(entries_.begin() + idx); + result.pending_removed = 1; + } + has_pending_evict_ = false; + } + + for (int i = (int)entries_.size() - 1; i >= 0; --i) { + if (entries_[(size_t)i].slot == slot) { + entries_.erase(entries_.begin() + i); + ++result.stale_slot_entries_removed; + } + } + + std::vector ids( + prompt_ids.begin(), prompt_ids.begin() + target_cut); + entries_.push_back({hash, slot, std::move(ids)}); + result.accepted = true; + return result; + } + + int abort(int slot) { + int removed = 0; + for (int i = (int)entries_.size() - 1; i >= 0; --i) { + if (entries_[(size_t)i].slot == slot) { + entries_.erase(entries_.begin() + i); + ++removed; + } + } + has_pending_evict_ = false; + return removed; + } + + // Returns false only when the supplied slot does not own the pending + // reservation. In that case state is left untouched. + bool cancel(int slot) { + if (has_pending_evict_) { + const int idx = find(pending_evict_key_); + if (idx >= 0 && entries_[(size_t)idx].slot != slot) return false; + } + has_pending_evict_ = false; + return true; + } + + void clear() { + entries_.clear(); + next_slot_ = 0; + has_pending_evict_ = false; + pending_evict_key_ = {}; + } + +private: + void move_to_end(int idx) { + if (idx < 0 || idx >= (int)entries_.size()) return; + auto entry = std::move(entries_[(size_t)idx]); + entries_.erase(entries_.begin() + idx); + entries_.push_back(std::move(entry)); + } + + int capacity_ = 0; + std::vector entries_; + int next_slot_ = 0; + PrefixHash pending_evict_key_{}; + bool has_pending_evict_ = false; +}; + +} // namespace dflash::common diff --git a/server/test/test_prefix_cache_state.cpp b/server/test/test_prefix_cache_state.cpp new file mode 100644 index 000000000..8bc0c6244 --- /dev/null +++ b/server/test/test_prefix_cache_state.cpp @@ -0,0 +1,201 @@ +#include "server/prefix_cache_state.h" + +#include +#include +#include +#include + +using dflash::common::InlinePrefixCacheState; +using dflash::common::PrefixHash; +using dflash::common::prefix_hash_equal; +using dflash::common::select_inline_evict_victim; + +namespace { + +PrefixHash key(uint8_t family, uint8_t depth) { + PrefixHash result{}; + result[0] = family; + result[1] = depth; + return result; +} + +void assert_invariants(const InlinePrefixCacheState & state) { + assert(state.capacity() >= 0); + assert(state.size() >= 0); + assert(state.size() <= state.capacity()); + if (state.capacity() > 0) { + assert(state.next_slot() >= 0); + assert(state.next_slot() < state.capacity()); + } + + const auto & entries = state.entries(); + for (size_t i = 0; i < entries.size(); ++i) { + assert(entries[i].slot >= 0); + assert(entries[i].slot < state.capacity()); + assert(!entries[i].ids.empty()); + for (size_t j = i + 1; j < entries.size(); ++j) { + assert(entries[i].slot != entries[j].slot); + assert(!prefix_hash_equal(entries[i].hash, entries[j].hash)); + } + } + + if (state.has_pending_eviction()) { + assert(state.contains(state.pending_eviction_key())); + } +} + +void test_round_robin_and_reuse() { + InlinePrefixCacheState state(2); + const std::vector a = {7, 10}; + const std::vector b = {8, 20}; + const std::vector c = {9, 30}; + + auto ra = state.prepare(key(1, 2), 2); + assert(ra.slot == 0); + assert(state.confirm(ra.slot, key(1, 2), 2, a).accepted); + + auto rb = state.prepare(key(2, 2), 2); + assert(rb.slot == 1); + assert(state.confirm(rb.slot, key(2, 2), 2, b).accepted); + assert_invariants(state); + + auto rc = state.prepare(key(3, 2), 2); + assert(rc.slot == 0); + assert(rc.victim_index == 0); + auto confirmed = state.confirm(rc.slot, key(3, 2), 2, c); + assert(confirmed.accepted); + assert(confirmed.pending_removed == 1); + assert(!state.contains(key(1, 2))); + assert(state.contains(key(2, 2))); + assert(state.contains(key(3, 2))); + assert_invariants(state); +} + +void test_abort_purges_reused_slot() { + InlinePrefixCacheState state(2); + const std::vector a = {7}; + const std::vector b = {8}; + + auto ra = state.prepare(key(1, 1), 1); + assert(state.confirm(ra.slot, key(1, 1), 1, a).accepted); + auto rb = state.prepare(key(2, 1), 1); + assert(state.confirm(rb.slot, key(2, 1), 1, b).accepted); + + auto pending = state.prepare(key(3, 1), 1); + assert(pending.slot >= 0); + state.abort(pending.slot); + for (const auto & entry : state.entries()) { + assert(entry.slot != pending.slot); + } + assert(!state.has_pending_eviction()); + assert_invariants(state); +} + +void test_abort_reuses_hole_before_occupied_slot() { + InlinePrefixCacheState state(2); + const std::vector a = {7}; + + auto committed = state.prepare(key(1, 1), 1); + assert(committed.slot == 0); + assert(state.confirm( + committed.slot, key(1, 1), 1, a).accepted); + + // Reserving the second slot advances the round-robin cursor back to slot + // zero. If that reservation aborts, slot one is a hole and slot zero still + // owns a valid snapshot. + auto failed = state.prepare(key(2, 1), 1); + assert(failed.slot == 1); + assert(state.abort(failed.slot) == 0); + assert(state.contains(key(1, 1))); + + // The HTTP layer immediately frees the slot returned by prepare(). It must + // therefore receive the free slot, not the occupied slot zero. + auto replacement = state.prepare(key(3, 1), 1); + assert(replacement.slot == failed.slot); + for (const auto & entry : state.entries()) { + assert(entry.slot != replacement.slot); + } + assert(state.contains(key(1, 1))); + assert_invariants(state); +} + +void test_cancel_preserves_entry() { + InlinePrefixCacheState state(1); + const std::vector ids = {7, 10}; + auto initial = state.prepare(key(1, 2), 2); + assert(state.confirm(initial.slot, key(1, 2), 2, ids).accepted); + + auto pending = state.prepare(key(2, 2), 2); + assert(pending.slot == initial.slot); + assert(state.has_pending_eviction()); + assert(state.cancel(pending.slot)); + assert(state.contains(key(1, 2))); + assert(!state.contains(key(2, 2))); + assert(!state.has_pending_eviction()); + assert_invariants(state); +} + +void test_stale_lookup_is_removed() { + InlinePrefixCacheState state(2); + const std::vector ids = {7, 10}; + auto reservation = state.prepare(key(1, 2), 2); + assert(state.confirm( + reservation.slot, key(1, 2), 2, ids).accepted); + + const auto stale = state.lookup_candidate(key(1, 2), 1); + assert(stale.stale_removed); + assert(stale.stale_slot == reservation.slot); + assert(stale.stale_committed_len == 2); + assert(state.size() == 0); + assert_invariants(state); +} + +void test_invalid_confirm_is_non_mutating() { + InlinePrefixCacheState state(2); + const std::vector ids = {7}; + assert(!state.confirm(-1, key(1, 1), 1, ids).accepted); + assert(!state.confirm(2, key(1, 1), 1, ids).accepted); + assert(!state.confirm(0, key(1, 2), 2, ids).accepted); + assert(state.size() == 0); + assert_invariants(state); +} + +void test_prefix_aware_eviction() { + const std::vector> chain = { + {7}, {7, 10}, {7, 10, 20}, + }; + assert(select_inline_evict_victim(chain) == 2); + + const std::vector> branch = { + {7}, {7, 10}, {7, 20}, + }; + assert(select_inline_evict_victim(branch) == 1); +} + +void test_clear_resets_allocator() { + InlinePrefixCacheState state(2); + const std::vector ids = {7}; + auto reservation = state.prepare(key(1, 1), 1); + assert(state.confirm( + reservation.slot, key(1, 1), 1, ids).accepted); + state.clear(); + assert(state.size() == 0); + assert(state.next_slot() == 0); + assert(!state.has_pending_eviction()); + assert_invariants(state); +} + +} // namespace + +int main() { + test_round_robin_and_reuse(); + test_abort_purges_reused_slot(); + test_abort_reuses_hole_before_occupied_slot(); + test_cancel_preserves_entry(); + test_stale_lookup_is_removed(); + test_invalid_confirm_is_non_mutating(); + test_prefix_aware_eviction(); + test_clear_resets_allocator(); + ::puts("prefix_cache_state: PASS"); + return 0; +}