diff --git a/src/ailego/buffer/block_eviction_queue.cc b/src/ailego/buffer/block_eviction_queue.cc index eff930127..571f09dbd 100644 --- a/src/ailego/buffer/block_eviction_queue.cc +++ b/src/ailego/buffer/block_eviction_queue.cc @@ -12,127 +12,990 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#else +#include +#include +#if !defined(MAP_ANONYMOUS) && defined(MAP_ANON) +#define MAP_ANONYMOUS MAP_ANON +#endif +#endif + namespace zvec { namespace ailego { -int BlockEvictionQueue::init() { - evict_batch_size_ = 512; - for (size_t i = 0; i < CACHE_QUEUE_NUM; i++) { - evict_queues_.push_back(ConcurrentQueue(evict_batch_size_ * 200)); +namespace { + +constexpr size_t kMinimumPageBytes = 4096UL; +constexpr size_t kSlabBytes = MemoryLimitPool::slab_size(); +constexpr size_t kSlabAlignment = MemoryLimitPool::slab_alignment(); +constexpr size_t kMaxSlabDataPages = kSlabBytes / kMinimumPageBytes - 1; + +static_assert((kSlabAlignment & (kSlabAlignment - 1)) == 0, + "slab alignment must be a power of two"); +static_assert(kSlabBytes == kSlabAlignment, + "slab owner lookup relies on size matching alignment"); +static_assert(kMaxSlabDataPages <= UINT16_MAX, + "slab page indexes must fit in uint16_t"); + +struct AlignedSlabMapping { + char *base{nullptr}; + void *reservation_base{nullptr}; + size_t reservation_size{0}; +}; + +AlignedSlabMapping reserve_aligned_slab(size_t page_size) { + constexpr size_t kReservationBytes = kSlabBytes + kSlabAlignment; +#if defined(_WIN32) + void *reservation = + ::VirtualAlloc(nullptr, kReservationBytes, MEM_RESERVE, PAGE_READWRITE); + if (reservation == nullptr) { + return {}; } - return 0; + const uintptr_t raw = reinterpret_cast(reservation); + const uintptr_t aligned = (raw + kSlabAlignment - 1) & ~(kSlabAlignment - 1); + void *header = ::VirtualAlloc(reinterpret_cast(aligned), page_size, + MEM_COMMIT, PAGE_READWRITE); + if (header == nullptr) { + ::VirtualFree(reservation, 0, MEM_RELEASE); + return {}; + } + return {reinterpret_cast(aligned), reservation, kReservationBytes}; +#else + (void)page_size; + void *reservation = ::mmap(nullptr, kReservationBytes, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (reservation == MAP_FAILED) { + return {}; + } + + const uintptr_t raw = reinterpret_cast(reservation); + const uintptr_t aligned = (raw + kSlabAlignment - 1) & ~(kSlabAlignment - 1); + const size_t prefix = aligned - raw; + const size_t suffix = kReservationBytes - prefix - kSlabBytes; + if (prefix != 0 && ::munmap(reservation, prefix) != 0) { + ::munmap(reservation, kReservationBytes); + return {}; + } + if (suffix != 0 && + ::munmap(reinterpret_cast(aligned + kSlabBytes), suffix) != 0) { + ::munmap(reinterpret_cast(aligned), kSlabBytes + suffix); + return {}; + } + return {reinterpret_cast(aligned), reinterpret_cast(aligned), + kSlabBytes}; +#endif +} + +void release_aligned_slab(void *reservation_base, size_t reservation_size) { +#if defined(_WIN32) + (void)reservation_size; + ::VirtualFree(reservation_base, 0, MEM_RELEASE); +#else + ::munmap(reservation_base, reservation_size); +#endif +} + +bool commit_slab_page(char *page, size_t page_size, bool reclaimed) { +#if defined(_WIN32) + (void)reclaimed; + return ::VirtualAlloc(page, page_size, MEM_COMMIT, PAGE_READWRITE) == page; +#elif defined(__APPLE__) && defined(MADV_FREE_REUSE) + if (!reclaimed) { + return true; + } + return ::madvise(page, page_size, MADV_FREE_REUSE) == 0; +#else + (void)page; + (void)page_size; + (void)reclaimed; + return true; +#endif +} + +bool discard_slab_page(char *page, size_t page_size) { +#if defined(_WIN32) + return ::VirtualFree(page, page_size, MEM_DECOMMIT) != 0; +#elif defined(__APPLE__) && defined(MADV_FREE_REUSABLE) + return ::madvise(page, page_size, MADV_FREE_REUSABLE) == 0; +#else + return ::madvise(page, page_size, MADV_DONTNEED) == 0; +#endif +} + +} // namespace + +struct MemoryLimitPool::ReclaimableSlab { + static constexpr uint64_t kMagic = 0x5A564543534C4142ULL; + + ReclaimableSlab(MemoryLimitPool *pool, void *reservation, + size_t reservation_bytes, size_t system_page_size) + : owner(pool), + reservation_base(reservation), + reservation_size(reservation_bytes), + page_size(system_page_size), + data_page_count( + static_cast(kSlabBytes / system_page_size - 1)) {} + + uint64_t magic{kMagic}; + MemoryLimitPool *owner{nullptr}; + ReclaimableSlab *next{nullptr}; + void *reservation_base{nullptr}; + size_t reservation_size{0}; + size_t page_size{0}; + std::mutex mutex; + uint16_t next_unused{1}; + uint16_t reclaimed_count{0}; + uint16_t data_page_count{0}; + size_t committed_pages{0}; + uint16_t reclaimed_pages[kMaxSlabDataPages]{}; +}; + +size_t MemoryLimitPool::page_buffer_size() { + static const size_t page_size = []() -> size_t { +#if defined(_WIN32) + SYSTEM_INFO info; + ::GetSystemInfo(&info); + return static_cast(info.dwPageSize); +#else + return static_cast(::getpagesize()); +#endif + }(); + assert(page_size >= kMinimumPageBytes && page_size < kSlabBytes && + (page_size & (page_size - 1)) == 0); + return page_size; } bool BlockEvictionQueue::evict_single_block(BlockType &item) { - bool found = false; - for (size_t i = 0; i < CACHE_QUEUE_NUM; i++) { - found = evict_queues_[i].try_dequeue(item); - if (found) { - break; + return evict_single_block(item, /*age_protected=*/false); +} + +bool BlockEvictionQueue::evict_single_block(BlockType &item, + bool age_protected) { + if (age_protected) { + const size_t probation = approximate_queue_sizes_[kProbationPriority].load( + std::memory_order_relaxed); + const size_t protected_pages = + approximate_queue_sizes_[kProtectedPriority].load( + std::memory_order_relaxed); + const size_t dominance_threshold = + probation > (std::numeric_limits::max() - 1) / + kProtectedDominanceRatio + ? std::numeric_limits::max() + : (probation + 1) * kProtectedDominanceRatio; + if (protected_pages > dominance_threshold && + evict_queues_[kProtectedPriority].try_dequeue(item)) { + approximate_queue_sizes_[kProtectedPriority].fetch_sub( + 1, std::memory_order_relaxed); + protected_aging_dequeues_.fetch_add(1, std::memory_order_relaxed); + return true; } } - return found; -} -bool BlockEvictionQueue::is_valid_and_alive(const BlockType &item) { - std::shared_lock lock(valid_owners_mutex_); - if (item.owner == nullptr || - valid_owners_.find(item.owner) == valid_owners_.end()) { - return false; + for (size_t i = 0; i < kQueueCount; i++) { + if (evict_queues_[i].try_dequeue(item)) { + approximate_queue_sizes_[i].fetch_sub(1, std::memory_order_relaxed); + return true; + } } - return !item.owner->is_dead_block(item.owner_key, item.version); + return false; } bool BlockEvictionQueue::evict_block(BlockType &item) { - bool ok = false; - do { - ok = evict_single_block(item); - if (!ok) { + size_t attempts = 0; + bool age_protected = true; + return evict_block(item, attempts, std::numeric_limits::max(), + age_protected); +} + +bool BlockEvictionQueue::evict_block(BlockType &item, size_t &attempts, + size_t max_attempts, bool &age_protected) { + while (attempts < max_attempts) { + if (!evict_single_block(item, age_protected)) { return false; } - } while (!is_valid_and_alive(item)); - return ok; + // Protected aging is deliberately a once-per-reclaim-batch decision. + age_protected = false; + ++attempts; + std::shared_lock lock(valid_owners_mutex_); + if (item.owner == nullptr || + valid_owners_.find(item.owner) == valid_owners_.end() || + item.owner->is_dead_block(item.owner_key, item.version)) { + continue; + } + const uint8_t current_priority = + item.owner->eviction_priority(item.owner_key); + if (item.priority != current_priority) { + item.priority = current_priority; + if (!add_single_block(item, static_cast(current_priority))) { + item.owner->eviction_requeue_failed(item.owner_key, item.version); + } + continue; + } + return true; + } + return false; } void BlockEvictionQueue::recycle() { BlockType item; - while (MemoryLimitPool::get_instance().is_full() && evict_block(item)) { + // A foreground page fault must not scan the whole global queue while its + // page remains in kLoadingRefCount. Try a small CLOCK sample and let the + // caller fall back or retry; background reclaim handles deep queue walks. + static constexpr size_t kForegroundReclaimAttempts = 20; + const size_t max_attempts = kForegroundReclaimAttempts; + size_t attempts = 0; + bool recovered = false; + bool age_protected = true; + // Page allocations can hit their admission limit before the process-wide + // pool is full because part of the budget is reserved for external cache + // consumers. Stop immediately after enough room for one page is reclaimed. + while (MemoryLimitPool::get_instance().is_page_full() && + attempts < max_attempts) { + if (!evict_block(item, attempts, max_attempts, age_protected)) { + if (attempts >= max_attempts) { + break; + } + if (recovered || recover_owner_queues() == 0) { + break; + } + recovered = true; + continue; + } + { + std::shared_lock lock(valid_owners_mutex_); + if (item.owner != nullptr && + valid_owners_.find(item.owner) != valid_owners_.end() && + !item.owner->is_dead_block(item.owner_key, item.version)) { + item.owner->evict_block(item.owner_key); + } + } + } +} + +size_t BlockEvictionQueue::batch_recycle(size_t count) { + size_t evicted = 0; + // Bound work when no page is currently evictable. + const size_t max_attempts = + count > (std::numeric_limits::max() - 16) / 4 + ? std::numeric_limits::max() + : count * 4 + 16; + size_t attempts = 0; + bool recovered = false; + bool age_protected = count >= kProtectedAgingMinBatch; + while (evicted < count && attempts < max_attempts) { + BlockType item; + if (!evict_block(item, attempts, max_attempts, age_protected)) { + if (attempts >= max_attempts) { + break; + } + if (recovered || recover_owner_queues() == 0) { + break; + } + recovered = true; + continue; + } std::shared_lock lock(valid_owners_mutex_); if (item.owner != nullptr && - valid_owners_.find(item.owner) != valid_owners_.end()) { - item.owner->evict_block(item.owner_key); + valid_owners_.find(item.owner) != valid_owners_.end() && + !item.owner->is_dead_block(item.owner_key, item.version) && + item.owner->evict_block(item.owner_key)) { + ++evicted; } } + return evicted; +} + +size_t BlockEvictionQueue::recover_owner_queues() { + std::shared_lock lock(valid_owners_mutex_); + size_t recovered = 0; + for (EvictableBlockOwner *owner : valid_owners_) { + recovered += owner->recover_eviction_queue(); + } + return recovered; } bool BlockEvictionQueue::add_single_block(const BlockType &block, int queue_index) { - bool ok = evict_queues_[queue_index].enqueue(block); + if (queue_index < 0 || queue_index >= static_cast(kQueueCount)) { + LOG_ERROR("invalid eviction priority: %d", queue_index); + return false; + } + BlockType queued = block; + queued.priority = static_cast(queue_index); + // Publish the depth first so a concurrent consumer cannot dequeue the item + // before its accounting is visible. Roll it back if enqueue fails. + approximate_queue_sizes_[queue_index].fetch_add(1, std::memory_order_relaxed); + bool ok = evict_queues_[queue_index].enqueue(queued); if (!ok) { + approximate_queue_sizes_[queue_index].fetch_sub(1, + std::memory_order_relaxed); LOG_ERROR("enqueue failed."); return false; } return true; } +MemoryLimitPool::~MemoryLimitPool() { + stop_background_evictor(); + drain_free_list(); +} + +void MemoryLimitPool::drain_free_list() { + size_t released = 0; + for (size_t i = 0; i < kNumFreeShards; ++i) { + std::lock_guard lk(free_shards_[i].mutex); + released += free_shards_[i].count.load(std::memory_order_relaxed); + free_shards_[i].head = nullptr; + free_shards_[i].count.store(0, std::memory_order_relaxed); + } + release_all_slabs_locked(); + if (released != 0) { + LOG_INFO("MemoryLimitPool: released %zu cached slab pages", released); + } +} + +size_t MemoryLimitPool::pick_shard() { + // Keep each thread on one shard for locality. + thread_local size_t idx = shard_seq_.fetch_add(1, std::memory_order_relaxed); + return idx % kNumFreeShards; +} + +bool MemoryLimitPool::try_reserve_used(size_t bytes) { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + size_t used = used_size_.load(std::memory_order_relaxed); + while (used <= capacity && bytes <= capacity - used) { + if (used_size_.compare_exchange_weak(used, used + bytes, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + return true; + } + } + return false; +} + +bool MemoryLimitPool::try_reserve_page_used(size_t bytes) { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + const size_t reserve = page_admission_reserve(); + const size_t external = external_used_size_.load(std::memory_order_relaxed); + const size_t remaining_reserve = reserve > external ? reserve - external : 0; + const size_t page_limit = capacity - remaining_reserve; + size_t used = used_size_.load(std::memory_order_relaxed); + while (used <= page_limit && bytes <= page_limit - used) { + if (used_size_.compare_exchange_weak(used, used + bytes, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + return true; + } + } + return false; +} + +bool MemoryLimitPool::try_reserve_committed(size_t bytes) { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + size_t committed = committed_size_.load(std::memory_order_relaxed); + while (committed <= capacity && bytes <= capacity - committed) { + if (committed_size_.compare_exchange_weak(committed, committed + bytes, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + return true; + } + } + return false; +} + +bool MemoryLimitPool::is_cacheable_buffer_size(size_t buffer_size) { + return buffer_size == page_buffer_size(); +} + +char *MemoryLimitPool::pop_free_buffer(size_t start_shard) { + for (size_t i = 0; i < kNumFreeShards; ++i) { + size_t shard = (start_shard + i) % kNumFreeShards; + std::lock_guard lock(free_shards_[shard].mutex); + char *buffer = free_shards_[shard].head; + if (buffer) { + free_shards_[shard].head = *reinterpret_cast(buffer); + free_shards_[shard].count.fetch_sub(1, std::memory_order_relaxed); + return buffer; + } + } + return nullptr; +} + +void MemoryLimitPool::push_free_buffer(char *buffer, size_t shard) { + shard %= kNumFreeShards; + std::lock_guard lock(free_shards_[shard].mutex); + *reinterpret_cast(buffer) = free_shards_[shard].head; + free_shards_[shard].head = buffer; + free_shards_[shard].count.fetch_add(1, std::memory_order_relaxed); +} + +char *MemoryLimitPool::acquire_slab_buffer() { + static_assert(sizeof(ReclaimableSlab) <= kMinimumPageBytes, + "slab metadata must fit in its header page"); + + std::lock_guard slabs_lock(slab_mutex_); + auto acquire_from = [](ReclaimableSlab *slab) -> char * { + std::lock_guard slab_lock(slab->mutex); + uint16_t page_index = 0; + bool reclaimed = false; + if (slab->reclaimed_count != 0) { + reclaimed = true; + page_index = slab->reclaimed_pages[--slab->reclaimed_count]; + } else if (slab->next_unused <= slab->data_page_count) { + page_index = slab->next_unused++; + } else { + return nullptr; + } + + char *page = reinterpret_cast(slab) + + static_cast(page_index) * slab->page_size; + if (!commit_slab_page(page, slab->page_size, reclaimed)) { + if (reclaimed) { + slab->reclaimed_pages[slab->reclaimed_count++] = page_index; + } else { + --slab->next_unused; + } + return nullptr; + } + ++slab->committed_pages; + return page; + }; + + if (allocation_slab_ != nullptr) { + if (char *page = acquire_from(allocation_slab_)) { + return page; + } + } + for (ReclaimableSlab *slab = slabs_; slab != nullptr; slab = slab->next) { + if (slab == allocation_slab_) { + continue; + } + if (char *page = acquire_from(slab)) { + allocation_slab_ = slab; + return page; + } + } + + const size_t page_size = page_buffer_size(); + AlignedSlabMapping mapping = reserve_aligned_slab(page_size); + if (mapping.base == nullptr) { + LOG_ERROR("MemoryLimitPool: failed to reserve an aligned slab"); + return nullptr; + } + ReclaimableSlab *slab = new (mapping.base) ReclaimableSlab( + this, mapping.reservation_base, mapping.reservation_size, page_size); + slab->next = slabs_; + slabs_ = slab; + allocation_slab_ = slab; + slab_count_.fetch_add(1, std::memory_order_relaxed); + slab_mapped_bytes_.fetch_add(mapping.reservation_size, + std::memory_order_relaxed); + char *page = acquire_from(slab); + if (page != nullptr) { + return page; + } + + slabs_ = slab->next; + allocation_slab_ = nullptr; + slab_count_.fetch_sub(1, std::memory_order_relaxed); + slab_mapped_bytes_.fetch_sub(mapping.reservation_size, + std::memory_order_relaxed); + slab->~ReclaimableSlab(); + release_aligned_slab(mapping.reservation_base, mapping.reservation_size); + return nullptr; +} + +bool MemoryLimitPool::reclaim_slab_buffer(char *buffer) { + const uintptr_t address = reinterpret_cast(buffer); + const uintptr_t slab_address = address & ~(kSlabAlignment - 1); + auto *slab = reinterpret_cast(slab_address); + const size_t offset = address - slab_address; + if (slab->magic != ReclaimableSlab::kMagic || slab->owner != this || + offset < slab->page_size || offset >= kSlabBytes || + offset % slab->page_size != 0) { + LOG_ERROR("MemoryLimitPool: invalid page returned to slab allocator"); + return false; + } + + const auto page_index = static_cast(offset / slab->page_size); + std::lock_guard slab_lock(slab->mutex); + if (!discard_slab_page(buffer, slab->page_size)) { + LOG_ERROR("MemoryLimitPool: failed to discard a free slab page"); + return false; + } + + assert(slab->committed_pages != 0); + --slab->committed_pages; + size_t previous = + committed_size_.fetch_sub(slab->page_size, std::memory_order_relaxed); + (void)previous; + assert(previous >= slab->page_size); + assert(slab->reclaimed_count < slab->data_page_count); + slab->reclaimed_pages[slab->reclaimed_count++] = page_index; + slab_reclaimed_pages_.fetch_add(1, std::memory_order_relaxed); + return true; +} + +void MemoryLimitPool::release_all_slabs_locked() { + std::lock_guard slabs_lock(slab_mutex_); + size_t released_bytes = 0; + ReclaimableSlab *slab = slabs_; + while (slab != nullptr) { + ReclaimableSlab *next = slab->next; + released_bytes += slab->committed_pages * slab->page_size; + void *reservation_base = slab->reservation_base; + size_t reservation_size = slab->reservation_size; + slab->~ReclaimableSlab(); + release_aligned_slab(reservation_base, reservation_size); + slab = next; + } + slabs_ = nullptr; + allocation_slab_ = nullptr; + slab_count_.store(0, std::memory_order_relaxed); + slab_mapped_bytes_.store(0, std::memory_order_relaxed); + if (released_bytes != 0) { + size_t previous = + committed_size_.fetch_sub(released_bytes, std::memory_order_relaxed); + (void)previous; + assert(previous >= released_bytes); + } +} + +size_t MemoryLimitPool::trim_free_buffers(size_t bytes_needed) { + if (bytes_needed == 0) { + return 0; + } + + size_t released_bytes = 0; + size_t shard = pick_shard(); + while (released_bytes < bytes_needed) { + char *buffer = pop_free_buffer(shard); + if (!buffer) { + break; + } + // Only publish capacity after the physical page has been discarded. + if (!reclaim_slab_buffer(buffer)) { + push_free_buffer(buffer, shard); + break; + } + released_bytes += page_buffer_size(); + } + return released_bytes; +} + int MemoryLimitPool::init(size_t pool_size) { - pool_size_ = 0; + std::unique_lock lifecycle_lock(lifecycle_mutex_); + // Re-publishing the active process-wide budget is a safe no-op. + if (initialized_.load(std::memory_order_acquire) && + pool_size_.load(std::memory_order_relaxed) == pool_size) { + return 0; + } + const size_t used = used_size_.load(std::memory_order_relaxed); + const size_t external = external_used_size_.load(std::memory_order_relaxed); + const size_t metadata = metadata_used_size_.load(std::memory_order_relaxed); + if (used != 0 || external != 0 || metadata != 0) { + LOG_ERROR( + "MemoryLimitPool reinitialization rejected while cache memory is " + "active: requested_capacity=%zu current_capacity=%zu used=%zu " + "external_used=%zu metadata_used=%zu", + pool_size, pool_size_.load(std::memory_order_relaxed), used, external, + metadata); + return -1; + } + + // Tear down the background evictor first: it reads pool_size_ and touches + // the free-list, both of which we are about to reset. + stop_background_evictor(); + pool_size_.store(0, std::memory_order_relaxed); BlockEvictionQueue::get_instance().recycle(); - pool_size_ = pool_size; - LOG_INFO("MemoryLimitPool initialized with pool size: %lu", pool_size_); + drain_free_list(); + pool_size_.store(pool_size, std::memory_order_relaxed); + initialized_.store(true, std::memory_order_release); + LOG_INFO("Shared cache initialized with capacity: %zu", pool_size); + if (pool_size > 0) { + try { + start_background_evictor(); + } catch (const std::exception &e) { + pool_size_.store(0, std::memory_order_relaxed); + initialized_.store(false, std::memory_order_release); + LOG_ERROR("Failed to start the shared-cache evictor: %s", e.what()); + return -1; + } catch (...) { + pool_size_.store(0, std::memory_order_relaxed); + initialized_.store(false, std::memory_order_release); + LOG_ERROR( + "Failed to start the shared-cache evictor with an unknown error"); + return -1; + } + } return 0; } +void MemoryLimitPool::start_background_evictor() { + bool expected = false; + if (!bg_running_.compare_exchange_strong(expected, true)) { + return; // already running + } + try { + bg_thread_ = std::thread([this] { background_evict_loop(); }); + } catch (...) { + bg_running_.store(false, std::memory_order_release); + throw; + } +} + +void MemoryLimitPool::stop_background_evictor() { + if (!bg_running_.exchange(false)) { + return; // not running + } + { std::lock_guard lk(bg_mutex_); } + bg_cv_.notify_all(); + if (bg_thread_.joinable()) { + bg_thread_.join(); + } +} + +void MemoryLimitPool::background_evict_loop() { + using std::chrono::milliseconds; + while (bg_running_.load()) { + { + std::unique_lock lk(bg_mutex_); + bg_cv_.wait_for(lk, milliseconds(5), [this] { + return !bg_running_.load() || should_background_reclaim(); + }); + } + if (!bg_running_.load()) break; + if (pool_size_.load(std::memory_order_relaxed) == 0) continue; + const size_t low = low_watermark(); + if (used_size_.load() > low) { + bg_evict_rounds_.fetch_add(1, std::memory_order_relaxed); + } + // Reclaim proactively down to the low watermark so the foreground path + // finds ready buffers on the free-list instead of evicting inline. + while (bg_running_.load() && used_size_.load() > low) { + size_t n = BlockEvictionQueue::get_instance().batch_recycle(64); + if (n == 0) { + // Back off when pressure remains but eviction makes no progress. + bg_no_progress_sleeps_.fetch_add(1, std::memory_order_relaxed); + std::unique_lock lk(bg_mutex_); + bg_cv_.wait_for(lk, milliseconds(5), + [this] { return !bg_running_.load(); }); + break; + } + bg_evicted_buffers_.fetch_add(n, std::memory_order_relaxed); + } + } +} + bool MemoryLimitPool::try_acquire_buffer(const size_t buffer_size, char *&buffer) { - size_t expected, desired; - do { - expected = used_size_.load(); - if (expected >= pool_size_) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + buffer = nullptr; + const bool cacheable = is_cacheable_buffer_size(buffer_size); + if (buffer_size == 0 || !(cacheable ? try_reserve_page_used(buffer_size) + : try_reserve_used(buffer_size))) { + // Out of budget: wake the background evictor so the next attempt is + // more likely to find a free buffer without inline eviction. + high_watermark_hits_.fetch_add(1, std::memory_order_relaxed); + bg_cv_.notify_one(); + return false; + } + + if (cacheable) { + buffer = pop_free_buffer(pick_shard()); + if (buffer) { + alloc_from_freelist_.fetch_add(1, std::memory_order_relaxed); + return true; + } + } + + if (!try_reserve_committed(buffer_size)) { + // This is primarily useful for a non-cacheable size. For the normal page + // size, all currently visible free buffers were already checked above. + trim_free_buffers(buffer_size); + if (!try_reserve_committed(buffer_size)) { + used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + high_watermark_hits_.fetch_add(1, std::memory_order_relaxed); return false; } - desired = expected + buffer_size; - } while (!used_size_.compare_exchange_weak(expected, desired)); - buffer = (char *)ailego_aligned_malloc(buffer_size, 4096); + } + buffer = cacheable ? acquire_slab_buffer() + : static_cast(ailego_aligned_malloc( + buffer_size, kMinimumPageBytes)); if (!buffer) { - used_size_.fetch_sub(buffer_size); + committed_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); return false; } + alloc_from_slab_.fetch_add(1, std::memory_order_relaxed); return true; } -void MemoryLimitPool::charge_external(const size_t buffer_size) { - size_t expected, desired; - do { - expected = used_size_.load(); - desired = expected + buffer_size; - } while (!used_size_.compare_exchange_weak(expected, desired)); +bool MemoryLimitPool::wait_for_available(const size_t buffer_size, + std::chrono::milliseconds timeout) { + if (buffer_size == 0) { + return true; + } + capacity_waits_.fetch_add(1, std::memory_order_relaxed); + std::unique_lock lock(capacity_mutex_); + const bool available = + capacity_cv_.wait_for(lock, timeout, [this, buffer_size] { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + const size_t used = used_size_.load(std::memory_order_relaxed); + return capacity >= used && buffer_size <= capacity - used; + }); + if (!available) { + capacity_wait_timeouts_.fetch_add(1, std::memory_order_relaxed); + } + return available; +} + +bool MemoryLimitPool::try_charge_external(const size_t buffer_size) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + return try_charge_fixed(buffer_size, &external_used_size_); +} + +bool MemoryLimitPool::try_charge_metadata(const size_t buffer_size) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + return try_charge_fixed(buffer_size, &metadata_used_size_); +} + +bool MemoryLimitPool::try_charge_fixed(const size_t buffer_size, + std::atomic *counter) { + if (buffer_size == 0) { + return true; + } + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + if (capacity == 0 || buffer_size > capacity) { + high_watermark_hits_.fetch_add(1, std::memory_order_relaxed); + return false; + } + + while (true) { + if (try_reserve_committed(buffer_size)) { + counter->fetch_add(buffer_size, std::memory_order_relaxed); + used_size_.fetch_add(buffer_size, std::memory_order_relaxed); + bg_cv_.notify_one(); + return true; + } + + size_t committed = committed_size_.load(std::memory_order_relaxed); + size_t available = committed >= capacity ? 0 : capacity - committed; + if (available >= buffer_size) { + continue; + } + if (trim_free_buffers(buffer_size - available) != 0) { + continue; + } + + // Reclaim until the reservation fits or eviction stops making progress. + if (BlockEvictionQueue::get_instance().batch_recycle(256) == 0) { + high_watermark_hits_.fetch_add(1, std::memory_order_relaxed); + return false; + } + } } void MemoryLimitPool::release_buffer(char *buffer, const size_t buffer_size) { - size_t expected, desired; - do { - expected = used_size_.load(); - desired = expected - buffer_size; - assert(expected >= buffer_size); - } while (!used_size_.compare_exchange_weak(expected, desired)); - ailego_free(buffer); + if (!buffer) { + size_t prev = used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)prev; + assert(prev >= buffer_size); + { std::lock_guard lock(capacity_mutex_); } + capacity_cv_.notify_one(); + return; + } + if (!is_cacheable_buffer_size(buffer_size)) { + ailego_free(buffer); + size_t committed_prev = + committed_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)committed_prev; + assert(committed_prev >= buffer_size); + size_t prev = used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)prev; + assert(prev >= buffer_size); + { std::lock_guard lock(capacity_mutex_); } + capacity_cv_.notify_one(); + return; + } + push_free_buffer(buffer, pick_shard()); + // Publish the free buffer before releasing its logical budget slot. + size_t prev = used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)prev; + assert(prev >= buffer_size); + { std::lock_guard lock(capacity_mutex_); } + capacity_cv_.notify_one(); } void MemoryLimitPool::release_external(const size_t buffer_size) { + release_fixed(buffer_size, &external_used_size_); +} + +void MemoryLimitPool::release_metadata(const size_t buffer_size) { + release_fixed(buffer_size, &metadata_used_size_); +} + +void MemoryLimitPool::release_fixed(const size_t buffer_size, + std::atomic *counter) { + if (buffer_size == 0) { + return; + } + // Unconditional subtract: single RMW instead of a CAS loop. + size_t prev = used_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)prev; + assert(prev >= buffer_size); + size_t counter_prev = + counter->fetch_sub(buffer_size, std::memory_order_relaxed); + (void)counter_prev; + assert(counter_prev >= buffer_size); + size_t committed_prev = + committed_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + (void)committed_prev; + assert(committed_prev >= buffer_size); + { std::lock_guard lock(capacity_mutex_); } + capacity_cv_.notify_all(); +} + +bool MemoryLimitPool::is_full() { + return used_size_.load(std::memory_order_relaxed) >= + pool_size_.load(std::memory_order_relaxed); +} + +size_t MemoryLimitPool::batch_acquire_buffers(size_t buffer_size, char **out, + size_t count) { + std::shared_lock lifecycle_lock(lifecycle_mutex_); + if (count == 0 || buffer_size == 0) return 0; + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + const bool cacheable = is_cacheable_buffer_size(buffer_size); + const size_t reserve = cacheable ? page_admission_reserve() : 0; + const size_t external = external_used_size_.load(std::memory_order_relaxed); + const size_t remaining_reserve = reserve > external ? reserve - external : 0; + const size_t admission_limit = capacity - remaining_reserve; + size_t total_size = 0; + size_t actual_count = count; size_t expected, desired; do { - expected = used_size_.load(); - desired = expected - buffer_size; - assert(expected >= buffer_size); - } while (!used_size_.compare_exchange_weak(expected, desired)); + expected = used_size_.load(std::memory_order_relaxed); + if (expected >= admission_limit) return 0; + size_t avail = (admission_limit - expected) / buffer_size; + if (avail == 0) return 0; + if (avail < actual_count) actual_count = avail; + total_size = actual_count * buffer_size; + desired = expected + total_size; + } while (!used_size_.compare_exchange_weak( + expected, desired, std::memory_order_relaxed, std::memory_order_relaxed)); + + size_t acquired = 0; + size_t s = pick_shard(); + if (cacheable) { + while (acquired < actual_count) { + out[acquired] = pop_free_buffer(s); + if (!out[acquired]) { + break; + } + ++acquired; + } + alloc_from_freelist_.fetch_add(acquired, std::memory_order_relaxed); + } + + while (acquired < actual_count) { + if (!try_reserve_committed(buffer_size)) { + trim_free_buffers(buffer_size); + if (!try_reserve_committed(buffer_size)) { + break; + } + } + out[acquired] = cacheable ? acquire_slab_buffer() + : static_cast(ailego_aligned_malloc( + buffer_size, kMinimumPageBytes)); + if (!out[acquired]) { + committed_size_.fetch_sub(buffer_size, std::memory_order_relaxed); + break; + } + alloc_from_slab_.fetch_add(1, std::memory_order_relaxed); + ++acquired; + } + if (acquired < actual_count) { + used_size_.fetch_sub((actual_count - acquired) * buffer_size, + std::memory_order_relaxed); + } + return acquired; } -bool MemoryLimitPool::is_full() { - return used_size_.load() >= pool_size_; +MemoryLimitPool::PoolStats MemoryLimitPool::stats() const { + PoolStats s; + s.pool_size = pool_size_.load(std::memory_order_relaxed); + s.used = used_size_.load(std::memory_order_relaxed); + s.committed = committed_size_.load(std::memory_order_relaxed); + s.external_used = external_used_size_.load(std::memory_order_relaxed); + s.metadata_used = metadata_used_size_.load(std::memory_order_relaxed); + const size_t fixed = s.external_used + s.metadata_used; + s.page_used = s.used >= fixed ? s.used - fixed : 0; + size_t free_buffers = 0; + for (size_t i = 0; i < kNumFreeShards; ++i) { + free_buffers += free_shards_[i].count.load(std::memory_order_relaxed); + } + s.free_buffers = free_buffers; + s.slab_count = slab_count_.load(std::memory_order_relaxed); + s.slab_mapped_bytes = slab_mapped_bytes_.load(std::memory_order_relaxed); + s.slab_header_bytes = s.slab_count * page_buffer_size(); + s.alloc_from_freelist = alloc_from_freelist_.load(std::memory_order_relaxed); + s.alloc_from_slab = alloc_from_slab_.load(std::memory_order_relaxed); + s.slab_reclaimed_pages = + slab_reclaimed_pages_.load(std::memory_order_relaxed); + s.bg_evict_rounds = bg_evict_rounds_.load(std::memory_order_relaxed); + s.bg_evicted_buffers = bg_evicted_buffers_.load(std::memory_order_relaxed); + s.bg_no_progress_sleeps = + bg_no_progress_sleeps_.load(std::memory_order_relaxed); + s.high_watermark_hits = high_watermark_hits_.load(std::memory_order_relaxed); + s.capacity_waits = capacity_waits_.load(std::memory_order_relaxed); + s.capacity_wait_timeouts = + capacity_wait_timeouts_.load(std::memory_order_relaxed); + return s; +} + +void MemoryLimitPool::log_stats() const { + PoolStats s = stats(); + LOG_INFO( + "Shared cache stats: capacity=%llu used=%llu committed=%llu " + "page_used=%llu external_used=%llu metadata_used=%llu " + "free_buffers=%llu slab_count=%llu slab_mapped_bytes=%llu " + "slab_header_bytes=%llu " + "alloc_from_freelist=%llu alloc_from_slab=%llu " + "slab_reclaimed_pages=%llu " + "bg_evict_rounds=%llu bg_evicted_buffers=%llu " + "bg_no_progress_sleeps=%llu " + "high_watermark_hits=%llu capacity_waits=%llu " + "capacity_wait_timeouts=%llu", + static_cast(s.pool_size), + static_cast(s.used), + static_cast(s.committed), + static_cast(s.page_used), + static_cast(s.external_used), + static_cast(s.metadata_used), + static_cast(s.free_buffers), + static_cast(s.slab_count), + static_cast(s.slab_mapped_bytes), + static_cast(s.slab_header_bytes), + static_cast(s.alloc_from_freelist), + static_cast(s.alloc_from_slab), + static_cast(s.slab_reclaimed_pages), + static_cast(s.bg_evict_rounds), + static_cast(s.bg_evicted_buffers), + static_cast(s.bg_no_progress_sleeps), + static_cast(s.high_watermark_hits), + static_cast(s.capacity_waits), + static_cast(s.capacity_wait_timeouts)); } } // namespace ailego diff --git a/src/ailego/buffer/vector_page_table.cc b/src/ailego/buffer/vector_page_table.cc index c875c32eb..b9d9e747c 100644 --- a/src/ailego/buffer/vector_page_table.cc +++ b/src/ailego/buffer/vector_page_table.cc @@ -12,7 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include +#include +#include #include #include #include @@ -21,7 +25,13 @@ #include #include +#if defined(__linux__) +#include +#include +#endif + #if defined(_MSC_VER) +#include #ifndef NOMINMAX #define NOMINMAX #endif @@ -68,100 +78,546 @@ namespace ailego { const size_t kVectorPageSize = MemoryHelper::PageSize(); +namespace { +constexpr size_t kBlockingAioBatchSize = VecBufferPool::kWritebackBatchPages; +} + +VecBufferPool::~VecBufferPool() { + // Finish queued writes before page buffers, latches, and descriptors go + // away. The synchronous pass below catches any prior writeback error. + stop_writeback(); + // Flush dirty pages before releasing memory and descriptors. + (void)this->flush_all(); + // Preserve final cache and writeback statistics after both persistence + // paths have drained. + log_stats(); + page_table_.force_evict_all_loaded(); + const size_t writable_metadata_bytes = block_mutex_metadata_bytes() + + writeback_staging_size_ + + writeback_io_staging_charge_; + block_mutexes_.reset(); + if (writeback_staging_ != nullptr) { + ailego_free(writeback_staging_); + writeback_staging_ = nullptr; + } +#if defined(__linux__) + writeback_io_uring_.reset(); +#endif + MemoryLimitPool::get_instance().release_metadata(writable_metadata_bytes); + block_mutex_count_ = 0; + writeback_staging_size_ = 0; + writeback_io_staging_charge_ = 0; + initialized_ = false; +#if defined(_MSC_VER) + _close(fd_); + _close(meta_fd_); +#else + close(fd_); + close(meta_fd_); +#endif +} + +version_t VectorPageTable::next_owner_version() { + return BlockEvictionQueue::get_instance().next_version(); +} + +size_t VectorPageTable::metadata_bytes_for_entries(size_t entry_num) { + if (entry_num > kMaxEntries) { + return std::numeric_limits::max(); + } + const size_t segment_count = + entry_num == 0 ? 0 : (entry_num - 1) / kSegmentSize + 1; + if (segment_count == 0) { + return 0; + } + if (segment_count > + (std::numeric_limits::max() - kSegmentDirectoryBytes) / + kSegmentMetadataBytes) { + return std::numeric_limits::max(); + } + return kSegmentDirectoryBytes + segment_count * kSegmentMetadataBytes; +} + +void VectorPageTable::initialize_segment(Entry *entries, + MetadataEntry *metadata_entries) { + for (size_t i = 0; i < kSegmentSize; ++i) { + entries[i].buffer.store(nullptr, std::memory_order_relaxed); + entries[i].ref_count.store(kUnloadedRefCount, std::memory_order_relaxed); + entries[i].in_evict_queue.store(false, std::memory_order_relaxed); + entries[i].referenced.store(false, std::memory_order_relaxed); + entries[i].evict_priority.store(0, std::memory_order_relaxed); + entries[i].ghost_state.store(kNoGhostHistory, std::memory_order_relaxed); + metadata_entries[i].next_loaded = kInvalidLoadedBlock; + metadata_entries[i].file_offset = 0; + metadata_entries[i].admission_state.store(0, std::memory_order_relaxed); + metadata_entries[i].is_dirty.store(false, std::memory_order_relaxed); + metadata_entries[i].writeback_pending.store(false, + std::memory_order_relaxed); + metadata_entries[i].ever_loaded.store(false, std::memory_order_relaxed); + } +} + +bool VectorPageTable::should_admit_miss(block_id_t block_id, uint32_t epoch) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &entry = entry_at(block_id); + + // Joining an existing residency transition preserves single-flight. A + // protected hint or hot ghost is also stronger evidence than miss count. + if (entry.ref_count.load(std::memory_order_acquire) != kUnloadedRefCount || + entry.evict_priority.load(std::memory_order_relaxed) >= kNormalPriority || + entry.ghost_state.load(std::memory_order_relaxed) == kEvictedHot) { + return true; + } + + MetadataEntry &metadata = metadata_entry_at(block_id); + static constexpr uint32_t kEpochMask = (uint32_t{1} << 24) - 1; + static constexpr uint8_t kAdmissionThreshold = 2; + epoch &= kEpochMask; + + uint32_t state = metadata.admission_state.load(std::memory_order_relaxed); + while (true) { + const uint32_t previous_epoch = state >> 8; + const uint8_t previous_count = static_cast(state); + const uint32_t age = (epoch - previous_epoch) & kEpochMask; + + uint8_t count = 1; + if (age == 0) { + count = previous_count == std::numeric_limits::max() + ? previous_count + : static_cast(previous_count + 1); + } else if (age == 1) { + count = static_cast(previous_count / 2 + 1); + } + const uint32_t updated = (epoch << 8) | count; + if (metadata.admission_state.compare_exchange_weak( + state, updated, std::memory_order_relaxed, + std::memory_order_relaxed)) { + // Close the race with a loader that claimed the page while its miss was + // being recorded; waiting for that load is preferable to duplicate I/O. + return count >= kAdmissionThreshold || + entry.ref_count.load(std::memory_order_acquire) != + kUnloadedRefCount; + } + } +} + bool VectorPageTable::init(size_t entry_num) { - size_t need_segments = (entry_num + kSegmentSize - 1) / kSegmentSize; - if (need_segments > kMaxSegments) { + if (entry_num > kMaxEntries) { LOG_ERROR( "VectorPageTable::init: entry_num=%zu exceeds capacity " - "(kMaxEntries=%zu, need_segments=%zu, kMaxSegments=%zu); " + "(kMaxEntries=%zu, kMaxSegments=%zu); " "refusing to init.", - entry_num, kMaxEntries, need_segments, kMaxSegments); + entry_num, kMaxEntries, kMaxSegments); + return false; + } + const size_t old_entry_num = entry_num_.load(std::memory_order_relaxed); + const size_t old_count = segment_count_.load(std::memory_order_relaxed); + if (old_count != 0) { + if (old_entry_num == entry_num) { + return true; + } + LOG_ERROR( + "VectorPageTable::init: refusing to replace an initialized table " + "(old_entries=%zu, requested_entries=%zu)", + old_entry_num, entry_num); + return false; + } + const size_t need_segments = + entry_num == 0 ? 0 : (entry_num - 1) / kSegmentSize + 1; + const size_t charge = metadata_bytes_for_entries(entry_num); + if (charge == std::numeric_limits::max()) { + LOG_ERROR("VectorPageTable::init: metadata size overflow for %zu entries", + entry_num); + return false; + } + if (!MemoryLimitPool::get_instance().try_charge_metadata(charge)) { + LOG_ERROR( + "VectorPageTable::init: shared memory budget cannot reserve %zu " + "metadata bytes for %zu entries", + charge, entry_num); return false; } - // Free old segments if any. init() is only called from VecBufferPool::init - // which is single-threaded with respect to other accesses, so a relaxed - // load of segment_count_ is sufficient here. - size_t old_count = segment_count_.load(std::memory_order_relaxed); - for (size_t i = 0; i < old_count; ++i) { - delete[] segments_[i]; - segments_[i] = nullptr; + + std::vector> new_segments; + std::vector> new_metadata_segments; + std::unique_ptr new_segment_directory; + std::unique_ptr new_metadata_segment_directory; + try { + if (need_segments != 0) { + new_segment_directory = std::make_unique(kMaxSegments); + new_metadata_segment_directory = + std::make_unique(kMaxSegments); + } + new_segments.reserve(need_segments); + new_metadata_segments.reserve(need_segments); + for (size_t s = 0; s < need_segments; ++s) { + auto entries = std::make_unique(kSegmentSize); + auto metadata_entries = std::make_unique(kSegmentSize); + initialize_segment(entries.get(), metadata_entries.get()); + new_segments.push_back(std::move(entries)); + new_metadata_segments.push_back(std::move(metadata_entries)); + } + } catch (const std::bad_alloc &) { + MemoryLimitPool::get_instance().release_metadata(charge); + LOG_ERROR( + "VectorPageTable::init: allocation failed for %zu entries (%zu " + "metadata bytes)", + entry_num, charge); + return false; } for (size_t s = 0; s < need_segments; ++s) { - segments_[s] = new Entry[kSegmentSize]; - for (size_t i = 0; i < kSegmentSize; ++i) { - segments_[s][i].ref_count.store(std::numeric_limits::min()); - segments_[s][i].in_evict_queue.store(false); - segments_[s][i].is_dirty.store(false); - segments_[s][i].buffer = nullptr; - segments_[s][i].file_offset = 0; - } - } - // Publish new segments to readers. segment_count_ is published first - // (release) so that a reader that acquire-loads segment_count_ before - // entry_num_ also sees a consistent segment table; entry_num_ is the - // primary synchronization point used by callers via entry_num(). + new_segment_directory[s] = new_segments[s].release(); + new_metadata_segment_directory[s] = new_metadata_segments[s].release(); + } + segments_ = std::move(new_segment_directory); + metadata_segments_ = std::move(new_metadata_segment_directory); + // Publish segments before the externally visible entry count. segment_count_.store(need_segments, std::memory_order_release); entry_num_.store(entry_num, std::memory_order_release); return true; } bool VectorPageTable::extend(size_t new_entry_num) { - // Relaxed read is fine: extend() is serialized by the caller (extend_file - // is invoked under the BufferStorage write latch). No other writer races - // with us on entry_num_ / segment_count_. + // The caller serializes page-table extension. if (new_entry_num <= entry_num_.load(std::memory_order_relaxed)) { return true; } - size_t new_segment_count = (new_entry_num + kSegmentSize - 1) / kSegmentSize; - if (new_segment_count > kMaxSegments) { + if (new_entry_num > kMaxEntries) { LOG_ERROR( "VectorPageTable::extend: new_entry_num=%zu exceeds capacity " - "(kMaxEntries=%zu, new_segment_count=%zu, kMaxSegments=%zu); " + "(kMaxEntries=%zu, kMaxSegments=%zu); " "refusing to extend.", - new_entry_num, kMaxEntries, new_segment_count, kMaxSegments); + new_entry_num, kMaxEntries, kMaxSegments); + return false; + } + const size_t new_segment_count = + new_entry_num == 0 ? 0 : (new_entry_num - 1) / kSegmentSize + 1; + const size_t old_count = segment_count_.load(std::memory_order_relaxed); + const size_t added_segments = new_segment_count - old_count; + const bool needs_directory = old_count == 0 && new_segment_count != 0; + if (added_segments > (std::numeric_limits::max() - + (needs_directory ? kSegmentDirectoryBytes : 0)) / + kSegmentMetadataBytes) { + LOG_ERROR( + "VectorPageTable::extend: metadata size overflow for %zu new " + "segments", + added_segments); + return false; + } + const size_t added_charge = added_segments * kSegmentMetadataBytes + + (needs_directory ? kSegmentDirectoryBytes : 0); + if (!MemoryLimitPool::get_instance().try_charge_metadata(added_charge)) { + LOG_ERROR( + "VectorPageTable::extend: shared memory budget cannot reserve %zu " + "additional metadata bytes (old_entries=%zu, new_entries=%zu)", + added_charge, entry_num_.load(std::memory_order_relaxed), + new_entry_num); return false; } - size_t old_count = segment_count_.load(std::memory_order_relaxed); + + std::vector> new_segments; + std::vector> new_metadata_segments; + std::unique_ptr new_segment_directory; + std::unique_ptr new_metadata_segment_directory; + try { + if (needs_directory) { + new_segment_directory = std::make_unique(kMaxSegments); + new_metadata_segment_directory = + std::make_unique(kMaxSegments); + } + new_segments.reserve(new_segment_count - old_count); + new_metadata_segments.reserve(new_segment_count - old_count); + for (size_t s = old_count; s < new_segment_count; ++s) { + auto entries = std::make_unique(kSegmentSize); + auto metadata_entries = std::make_unique(kSegmentSize); + initialize_segment(entries.get(), metadata_entries.get()); + new_segments.push_back(std::move(entries)); + new_metadata_segments.push_back(std::move(metadata_entries)); + } + } catch (const std::bad_alloc &) { + MemoryLimitPool::get_instance().release_metadata(added_charge); + LOG_ERROR( + "VectorPageTable::extend: allocation failed for %zu new entries " + "(%zu additional metadata bytes)", + new_entry_num, added_charge); + return false; + } + Entry **segment_directory = + needs_directory ? new_segment_directory.get() : segments_.get(); + MetadataEntry **metadata_segment_directory = + needs_directory ? new_metadata_segment_directory.get() + : metadata_segments_.get(); for (size_t s = old_count; s < new_segment_count; ++s) { - segments_[s] = new Entry[kSegmentSize]; - for (size_t i = 0; i < kSegmentSize; ++i) { - segments_[s][i].ref_count.store(std::numeric_limits::min()); - segments_[s][i].in_evict_queue.store(false); - segments_[s][i].is_dirty.store(false); - segments_[s][i].buffer = nullptr; - segments_[s][i].file_offset = 0; - } - } - // Publish in the same order as init(): segment_count_ first, entry_num_ - // last. Both are release-stores so that the prior segment allocation / - // Entry initialization is visible to any reader that acquire-loads either - // counter (typically via entry_num()). + const size_t idx = s - old_count; + segment_directory[s] = new_segments[idx].release(); + metadata_segment_directory[s] = new_metadata_segments[idx].release(); + } + if (needs_directory) { + segments_ = std::move(new_segment_directory); + metadata_segments_ = std::move(new_metadata_segment_directory); + } + // Match init() publication order: segments first, entry count last. segment_count_.store(new_segment_count, std::memory_order_release); entry_num_.store(new_entry_num, std::memory_order_release); return true; } -char *VectorPageTable::acquire_block(block_id_t block_id) { +bool VectorPageTable::rollback_extend(size_t old_entry_num) { + const size_t current_entry_num = entry_num_.load(std::memory_order_relaxed); + if (old_entry_num > current_entry_num) { + return false; + } + if (old_entry_num == current_entry_num) { + return true; + } + for (size_t i = old_entry_num; i < current_entry_num; ++i) { + if (entry_at(i).buffer.load(std::memory_order_relaxed) != nullptr || + entry_at(i).ref_count.load(std::memory_order_relaxed) != + std::numeric_limits::min() || + metadata_entry_at(i).ever_loaded.load(std::memory_order_relaxed)) { + LOG_ERROR( + "VectorPageTable::rollback_extend: new entry %zu is already in " + "use; refusing rollback", + i); + return false; + } + } + + const size_t old_segment_count = + old_entry_num == 0 ? 0 : (old_entry_num - 1) / kSegmentSize + 1; + const size_t current_segment_count = + segment_count_.load(std::memory_order_relaxed); + entry_num_.store(old_entry_num, std::memory_order_release); + segment_count_.store(old_segment_count, std::memory_order_release); + for (size_t s = old_segment_count; s < current_segment_count; ++s) { + delete[] segments_[s]; + segments_[s] = nullptr; + delete[] metadata_segments_[s]; + metadata_segments_[s] = nullptr; + } + size_t released_charge = + (current_segment_count - old_segment_count) * kSegmentMetadataBytes; + if (old_segment_count == 0) { + segments_.reset(); + metadata_segments_.reset(); + released_charge += kSegmentDirectoryBytes; + } + MemoryLimitPool::get_instance().release_metadata(released_charge); + return true; +} + +char *VectorPageTable::acquire_block(block_id_t block_id, bool record_reuse) { assert(block_id < entry_num_.load(std::memory_order_relaxed)); Entry &e = entry_at(block_id); + // Pin only resident pages; negative values are transition sentinels. + int count = e.ref_count.load(std::memory_order_acquire); + while (ailego_likely(count >= 0)) { + if (e.ref_count.compare_exchange_weak(count, count + 1, + std::memory_order_acquire, + std::memory_order_relaxed)) { + if (record_reuse) { + const uint32_t sample = next_hit_sample(); + // Reuse policy is approximate: a genuinely hot page is sampled + // quickly, while the common hit path avoids repeated atomic metadata + // updates. Also stop policy work after pressure has subsided. + if ((sample & (kReusePolicySampleRate - 1)) == 0 && + adaptive_priority_enabled_ && + has_evicted_.load(std::memory_order_relaxed)) { + const uint8_t ghost_state = + e.ghost_state.load(std::memory_order_relaxed); + const uint8_t priority = + e.evict_priority.load(std::memory_order_relaxed); + // Most HNSW hits are already protected by the one-time hot-set hint. + // Avoid global pressure checks and no-op promotion attempts for + // those pages. A ghost-admitted protected page remains eligible so + // one sampled reuse can validate its renewed hot history. + const bool needs_policy_update = + ghost_state == kGhostAdmitted || priority < kNormalPriority; + if (needs_policy_update && + MemoryLimitPool::get_instance().under_cache_pressure()) { + // A sampled reuse after ghost admission validates that the page is + // still hot. Its next protected residency may leave another ghost. + if (ghost_state == kGhostAdmitted) { + uint8_t ghost_admitted = kGhostAdmitted; + (void)e.ghost_state.compare_exchange_strong( + ghost_admitted, kNoGhostHistory, std::memory_order_relaxed, + std::memory_order_relaxed); + } + if (priority < kNormalPriority) { + (void)promote_evict_priority(block_id, kNormalPriority); + } + if (!e.referenced.load(std::memory_order_relaxed)) { + e.referenced.store(true, std::memory_order_relaxed); + } + } + } + // Sample the observability counter and CLOCK reference bit together. + if ((sample & (kHitSampleRate - 1)) == 0) { + if (!e.referenced.load(std::memory_order_relaxed)) { + e.referenced.store(true, std::memory_order_relaxed); + } + inc_sampled_hit(); + } + } + return e.buffer.load(std::memory_order_acquire); + } + } + return nullptr; +} + +VectorPageTable::LoadClaimResult VectorPageTable::try_claim_block_load( + block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &entry = entry_at(block_id); + int state = entry.ref_count.load(std::memory_order_acquire); while (true) { - int current_count = e.ref_count.load(std::memory_order_acquire); - if (current_count < 0) { - return nullptr; + if (state >= 0) { + return LoadClaimResult::kResident; + } + if (state == kLoadingRefCount) { + return LoadClaimResult::kLoading; + } + if (state != kUnloadedRefCount) { + return LoadClaimResult::kEvicting; + } + if (entry.ref_count.compare_exchange_weak(state, kLoadingRefCount, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return LoadClaimResult::kClaimed; + } + } +} + +bool VectorPageTable::wait_for_block_transition(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + const Entry &entry = entry_at(block_id); + using clock = std::chrono::steady_clock; + const auto wait_start = clock::now(); + auto last_log = wait_start; + unsigned spin_count = 0; + bool warned = false; + static constexpr auto kHardTimeout = std::chrono::seconds(30); + while (true) { + const int state = entry.ref_count.load(std::memory_order_acquire); + if (state != kLoadingRefCount && state != kEvictingRefCount) { + return true; + } + + ++spin_count; + if (spin_count < 64) { + } else if (spin_count < 1024) { + std::this_thread::yield(); + } else if (spin_count < 8192) { + std::this_thread::sleep_for(std::chrono::microseconds(100)); + } else { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); } - if (e.ref_count.compare_exchange_weak(current_count, current_count + 1, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - return e.buffer; + + const auto now = clock::now(); + const auto elapsed = now - wait_start; + if (!warned && elapsed >= std::chrono::milliseconds(100)) { + LOG_WARN( + "wait_for_block_transition: long wait on block_id=%zu state=%d " + "(>=100ms)", + static_cast(block_id), state); + warned = true; + } + if (elapsed >= kHardTimeout) { + LOG_ERROR( + "wait_for_block_transition: hard timeout (%lld s) on block_id=%zu " + "state=%d", + static_cast( + std::chrono::duration_cast(elapsed) + .count()), + static_cast(block_id), state); + return false; + } + if (elapsed >= std::chrono::seconds(1) && + (now - last_log) >= std::chrono::seconds(1)) { + const auto secs = + std::chrono::duration_cast(elapsed).count(); + LOG_ERROR( + "wait_for_block_transition: block_id=%zu state=%d still busy after " + "%lld s", + static_cast(block_id), state, static_cast(secs)); + last_log = now; + } + } +} + +char *VectorPageTable::publish_claimed_block(block_id_t block_id, char *buffer, + size_t file_offset) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + assert(buffer != nullptr); + Entry &entry = entry_at(block_id); + MetadataEntry &metadata = metadata_entry_at(block_id); + if (entry.ref_count.load(std::memory_order_acquire) != kLoadingRefCount) { + LOG_ERROR( + "publish_claimed_block: block_id=%zu is not owned by a loader, " + "state=%d", + static_cast(block_id), + entry.ref_count.load(std::memory_order_relaxed)); + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + return nullptr; + } + + metadata.file_offset = file_offset; + metadata.is_dirty.store(false, std::memory_order_relaxed); + entry.referenced.store(false, std::memory_order_relaxed); + metadata.admission_state.store(0, std::memory_order_relaxed); + if (adaptive_priority_enabled_) { + uint8_t evicted_hot = kEvictedHot; + if (entry.ghost_state.compare_exchange_strong(evicted_hot, kGhostAdmitted, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + // A ghost hit is admitted directly to protected. It must be reused + // while resident before it is allowed to leave another ghost. + (void)promote_evict_priority(block_id, kNormalPriority); + ghost_hot_hits_.fetch_add(1, std::memory_order_relaxed); } } + if (!metadata.ever_loaded.exchange(true, std::memory_order_acq_rel)) { + size_t head = loaded_head_.load(std::memory_order_relaxed); + do { + metadata.next_loaded = head; + } while (!loaded_head_.compare_exchange_weak( + head, block_id, std::memory_order_release, std::memory_order_relaxed)); + } + entry.buffer.store(buffer, std::memory_order_release); + entry.in_evict_queue.store(true, std::memory_order_relaxed); + entry.ref_count.store(1, std::memory_order_release); + + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast( + entry.evict_priority.load(std::memory_order_relaxed)))) { + // The final release will take the rare fallback registration path. + eviction_requeue_failed(block_id, owner_version_); + } + return buffer; +} + +bool VectorPageTable::cancel_block_load(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &entry = entry_at(block_id); + int expected = kLoadingRefCount; + return entry.ref_count.compare_exchange_strong(expected, kUnloadedRefCount, + std::memory_order_release, + std::memory_order_relaxed); } void VectorPageTable::release_block(block_id_t block_id) { assert(block_id < entry_num_.load(std::memory_order_relaxed)); Entry &e = entry_at(block_id); + // Installation normally registers the page; retry only after queue failure. if (e.ref_count.fetch_sub(1, std::memory_order_release) == 1) { + if (e.in_evict_queue.load(std::memory_order_relaxed)) { + return; + } std::atomic_thread_fence(std::memory_order_acquire); bool expected = false; if (e.in_evict_queue.compare_exchange_strong(expected, true, @@ -170,105 +626,288 @@ void VectorPageTable::release_block(block_id_t block_id) { BlockEvictionQueue::BlockType block; block.owner = this; block.owner_key = block_id; - block.version = 0; - BlockEvictionQueue::get_instance().add_single_block(block, 0); + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast( + e.evict_priority.load(std::memory_order_relaxed)))) { + eviction_requeue_failed(block_id, owner_version_); + } } } } -void VectorPageTable::evict_block(block_id_t block_id) { +bool VectorPageTable::evict_block(block_id_t block_id) { + return do_evict_block(block_id, /*force=*/false); +} + +bool VectorPageTable::force_evict_block(block_id_t block_id) { + return do_evict_block(block_id, /*force=*/true); +} + +bool VectorPageTable::reclaim_clean_block(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_relaxed)); + Entry &entry = entry_at(block_id); + int expected = 0; + if (!entry.ref_count.compare_exchange_strong(expected, kEvictingRefCount)) { + return false; + } + if (metadata_entry_at(block_id).is_dirty.load(std::memory_order_acquire)) { + entry.ref_count.store(0, std::memory_order_release); + return false; + } + + char *buffer = entry.buffer.exchange(nullptr, std::memory_order_acq_rel); + if (buffer != nullptr) { + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + } + inc_evict(entry.evict_priority.load(std::memory_order_relaxed)); + entry.in_evict_queue.store(false, std::memory_order_relaxed); + entry.ref_count.store(kUnloadedRefCount, std::memory_order_release); + return true; +} + +void VectorPageTable::force_evict_all_loaded() { + size_t block_id = loaded_head_.load(std::memory_order_acquire); + while (block_id != kInvalidLoadedBlock) { + const size_t next = metadata_entry_at(block_id).next_loaded; + assert(is_released(block_id)); + (void)force_evict_block(block_id); + block_id = next; + } +} + +size_t VectorPageTable::recover_eviction_queue() { + if (!eviction_recovery_needed_.exchange(false, std::memory_order_acq_rel)) { + return 0; + } + size_t recovered = 0; + size_t block_id = loaded_head_.load(std::memory_order_acquire); + while (block_id != kInvalidLoadedBlock) { + Entry &entry = entry_at(block_id); + const size_t next = metadata_entry_at(block_id).next_loaded; + if (entry.buffer.load(std::memory_order_acquire) != nullptr && + entry.ref_count.load(std::memory_order_acquire) == 0) { + bool expected = false; + if (entry.in_evict_queue.compare_exchange_strong( + expected, true, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (BlockEvictionQueue::get_instance().add_single_block( + block, static_cast(entry.evict_priority.load( + std::memory_order_relaxed)))) { + ++recovered; + } else { + entry.in_evict_queue.store(false, std::memory_order_release); + eviction_recovery_needed_.store(true, std::memory_order_release); + } + } + } + block_id = next; + } + return recovered; +} + +std::array +VectorPageTable::resident_pages_by_priority() const { + std::array resident{}; + size_t block_id = loaded_head_.load(std::memory_order_acquire); + while (block_id != kInvalidLoadedBlock) { + const Entry &entry = entry_at(block_id); + const size_t next = metadata_entry_at(block_id).next_loaded; + if (entry.buffer.load(std::memory_order_acquire) != nullptr) { + const uint8_t priority = + entry.evict_priority.load(std::memory_order_relaxed); + if (priority < kPriorityCount) { + ++resident[priority]; + } + } + block_id = next; + } + return resident; +} + +bool VectorPageTable::do_evict_block(block_id_t block_id, bool force) { assert(block_id < entry_num_.load(std::memory_order_relaxed)); Entry &e = entry_at(block_id); int expected = 0; - // Two-phase eviction to prevent data race on e.buffer with - // set_block_acquired. We first CAS to kEvicting (-1), which causes - // set_block_acquired to spin-wait; then do the actual work (flush, free, - // null buffer); finally store INT_MIN ("evicted") which unblocks - // set_block_acquired. - static constexpr int kEvicting = -1; - if (e.ref_count.compare_exchange_strong(expected, kEvicting)) { - char *buffer = e.buffer; - if (buffer && e.is_dirty.load(std::memory_order_relaxed) && - flush_callback_) { - flush_callback_(block_id, buffer, kVectorPageSize, e.file_offset); - e.is_dirty.store(false, std::memory_order_relaxed); + if (e.ref_count.compare_exchange_strong(expected, kEvictingRefCount)) { + // CLOCK gives recently referenced pages one more queue turn. + if (!force && e.referenced.load(std::memory_order_relaxed)) { + e.referenced.store(false, std::memory_order_relaxed); + inc_second_chance(); + // Preserve logical membership while moving the page to the tail. + e.ref_count.store(0, std::memory_order_release); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast( + e.evict_priority.load(std::memory_order_relaxed)))) { + eviction_requeue_failed(block_id, owner_version_); + } + return false; // spared, not reclaimed + } + if (!force && adaptive_priority_enabled_) { + uint8_t expected_priority = kNormalPriority; + if (e.evict_priority.compare_exchange_strong( + expected_priority, kLowPriority, std::memory_order_relaxed, + std::memory_order_relaxed)) { + inc_priority_demotion(kLowPriority); + uint8_t ghost_state = e.ghost_state.load(std::memory_order_relaxed); + if (ghost_state == kGhostAdmitted) { + // A ghost-admitted page that was not reused is stale. Do not let it + // renew itself indefinitely through repeated reloads. + e.ghost_state.store(kNoGhostHistory, std::memory_order_relaxed); + } else if (ghost_state != kEvictedHot) { + e.ghost_state.store(kEvictedHot, std::memory_order_relaxed); + ghost_hot_marks_.fetch_add(1, std::memory_order_relaxed); + } + // Demotion already gives the page another queue turn. A real reuse can + // set CLOCK again; an unconditional second chance only amplifies CPU + // work during sustained pressure. + e.referenced.store(false, std::memory_order_relaxed); + e.ref_count.store(0, std::memory_order_release); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast(kLowPriority))) { + eviction_requeue_failed(block_id, owner_version_); + } + return false; + } + } + MetadataEntry &metadata = metadata_entry_at(block_id); + char *buffer = e.buffer.load(std::memory_order_acquire); + if (buffer && metadata.is_dirty.load(std::memory_order_relaxed)) { + if (!force && writeback_callback_) { + bool scheduled = false; + try { + scheduled = writeback_callback_(block_id); + } catch (...) { + LOG_ERROR( + "VectorPageTable::evict_block: writeback callback threw for " + "block_id=%zu", + static_cast(block_id)); + } + if (scheduled) { + // Persistence belongs to the pool's writeback worker. Keep this + // page resident and queued until the worker makes it clean. + e.ref_count.store(0, std::memory_order_release); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast(e.evict_priority.load( + std::memory_order_relaxed)))) { + e.in_evict_queue.store(false, std::memory_order_relaxed); + eviction_recovery_needed_.store(true, std::memory_order_release); + } + return false; + } + } + int flush_rc = -1; + if (flush_callback_) { + try { + flush_rc = flush_callback_(block_id, buffer, kVectorPageSize, + metadata.file_offset); + } catch (...) { + LOG_ERROR( + "VectorPageTable::evict_block: flush callback threw for " + "block_id=%zu", + static_cast(block_id)); + } + } else { + LOG_ERROR( + "VectorPageTable::evict_block: dirty block %zu has no flush " + "callback", + static_cast(block_id)); + } + if (flush_rc != 0 && !force) { + // Keep a dirty page resident when writeback fails. + e.ref_count.store(0, std::memory_order_release); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, static_cast( + e.evict_priority.load(std::memory_order_relaxed)))) { + e.in_evict_queue.store(false, std::memory_order_relaxed); + eviction_recovery_needed_.store(true, std::memory_order_release); + } + return false; + } + if (flush_rc == 0) { + metadata.is_dirty.store(false, std::memory_order_relaxed); + inc_dirty_flush(); + } else { + LOG_ERROR( + "VectorPageTable::force_evict_block: discarding dirty block %zu " + "after flush failure during teardown", + static_cast(block_id)); + } } + buffer = e.buffer.exchange(nullptr, std::memory_order_acq_rel); if (buffer) { - e.buffer = nullptr; MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); } - // Transition to fully-evicted state. Use release so that the - // set_block_acquired acquire-load sees e.buffer == nullptr. - e.ref_count.store(std::numeric_limits::min(), - std::memory_order_release); + inc_evict(e.evict_priority.load(std::memory_order_relaxed)); + // Clear old membership before publishing the unloaded sentinel. + e.in_evict_queue.store(false, std::memory_order_relaxed); + e.ref_count.store(kUnloadedRefCount, std::memory_order_release); + return true; } - e.in_evict_queue.store(false, std::memory_order_relaxed); + + // Do not queue unloaded or transitioning entries. + if (expected < 0) { + return false; + } + + // Move pinned pages to the tail without duplicating membership. + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = owner_version_; + if (!BlockEvictionQueue::get_instance().add_single_block( + block, + static_cast(e.evict_priority.load(std::memory_order_relaxed)))) { + // Let release_block() retry registration when the last pin is dropped. + eviction_requeue_failed(block_id, owner_version_); + } + return false; } char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, size_t file_offset) { assert(block_id < entry_num_.load(std::memory_order_acquire)); - Entry &e = entry_at(block_id); - // Diagnostics for the kEvicting wait. The wait itself never gives up: - // the only thread that can transition kEvicting -> INT_MIN is the - // evict_block() owner, so abandoning the spin here would orphan the - // entry in kEvicting forever. Instead, we use bounded backoff and emit - // tiered logs so a stuck eviction is observable. - using clock = std::chrono::steady_clock; - const auto wait_start = clock::now(); - auto last_log = wait_start; - unsigned spin_count = 0; - bool warned = false; while (true) { - int current_count = e.ref_count.load(std::memory_order_acquire); - if (current_count >= 0) { - if (e.ref_count.compare_exchange_weak(current_count, current_count + 1, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); - return e.buffer; - } - } else if (current_count == std::numeric_limits::min()) { - // Fully evicted — safe to claim this entry for our new buffer. - e.buffer = buffer; - e.file_offset = file_offset; - e.in_evict_queue.store(false, std::memory_order_relaxed); - e.is_dirty.store(false, std::memory_order_relaxed); - e.ref_count.store(1, std::memory_order_release); - return e.buffer; - } else { - // kEvicting (-1): eviction is in progress on this entry. - // Tiered backoff: hot spin first, then short sleep, then longer sleep. - ++spin_count; - if (spin_count < 64) { - // Pure busy wait for the common ~μs case. - } else if (spin_count < 1024) { - std::this_thread::yield(); - } else if (spin_count < 8192) { - std::this_thread::sleep_for(std::chrono::microseconds(100)); - } else { - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } - // Tiered diagnostics: warn once after 100ms, error every 1s after 1s. - const auto now = clock::now(); - const auto elapsed = now - wait_start; - if (!warned && elapsed >= std::chrono::milliseconds(100)) { - LOG_WARN( - "set_block_acquired: long kEvicting wait on block_id=%zu " - "(>=100ms); evict_block may be slow", - static_cast(block_id)); - warned = true; - } - if (elapsed >= std::chrono::seconds(1) && - (now - last_log) >= std::chrono::seconds(1)) { - const auto secs = - std::chrono::duration_cast(elapsed).count(); - LOG_ERROR( - "set_block_acquired: stuck in kEvicting on block_id=%zu for " - "%lld s; evict_block owner may be hung or starved", - static_cast(block_id), static_cast(secs)); - last_log = now; + switch (try_claim_block_load(block_id)) { + case LoadClaimResult::kClaimed: + return publish_claimed_block(block_id, buffer, file_offset); + case LoadClaimResult::kResident: { + char *resident = acquire_block(block_id, /*record_reuse=*/false); + if (resident != nullptr) { + MemoryLimitPool::get_instance().release_buffer(buffer, + kVectorPageSize); + return resident; + } + break; } + case LoadClaimResult::kLoading: + case LoadClaimResult::kEvicting: + if (!wait_for_block_transition(block_id)) { + MemoryLimitPool::get_instance().release_buffer(buffer, + kVectorPageSize); + return nullptr; + } + break; } } } @@ -276,90 +915,687 @@ char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, VecBufferPool::VecBufferPool(const std::string &filename, bool writable) { file_name_ = filename; writable_ = writable; + page_table_.set_adaptive_priority(!writable_); #if defined(_MSC_VER) int flags = writable_ ? (O_RDWR | _O_BINARY) : (O_RDONLY | _O_BINARY); const std::wstring wide_filename = FileHelper::Utf8ToWide(filename); fd_ = wide_filename.empty() ? -1 : _wopen(wide_filename.c_str(), flags, 0644); + meta_fd_ = + wide_filename.empty() ? -1 : _wopen(wide_filename.c_str(), flags, 0644); #else - int flags = writable_ ? O_RDWR : O_RDONLY; - fd_ = ::open(filename.c_str(), flags, 0644); + int base_flags = writable_ ? O_RDWR : O_RDONLY; + // Buffered channel for unaligned metadata I/O. + meta_fd_ = ::open(filename.c_str(), base_flags, 0644); + // Keep metadata buffered, but bypass the kernel page cache for page data. + // Linux uses O_DIRECT; Darwin provides the equivalent through F_NOCACHE. + int data_flags = base_flags; +#ifdef O_DIRECT + data_flags |= O_DIRECT; #endif + fd_ = ::open(filename.c_str(), data_flags, 0644); +#ifdef O_DIRECT if (fd_ < 0) { + LOG_WARN( + "VecBufferPool: open with O_DIRECT failed for file[%s] (errno=%d), " + "falling back to buffered IO", + filename.c_str(), errno); + fd_ = ::open(filename.c_str(), base_flags, 0644); + direct_io_enabled_ = false; + } else { + direct_io_enabled_ = true; + } +#elif defined(F_NOCACHE) + if (fd_ >= 0) { + if (::fcntl(fd_, F_NOCACHE, 1) != 0) { + const int error = errno; + LOG_ERROR( + "VecBufferPool: failed to enable F_NOCACHE for file[%s] " + "(errno=%d)", + filename.c_str(), error); + ::close(fd_); + fd_ = -1; + errno = error; + } else { + direct_io_enabled_ = true; + } + } +#else + direct_io_enabled_ = false; +#endif +#endif + if (fd_ < 0 || meta_fd_ < 0) { + if (fd_ >= 0) { +#if defined(_MSC_VER) + _close(fd_); +#else + ::close(fd_); +#endif + } + if (meta_fd_ >= 0) { +#if defined(_MSC_VER) + _close(meta_fd_); +#else + ::close(meta_fd_); +#endif + } throw std::runtime_error("Failed to open file: " + filename); } #if defined(_MSC_VER) struct _stat64 st; if (_fstat64(fd_, &st) < 0) { _close(fd_); + _close(meta_fd_); #else struct stat st; if (fstat(fd_, &st) < 0) { ::close(fd_); + ::close(meta_fd_); #endif throw std::runtime_error("Failed to stat file: " + filename); } file_size_ = st.st_size; + initial_file_size_ = file_size_; +#if defined(__linux__) + // Select the process-wide backend; thread-local contexts are created lazily. + io_backend_type_ = direct_io_enabled_ ? IOBackend::Instance().available() + : IOBackendType::kPread; + aio_enabled_ = io_backend_type_ != IOBackendType::kPread; +#endif } -int VecBufferPool::init() { - size_t block_num = (file_size_ + kVectorPageSize - 1) / kVectorPageSize; - if (!page_table_.init(block_num)) { - LOG_ERROR( - "VecBufferPool::init: page_table_ init failed for file[%s], " - "file_size=%zu, block_num=%zu (exceeds " - "VectorPageTable::kMaxEntries=%zu)", - file_name_.c_str(), file_size_, block_num, - VectorPageTable::kMaxEntries); - return -1; +size_t VecBufferPool::metadata_bytes_for_page_count(size_t page_count, + bool writable) { + const size_t page_table_bytes = + VectorPageTable::metadata_bytes_for_entries(page_count); + if (page_table_bytes == std::numeric_limits::max()) { + return page_table_bytes; } - block_mutexes_ = - std::make_unique(VecBufferPool::kMutexBucketCount); - LOG_DEBUG("entry num: %zu, file_size: %zu", page_table_.entry_num(), - file_size_); - - // In writable mode, inject a flush callback into the page table so that - // evict_block()/flush_block()/flush_all() can pwrite dirty blocks back to - // the backing file without needing to know about fd_ directly. - if (writable_) { - int fd = fd_; - const std::string &name = file_name_; - page_table_.set_flush_callback([fd, &name](block_id_t /*block_id*/, - char *buf, size_t sz, - size_t off) -> int { - ssize_t w = zvec_pwrite(fd, buf, sz, off); - if (w != static_cast(sz)) { - LOG_ERROR( - "Buffer pool flush failed: file[%s], offset[%zu], " - "expected[%zu], got[%zd]", - name.c_str(), off, sz, w); - return -1; - } - return 0; - }); + // Writable files can grow after the pool opens. The mutex array cannot be + // replaced while readers and writers hold stripes, so allocate the stable + // maximum up front. Besides avoiding cross-page write contention, this + // lets a 128-page writeback batch hold distinct stripes after extend_file(). + const size_t mutex_count = writable ? kMutexBucketCount : 0; + const size_t mutex_bytes = mutex_count * sizeof(std::shared_mutex); + const size_t staging_bytes = + writable ? kBlockingAioBatchSize * kVectorPageSize : 0; + size_t io_staging_bytes = 0; +#if defined(__linux__) + if (writable && + IOBackend::Instance().available() == IOBackendType::kIoUring) { + io_staging_bytes = kBlockingAioBatchSize * kVectorPageSize; } - return 0; +#endif + if (mutex_bytes > std::numeric_limits::max() - staging_bytes || + mutex_bytes + staging_bytes > + std::numeric_limits::max() - io_staging_bytes) { + return std::numeric_limits::max(); + } + const size_t writable_bytes = mutex_bytes + staging_bytes + io_staging_bytes; + if (page_table_bytes > std::numeric_limits::max() - writable_bytes) { + return std::numeric_limits::max(); + } + return page_table_bytes + writable_bytes; } -VecBufferPoolHandle VecBufferPool::get_handle() { - return VecBufferPoolHandle(*this); -} +int VecBufferPool::init() { + if (initialized_) { + return 0; + } + if (writable_) { + // Configure the potentially allocating callback before reserving metadata. + try { + int fd = fd_; + page_table_.set_flush_callback( + [fd, &fn = file_name_](block_id_t /*block_id*/, char *buf, size_t sz, + size_t off) -> int { + ssize_t w = zvec_pwrite(fd, buf, sz, off); + if (w != static_cast(sz)) { + LOG_ERROR( + "Buffer pool flush failed: file[%s], offset[%zu], " + "expected[%zu], got[%zd]", + fn.c_str(), off, sz, w); + return -1; + } + return 0; + }); + page_table_.set_writeback_callback( + [this](block_id_t block_id) { return enqueue_writeback(block_id); }); + } catch (const std::bad_alloc &) { + LOG_ERROR( + "VecBufferPool::init: failed to allocate flush callback for file[%s]", + file_name_.c_str()); + return -1; + } + } -char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) { - assert(page_id < page_table_.entry_num()); - char *buffer = page_table_.acquire_block(page_id); - if (buffer) { - return buffer; + const size_t block_num = + file_size_ == 0 ? 0 : (file_size_ - 1) / kVectorPageSize + 1; + if (block_num > VectorPageTable::kMaxEntries) { + LOG_ERROR( + "VecBufferPool::init: file[%s] needs %zu entries, exceeding " + "VectorPageTable::kMaxEntries=%zu", + file_name_.c_str(), block_num, VectorPageTable::kMaxEntries); + return -1; } - std::lock_guard lock( - block_mutexes_[page_id % VecBufferPool::kMutexBucketCount]); - buffer = page_table_.acquire_block(page_id); - if (buffer) { - return buffer; + // Writable files grow in place. Keep stripe addresses stable for the pool's + // lifetime and avoid collapsing future writeback batches onto the few pages + // present when the file was opened. + const size_t mutex_count = writable_ ? kMutexBucketCount : 0; + const size_t mutex_charge = mutex_count * sizeof(std::shared_mutex); + const size_t staging_charge = + writable_ ? kBlockingAioBatchSize * kVectorPageSize : 0; + size_t io_staging_charge = 0; +#if defined(__linux__) + std::unique_ptr writeback_io_uring; + if (writable_ && io_backend_type_ == IOBackendType::kIoUring) { + // Keep the estimator and actual reservation stable even if creating this + // pool's ring fails and it must fall back to pwrite. + io_staging_charge = kBlockingAioBatchSize * kVectorPageSize; + try { + writeback_io_uring = std::make_unique(); + if (!writeback_io_uring->setup(kBlockingAioBatchSize)) { + writeback_io_uring.reset(); + LOG_WARN( + "VecBufferPool::init: io_uring writeback setup failed for " + "file[%s], falling back to pwrite", + file_name_.c_str()); + } + } catch (const std::bad_alloc &) { + writeback_io_uring.reset(); + LOG_WARN( + "VecBufferPool::init: cannot allocate io_uring writeback context " + "for file[%s], falling back to pwrite", + file_name_.c_str()); + } } - { - bool found = MemoryLimitPool::get_instance().try_acquire_buffer( +#endif + const size_t writable_metadata_charge = + mutex_charge + staging_charge + io_staging_charge; + if (writable_metadata_charge != 0 && + !MemoryLimitPool::get_instance().try_charge_metadata( + writable_metadata_charge)) { + LOG_ERROR( + "VecBufferPool::init: shared memory budget cannot reserve %zu bytes " + "for %zu page-lock stripes and writeback staging (file=%s)", + writable_metadata_charge, mutex_count, file_name_.c_str()); + return -1; + } + std::unique_ptr mutexes; + if (mutex_count != 0) { + try { + mutexes = std::make_unique(mutex_count); + } catch (const std::bad_alloc &) { + MemoryLimitPool::get_instance().release_metadata( + writable_metadata_charge); + LOG_ERROR( + "VecBufferPool::init: failed to allocate %zu page-lock stripes " + "(file=%s)", + mutex_count, file_name_.c_str()); + return -1; + } + } + char *writeback_staging = nullptr; + if (staging_charge != 0) { + writeback_staging = static_cast( + ailego_aligned_malloc(staging_charge, kVectorPageSize)); + if (writeback_staging == nullptr) { + MemoryLimitPool::get_instance().release_metadata( + writable_metadata_charge); + LOG_ERROR( + "VecBufferPool::init: failed to allocate %zu bytes of writeback " + "staging (file=%s)", + staging_charge, file_name_.c_str()); + return -1; + } + } + if (!page_table_.init(block_num)) { + if (writeback_staging != nullptr) { + ailego_free(writeback_staging); + } + MemoryLimitPool::get_instance().release_metadata(writable_metadata_charge); + LOG_ERROR( + "VecBufferPool::init: page_table_ init failed for file[%s], " + "file_size=%zu, block_num=%zu, required_metadata=%zu", + file_name_.c_str(), file_size_, block_num, + metadata_bytes_for_page_count(block_num, writable_)); + return -1; + } + block_mutexes_ = std::move(mutexes); + block_mutex_count_ = mutex_count; + writeback_staging_ = writeback_staging; + writeback_staging_size_ = staging_charge; + writeback_io_staging_charge_ = io_staging_charge; +#if defined(__linux__) + writeback_io_uring_ = std::move(writeback_io_uring); +#endif + LOG_DEBUG("entry num: %zu, file_size: %zu", page_table_.entry_num(), + file_size_); + + initialized_ = true; + if (writable_) { + try { + start_writeback(); + } catch (const std::exception &e) { + page_table_.set_writeback_callback({}); + LOG_WARN( + "VecBufferPool::init: failed to start background writeback for " + "file[%s], falling back to synchronous dirty eviction: %s", + file_name_.c_str(), e.what()); + } catch (...) { + page_table_.set_writeback_callback({}); + LOG_WARN( + "VecBufferPool::init: failed to start background writeback for " + "file[%s], falling back to synchronous dirty eviction", + file_name_.c_str()); + } + } + return 0; +} + +VecBufferPoolHandle VecBufferPool::get_handle() { + return VecBufferPoolHandle(*this); +} + +bool VecBufferPool::enqueue_writeback(block_id_t page_id) { + if (writeback_error() != 0) { + return false; + } + if (!page_table_.try_mark_writeback_pending(page_id)) { + return true; + } + + try { + { + std::lock_guard lock(writeback_mutex_); + if (writeback_stopping_) { + page_table_.clear_writeback_pending(page_id); + return false; + } + writeback_queue_.push_back(page_id); + writeback_requests_.fetch_add(1, std::memory_order_relaxed); + const uint64_t pending = + writeback_pending_.fetch_add(1, std::memory_order_relaxed) + 1; + uint64_t peak = writeback_peak_pending_.load(std::memory_order_relaxed); + while (peak < pending && !writeback_peak_pending_.compare_exchange_weak( + peak, pending, std::memory_order_relaxed, + std::memory_order_relaxed)) { + } + } + } catch (...) { + page_table_.clear_writeback_pending(page_id); + return false; + } + writeback_cv_.notify_one(); + return true; +} + +void VecBufferPool::start_writeback() { + std::lock_guard lock(writeback_mutex_); + if (writeback_thread_.joinable()) { + return; + } + writeback_stopping_ = false; + writeback_error_.store(0, std::memory_order_release); + writeback_thread_ = std::thread([this] { writeback_loop(); }); +} + +void VecBufferPool::stop_writeback() { + { + std::lock_guard lock(writeback_mutex_); + if (!writeback_thread_.joinable()) { + return; + } + writeback_stopping_ = true; + } + writeback_cv_.notify_all(); + writeback_thread_.join(); +} + +void VecBufferPool::drain_writeback() { + std::unique_lock lock(writeback_mutex_); + if (!writeback_thread_.joinable()) { + return; + } + writeback_drained_cv_.wait(lock, [this] { + return writeback_queue_.empty() && writeback_inflight_ == 0; + }); +} + +void VecBufferPool::writeback_loop() { + std::vector page_ids; + page_ids.reserve(kBlockingAioBatchSize); + while (true) { + { + std::unique_lock lock(writeback_mutex_); + writeback_cv_.wait(lock, [this] { + return writeback_stopping_ || !writeback_queue_.empty(); + }); + if (writeback_queue_.empty()) { + if (writeback_stopping_) { + break; + } + continue; + } + page_ids.clear(); + while (!writeback_queue_.empty() && + page_ids.size() < kBlockingAioBatchSize) { + page_ids.push_back(writeback_queue_.front()); + writeback_queue_.pop_front(); + } + writeback_inflight_ += page_ids.size(); + } + + try { + std::lock_guard flush_lock(writeback_flush_mutex_); + flush_writeback_batch(page_ids, writeback_staging_); + } catch (...) { + writeback_failures_.fetch_add(page_ids.size(), std::memory_order_relaxed); + int expected = 0; + (void)writeback_error_.compare_exchange_strong( + expected, EIO, std::memory_order_release, std::memory_order_relaxed); + LOG_ERROR("VecBufferPool writeback threw: file[%s], pages[%zu]", + file_name_.c_str(), page_ids.size()); + } + + for (block_id_t page_id : page_ids) { + page_table_.clear_writeback_pending(page_id); + } + writeback_pending_.fetch_sub(page_ids.size(), std::memory_order_relaxed); + for (block_id_t page_id : page_ids) { + (void)page_table_.reclaim_clean_block(page_id); + } + + { + std::lock_guard lock(writeback_mutex_); + writeback_inflight_ -= page_ids.size(); + if (writeback_queue_.empty() && writeback_inflight_ == 0) { + writeback_drained_cv_.notify_all(); + } + } + } + + { + std::lock_guard lock(writeback_mutex_); + if (writeback_queue_.empty() && writeback_inflight_ == 0) { + writeback_drained_cv_.notify_all(); + } + } +} + +bool VecBufferPool::flush_writeback_batch(std::vector &page_ids, + char *staging) { + if (page_ids.empty()) { + return true; + } + std::sort(page_ids.begin(), page_ids.end()); + page_ids.erase(std::unique(page_ids.begin(), page_ids.end()), page_ids.end()); + +#if defined(__linux__) + if (writeback_io_uring_ && writeback_io_uring_->is_valid()) { + bool all_ok = true; + size_t pos = 0; + while (pos < page_ids.size()) { + std::array selected_pages{}; + std::array buffers{}; + std::array locked_stripes{}; + std::array, kBlockingAioBatchSize> + locks; + size_t selected = 0; + + while (pos < page_ids.size() && selected < kBlockingAioBatchSize) { + const block_id_t page_id = page_ids[pos]; + if (!page_table_.is_block_dirty(page_id)) { + ++pos; + continue; + } + + const size_t stripe = page_id % block_mutex_count_; + bool stripe_already_locked = false; + for (size_t i = 0; i < selected; ++i) { + if (locked_stripes[i] == stripe) { + stripe_already_locked = true; + break; + } + } + // std::shared_mutex does not guarantee recursive shared ownership. + // Submit the current group before taking the same stripe again. + if (stripe_already_locked) { + break; + } + + char *buffer = page_table_.acquire_block(page_id, + /*record_reuse=*/false); + if (buffer == nullptr) { + ++pos; + continue; + } + locks[selected] = + std::shared_lock(block_mutexes_[stripe]); + if (!page_table_.is_block_dirty(page_id)) { + locks[selected].unlock(); + page_table_.release_block(page_id); + ++pos; + continue; + } + selected_pages[selected] = page_id; + buffers[selected] = buffer; + locked_stripes[selected] = stripe; + ++selected; + ++pos; + } + + if (selected == 0) { + continue; + } + + std::array requests{}; + for (size_t i = 0; i < selected; ++i) { + requests[i] = IoUringWrite(selected_pages[i] * kVectorPageSize, + kVectorPageSize, buffers[i]); + } + + writeback_batches_.fetch_add(1, std::memory_order_relaxed); + writeback_aio_batches_.fetch_add(1, std::memory_order_relaxed); + writeback_aio_pages_.fetch_add(selected, std::memory_order_relaxed); + bool aio_ok = writeback_io_uring_->execute_writes(fd_, requests.data(), + selected) == 0; + if (!aio_ok) { + writeback_aio_fallbacks_.fetch_add(1, std::memory_order_relaxed); + } + + size_t flushed = 0; + size_t failed = 0; + int batch_error = 0; + for (size_t i = 0; i < selected; ++i) { + bool page_ok = aio_ok; + if (!page_ok) { + errno = 0; + const ssize_t written = + zvec_pwrite(fd_, buffers[i], kVectorPageSize, + selected_pages[i] * kVectorPageSize); + page_ok = written == static_cast(kVectorPageSize); + if (!page_ok) { + const int error = errno != 0 ? errno : EIO; + if (batch_error == 0) { + batch_error = error; + } + LOG_ERROR( + "VecBufferPool writeback fallback failed: file[%s], " + "page[%zu], expected[%zu], got[%zd], errno[%d]", + file_name_.c_str(), static_cast(selected_pages[i]), + kVectorPageSize, written, error); + } + } + if (page_ok) { + page_table_.clear_dirty(selected_pages[i]); + ++flushed; + } else { + ++failed; + } + } + if (flushed != 0) { + page_table_.record_dirty_flush(flushed); + writeback_pages_.fetch_add(flushed, std::memory_order_relaxed); + } + if (failed != 0) { + writeback_failures_.fetch_add(failed, std::memory_order_relaxed); + int expected = 0; + (void)writeback_error_.compare_exchange_strong( + expected, batch_error != 0 ? batch_error : EIO, + std::memory_order_release, std::memory_order_relaxed); + all_ok = false; + } + + for (size_t i = 0; i < selected; ++i) { + locks[i].unlock(); + page_table_.release_block(selected_pages[i]); + } + } + return all_ok; + } +#endif + + bool all_ok = true; + const size_t max_run = + std::max(1, std::min(kBlockingAioBatchSize, block_mutex_count_)); + const size_t run_limit = staging != nullptr ? max_run : 1; + std::array buffers{}; + std::array, kBlockingAioBatchSize> locks; + + size_t pos = 0; + while (pos < page_ids.size()) { + const block_id_t run_start = page_ids[pos]; + size_t run_count = 0; + while (pos + run_count < page_ids.size() && run_count < run_limit && + page_ids[pos + run_count] == run_start + run_count) { + const block_id_t page_id = page_ids[pos + run_count]; + if (!page_table_.is_block_dirty(page_id)) { + break; + } + char *buffer = page_table_.acquire_block(page_id, + /*record_reuse=*/false); + if (buffer == nullptr) { + break; + } + locks[run_count] = std::shared_lock( + block_mutexes_[page_id % block_mutex_count_]); + if (!page_table_.is_block_dirty(page_id)) { + locks[run_count].unlock(); + page_table_.release_block(page_id); + break; + } + buffers[run_count] = buffer; + if (staging != nullptr) { + std::memcpy(staging + run_count * kVectorPageSize, buffer, + kVectorPageSize); + } + ++run_count; + } + + if (run_count == 0) { + ++pos; + continue; + } + + const char *write_buffer = staging != nullptr ? staging : buffers[0]; + // Without staging, resident pages are not contiguous; preserve correctness + // by submitting one page at a time. + const size_t submitted_pages = staging != nullptr ? run_count : 1; + const size_t submitted_size = submitted_pages * kVectorPageSize; + writeback_batches_.fetch_add(1, std::memory_order_relaxed); + const ssize_t written = zvec_pwrite(fd_, write_buffer, submitted_size, + run_start * kVectorPageSize); + const bool ok = written == static_cast(submitted_size); + if (ok) { + for (size_t i = 0; i < submitted_pages; ++i) { + page_table_.clear_dirty(run_start + i); + } + page_table_.record_dirty_flush(submitted_pages); + writeback_pages_.fetch_add(submitted_pages, std::memory_order_relaxed); + } else { + writeback_failures_.fetch_add(submitted_pages, std::memory_order_relaxed); + int expected = 0; + const int error = errno != 0 ? errno : EIO; + (void)writeback_error_.compare_exchange_strong(expected, error, + std::memory_order_release, + std::memory_order_relaxed); + LOG_ERROR( + "VecBufferPool writeback failed: file[%s], offset[%zu], " + "expected[%zu], got[%zd], errno[%d]", + file_name_.c_str(), run_start * kVectorPageSize, submitted_size, + written, error); + all_ok = false; + } + + for (size_t i = 0; i < run_count; ++i) { + locks[i].unlock(); + page_table_.release_block(run_start + i); + } + pos += staging != nullptr ? run_count : submitted_pages; + } + return all_ok; +} + +char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry, + bool record_reuse) { + assert(page_id < page_table_.entry_num()); + while (true) { + char *buffer = page_table_.acquire_block(page_id, record_reuse); + if (buffer) { + return buffer; + } + + const auto claim = page_table_.try_claim_block_load(page_id); + if (claim != VectorPageTable::LoadClaimResult::kClaimed) { + if (claim == VectorPageTable::LoadClaimResult::kResident) { + continue; + } + if (claim == VectorPageTable::LoadClaimResult::kLoading) { + singleflight_waits_.fetch_add(1, std::memory_order_relaxed); + } + // Recheck the stable state from the beginning after it completes. + if (!page_table_.wait_for_block_transition(page_id)) { + return nullptr; + } + continue; + } + + bool found = MemoryLimitPool::get_instance().try_acquire_buffer( kVectorPageSize, buffer); - if (!found) { + if (!found && writable_ && retry > 0) { + int no_progress_waits = 0; + uint64_t completed = writeback_pages_.load(std::memory_order_relaxed); + while (!found && no_progress_waits < retry) { + // Bound foreground queue scanning. Dirty candidates are only queued; + // disk I/O belongs to the writeback worker. + (void)BlockEvictionQueue::get_instance().batch_recycle(64); + found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer); + if (found || writeback_error() != 0) { + break; + } + + writeback_waits_.fetch_add(1, std::memory_order_relaxed); + const auto wait_start = std::chrono::steady_clock::now(); + const bool capacity_released = + MemoryLimitPool::get_instance().wait_for_available( + kVectorPageSize, std::chrono::milliseconds(100)); + const auto wait_end = std::chrono::steady_clock::now(); + writeback_wait_us_.fetch_add( + static_cast( + std::chrono::duration_cast( + wait_end - wait_start) + .count()), + std::memory_order_relaxed); + const uint64_t now = writeback_pages_.load(std::memory_order_relaxed); + if (capacity_released || now != completed) { + no_progress_waits = 0; + completed = now; + } else { + ++no_progress_waits; + } + } + } else if (!found) { for (int i = 0; i < retry; i++) { BlockEvictionQueue::get_instance().recycle(); found = MemoryLimitPool::get_instance().try_acquire_buffer( @@ -370,31 +1606,197 @@ char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) { } } if (!found) { - LOG_ERROR("Buffer pool failed to get free buffer: file[%s], page_id[%zu]", - file_name_.c_str(), page_id); + const auto memory_stats = MemoryLimitPool::get_instance().stats(); + const auto page_stats = page_table_.stats(); + const int error = writeback_error(); + if (error != 0) { + LOG_ERROR( + "Buffer pool allocation stopped after writeback failure: " + "file[%s], page_id[%zu], error[%d], used[%zu], " + "committed[%zu], free_buffers[%zu]", + file_name_.c_str(), page_id, error, memory_stats.used, + memory_stats.committed, memory_stats.free_buffers); + } else if (writable_) { + LOG_WARN( + "Buffer pool allocation made no progress: file[%s], " + "page_id[%zu], used[%zu], committed[%zu], free_buffers[%zu], " + "evict[%llu], second_chance[%llu]", + file_name_.c_str(), page_id, memory_stats.used, + memory_stats.committed, memory_stats.free_buffers, + static_cast(page_stats.evict), + static_cast(page_stats.second_chance)); + } else { + LOG_DEBUG( + "Buffer pool failed to get free buffer: file[%s], page_id[%zu], " + "used[%zu], committed[%zu], free_buffers[%zu], evict[%llu], " + "second_chance[%llu]", + file_name_.c_str(), page_id, memory_stats.used, + memory_stats.committed, memory_stats.free_buffers, + static_cast(page_stats.evict), + static_cast(page_stats.second_chance)); + } + (void)page_table_.cancel_block_load(page_id); return nullptr; } + + const size_t page_offset = page_id * kVectorPageSize; + // Count one miss per page for which this thread won the load claim. + miss_count_.fetch_add(1, std::memory_order_relaxed); + // Newly extended pages start zeroed; reload evicted pages from disk. + if (writable_ && page_offset >= initial_file_size_ && + !page_table_.is_ever_loaded(page_id)) { + std::memset(buffer, 0, kVectorPageSize); + } else { + // Accept and zero-pad an unaligned final page. + const size_t read_len = + direct_io_enabled_ + ? kVectorPageSize + : std::min(kVectorPageSize, file_size_ - page_offset); + if (read_len < kVectorPageSize) { + std::memset(buffer + read_len, 0, kVectorPageSize - read_len); + } + const ssize_t read_bytes = zvec_pread(fd_, buffer, read_len, page_offset); + if (read_bytes != static_cast(read_len)) { + // Accept short read at EOF: last page may not be full kVectorPageSize. + if (read_bytes > 0 && + (page_offset + static_cast(read_bytes) >= file_size_)) { + std::memset(buffer + read_bytes, 0, kVectorPageSize - read_bytes); + } else { + LOG_ERROR( + "Buffer pool failed to read file at offset: file[%s], " + "page_id[%zu], offset[%zu], expected[%zu], got[%zd]", + file_name_.c_str(), page_id, page_offset, read_len, read_bytes); + MemoryLimitPool::get_instance().release_buffer(buffer, + kVectorPageSize); + (void)page_table_.cancel_block_load(page_id); + return nullptr; + } + } + } + return page_table_.publish_claimed_block(page_id, buffer, page_offset); } +} - size_t page_offset = page_id * kVectorPageSize; - size_t expected_bytes = std::min(kVectorPageSize, file_size_ - page_offset); - if (expected_bytes < kVectorPageSize) { - std::memset(buffer + expected_bytes, 0, kVectorPageSize - expected_bytes); +bool VecBufferPool::try_acquire_resident_pages(const block_id_t *page_ids, + size_t count, char **pages) { + if (count == 0) return true; + if (!page_ids || !pages) return false; + + std::fill_n(pages, count, nullptr); + for (size_t i = 0; i < count; ++i) { + if (page_ids[i] >= page_table_.entry_num()) { + release_pages(page_ids, i); + std::fill_n(pages, i, nullptr); + return false; + } + pages[i] = try_acquire_buffer(page_ids[i]); + if (!pages[i]) { + release_pages(page_ids, i); + std::fill_n(pages, i, nullptr); + return false; + } } - ssize_t read_bytes = zvec_pread(fd_, buffer, expected_bytes, page_offset); - if (read_bytes != static_cast(expected_bytes)) { - LOG_ERROR( - "Buffer pool failed to read file at offset: file[%s], page_id[%zu], " - "offset[%zu], expected[%zu], got[%zd]", - file_name_.c_str(), page_id, page_offset, expected_bytes, read_bytes); - MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); - return nullptr; + return true; +} + +bool VecBufferPool::acquire_pages(const block_id_t *page_ids, size_t count, + char **pages) { + if (count == 0) return true; + if (!page_ids || !pages) return false; + + std::fill_n(pages, count, nullptr); + std::array miss_batch{}; + size_t miss_count = 0; + + // Pin hits before I/O so eviction cannot reclaim them before delivery. + for (size_t i = 0; i < count; ++i) { + if (page_ids[i] >= page_table_.entry_num()) { + for (size_t j = 0; j < i; ++j) { + if (pages[j]) { + page_table_.release_block(page_ids[j]); + pages[j] = nullptr; + } + } + return false; + } + pages[i] = try_acquire_buffer(page_ids[i]); + if (!pages[i]) { + miss_batch[miss_count++] = page_ids[i]; + if (miss_count == miss_batch.size()) { + (void)load_pages_aio(miss_batch.data(), miss_count, kLowPriority); + miss_count = 0; + } + } } - return page_table_.set_block_acquired(page_id, buffer, page_offset); + if (miss_count != 0) { + (void)load_pages_aio(miss_batch.data(), miss_count, kLowPriority); + } + + // Resolve and pin every output, including duplicates, after population. + for (size_t i = 0; i < count; ++i) { + if (pages[i]) continue; + // Population releases its installation pin before this resolution pass. + // Acquiring it here completes the original miss; it is not evidence of a + // later reuse and must leave the page in probation. + // A batch caller can roll back and use its direct-I/O fallback. Keep only + // one bounded foreground reclaim attempt here so a capacity miss does not + // leave the page in kLoadingRefCount while scanning the global queue. + pages[i] = acquire_buffer(page_ids[i], 1, /*record_reuse=*/false); + if (!pages[i]) { + for (size_t j = 0; j < count; ++j) { + if (pages[j]) { + page_table_.release_block(page_ids[j]); + pages[j] = nullptr; + } + } + return false; + } + } + return true; +} + +void VecBufferPool::release_pages(const block_id_t *page_ids, size_t count) { + if (!page_ids) return; + for (size_t i = 0; i < count; ++i) { + if (page_ids[i] < page_table_.entry_num()) { + page_table_.release_block(page_ids[i]); + } + } +} + +bool VecBufferPool::should_admit_page(block_id_t page_id) { + if (page_id >= page_table_.entry_num()) { + return false; + } + if (writable_ || !MemoryLimitPool::get_instance().under_cache_pressure()) { + return true; + } + + // Move the aging epoch every 64K evaluated cold misses. Exact per-page + // counters live in existing page-table padding, so this adds no side hash. + static constexpr uint64_t kObservationsPerEpoch = uint64_t{1} << 16; + const uint64_t observation = + admission_observations_.fetch_add(1, std::memory_order_relaxed); + const uint32_t epoch = + static_cast(observation / kObservationsPerEpoch); + const bool admitted = page_table_.should_admit_miss(page_id, epoch); + if (admitted) { + admission_admitted_.fetch_add(1, std::memory_order_relaxed); + } else { + admission_rejected_.fetch_add(1, std::memory_order_relaxed); + } + return admitted; } int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) { - ssize_t read_bytes = zvec_pread(fd_, buffer, length, offset); + if (length == 0) { + return 0; + } + if (buffer == nullptr || offset > file_size_ || + length > file_size_ - offset) { + return -1; + } + ssize_t read_bytes = zvec_pread(meta_fd_, buffer, length, offset); if (read_bytes != static_cast(length)) { LOG_ERROR( "Buffer pool failed to read file at offset: file[%s], offset[%zu], " @@ -405,6 +1807,64 @@ int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) { return 0; } +bool VecBufferPool::read_range_bypass(size_t file_offset, size_t length, + char *buffer) { + if (length == 0) { + return true; + } + if (buffer == nullptr || file_offset > file_size_ || + length > file_size_ - file_offset) { + return false; + } + + struct BypassScratch { + ~BypassScratch() { + if (page != nullptr) { + ailego_free(page); + } + } + char *page{nullptr}; + }; + static thread_local BypassScratch scratch; + if (scratch.page == nullptr) { + scratch.page = static_cast( + ailego_aligned_malloc(kVectorPageSize, kVectorPageSize)); + } + if (scratch.page == nullptr) { + return false; + } + char *page = scratch.page; + + size_t copied = 0; + size_t io_requests = 0; + bool ok = true; + while (copied < length) { + const size_t absolute = file_offset + copied; + const size_t page_offset = (absolute / kVectorPageSize) * kVectorPageSize; + const size_t within_page = absolute - page_offset; + const size_t copy_size = + std::min(length - copied, kVectorPageSize - within_page); + const size_t available = file_size_ - page_offset; + const size_t read_size = direct_io_enabled_ + ? kVectorPageSize + : std::min(kVectorPageSize, available); + + ++io_requests; + const ssize_t read_bytes = zvec_pread(fd_, page, read_size, page_offset); + if (read_bytes <= 0 || + within_page + copy_size > static_cast(read_bytes)) { + ok = false; + break; + } + std::memcpy(buffer + copied, page + within_page, copy_size); + copied += copy_size; + } + if (ok) { + record_bypass_read(length, io_requests); + } + return ok; +} + int VecBufferPool::write_range(size_t file_offset, size_t length, const char *src) { if (!writable_) { @@ -415,20 +1875,28 @@ int VecBufferPool::write_range(size_t file_offset, size_t length, if (length == 0) { return 0; } + if (src == nullptr || file_offset > file_size_ || + length > file_size_ - file_offset) { + LOG_ERROR( + "write_range exceeds file bounds: file[%s], offset[%zu], " + "length[%zu], file_size[%zu]", + file_name_.c_str(), file_offset, length, file_size_); + return -1; + } size_t first_page = file_offset / kVectorPageSize; size_t last_page = (file_offset + length - 1) / kVectorPageSize; size_t remaining = length; size_t src_cursor = 0; for (size_t pg = first_page; pg <= last_page; ++pg) { - // Loading the page ensures we do not clobber unrelated bytes within the - // same page when the write is not page-aligned. acquire_buffer() pre-fills - // from the backing file (or zero-pads beyond EOF). + // Load partial pages before modifying them. char *page = this->acquire_buffer(pg, 50); if (!page) { LOG_ERROR("write_range acquire failed: file[%s], page[%zu]", file_name_.c_str(), pg); return -1; } + std::unique_lock page_lock( + block_mutexes_[pg % block_mutex_count_]); size_t page_start = pg * kVectorPageSize; size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); @@ -441,6 +1909,59 @@ int VecBufferPool::write_range(size_t file_offset, size_t length, return 0; } +int VecBufferPool::write_fragments(const VecBufferWriteFragment *fragments, + size_t count) { + if (!writable_ || (count != 0 && fragments == nullptr)) { + return -1; + } + if (count == 0) { + return 0; + } + + size_t page_id = std::numeric_limits::max(); + bool has_data = false; + for (size_t i = 0; i < count; ++i) { + const auto &fragment = fragments[i]; + if (fragment.length == 0) { + continue; + } + if (fragment.src == nullptr || fragment.file_offset > file_size_ || + fragment.length > file_size_ - fragment.file_offset) { + return -1; + } + const size_t fragment_page = fragment.file_offset / kVectorPageSize; + const size_t offset_in_page = fragment.file_offset % kVectorPageSize; + if (fragment.length > kVectorPageSize - offset_in_page || + (has_data && fragment_page != page_id)) { + return -1; + } + page_id = fragment_page; + has_data = true; + } + if (!has_data) { + return 0; + } + + char *page = acquire_buffer(static_cast(page_id), 50); + if (page == nullptr) { + return -1; + } + { + std::unique_lock page_lock( + block_mutexes_[page_id % block_mutex_count_]); + for (size_t i = 0; i < count; ++i) { + const auto &fragment = fragments[i]; + if (fragment.length != 0) { + std::memcpy(page + fragment.file_offset % kVectorPageSize, fragment.src, + fragment.length); + } + } + page_table_.mark_dirty(page_id); + } + page_table_.release_block(page_id); + return 0; +} + int VecBufferPool::write_meta(size_t offset, size_t length, const char *buffer) { if (!writable_) { @@ -448,7 +1969,14 @@ int VecBufferPool::write_meta(size_t offset, size_t length, file_name_.c_str()); return -1; } - ssize_t w = zvec_pwrite(fd_, buffer, length, offset); + if (length == 0) { + return 0; + } + if (buffer == nullptr || offset > file_size_ || + length > file_size_ - offset) { + return -1; + } + ssize_t w = zvec_pwrite(meta_fd_, buffer, length, offset); if (w != static_cast(length)) { LOG_ERROR( "Buffer pool failed to write meta: file[%s], offset[%zu], " @@ -463,27 +1991,47 @@ int VecBufferPool::flush_all() { if (!writable_) { return 0; } + // Establish one persistence owner before the full scan. This also gives + // callers a deterministic drain point for all previously queued pages. + drain_writeback(); + std::unique_lock flush_lock(writeback_flush_mutex_); + const size_t total = page_table_.entry_num(); + if (total == 0) { + return 0; + } + int rc = 0; size_t total_dirty = 0; - size_t fail_count = 0; - for (size_t i = 0; i < page_table_.entry_num(); ++i) { - if (page_table_.is_block_dirty(i)) { - ++total_dirty; - int r = page_table_.flush_block(i); - if (r != 0) { - rc = r; - ++fail_count; + size_t failed_batches = 0; + std::vector dirty_pages; + dirty_pages.reserve(kBlockingAioBatchSize); + for (size_t page_id = 0; page_id < total; ++page_id) { + if (page_table_.is_block_dirty(page_id)) { + dirty_pages.push_back(page_id); + } + if (dirty_pages.size() == kBlockingAioBatchSize || + (page_id + 1 == total && !dirty_pages.empty())) { + total_dirty += dirty_pages.size(); + if (!flush_writeback_batch(dirty_pages, writeback_staging_)) { + rc = -1; + ++failed_batches; } + dirty_pages.clear(); } } - if (fail_count != 0) { - // Aggregated diagnostic so that callers (notably ~VecBufferPool, which - // discards the return value) cannot silently lose dirty pages: any - // unflushed page at this point means the on-disk image is now stale. + + if (failed_batches != 0) { LOG_ERROR( - "VecBufferPool::flush_all: %zu/%zu dirty page(s) failed to flush, " - "file[%s] last_rc=%d -- on-disk data may be stale.", - fail_count, total_dirty, file_name_.c_str(), rc); + "VecBufferPool::flush_all: %zu writeback batch(es) covering %zu dirty " + "page(s) failed, file[%s] last_rc=%d -- on-disk data may be stale.", + failed_batches, total_dirty, file_name_.c_str(), rc); + } else { + writeback_error_.store(0, std::memory_order_release); + } + flush_lock.unlock(); + drain_writeback(); + if (writeback_error() != 0) { + rc = -1; } return rc; } @@ -497,11 +2045,16 @@ bool VecBufferPool::extend_file(size_t new_size) { if (new_size <= file_size_) { return true; } - // Pre-validate against the page table's static capacity BEFORE mutating - // any on-disk state. Otherwise a successful ftruncate followed by a - // failed page_table_.extend() would leave the file size and the page - // table out of sync (file grew, but no Entry slots cover the new range). - size_t new_entry_num = (new_size + kVectorPageSize - 1) / kVectorPageSize; + // O_DIRECT requires page-aligned backing-file growth. + if (new_size % kVectorPageSize != 0) { + LOG_ERROR( + "extend_file target must be page-aligned: file[%s], new_size[%zu], " + "page_size[%zu]", + file_name_.c_str(), new_size, kVectorPageSize); + return false; + } + // Validate page-table capacity before changing the file. + const size_t new_entry_num = (new_size - 1) / kVectorPageSize + 1; if (new_entry_num > VectorPageTable::kMaxEntries) { LOG_ERROR( "extend_file: requested new_size=%zu would require %zu page entries, " @@ -510,44 +2063,51 @@ bool VecBufferPool::extend_file(size_t new_size) { file_name_.c_str()); return false; } + const size_t old_entry_num = page_table_.entry_num(); + if (new_entry_num > old_entry_num && !page_table_.extend(new_entry_num)) { + LOG_ERROR( + "extend_file: page_table_.extend(%zu) failed before resizing " + "file=%s to %zu bytes", + new_entry_num, file_name_.c_str(), new_size); + return false; + } + #if defined(_MSC_VER) if (_chsize_s(fd_, static_cast(new_size)) != 0) { LOG_ERROR("extend_file _chsize_s failed: file[%s], new_size[%zu]", file_name_.c_str(), new_size); + if (!page_table_.rollback_extend(old_entry_num)) { + LOG_ERROR("extend_file: failed to roll back page table for file[%s]", + file_name_.c_str()); + } return false; } #else if (::ftruncate(fd_, static_cast(new_size)) != 0) { LOG_ERROR("extend_file ftruncate failed: file[%s], new_size[%zu]", file_name_.c_str(), new_size); + if (!page_table_.rollback_extend(old_entry_num)) { + LOG_ERROR("extend_file: failed to roll back page table for file[%s]", + file_name_.c_str()); + } return false; } #endif file_size_ = new_size; - // Extend the page table to cover the new file range. Existing entries - // stay at their original addresses so concurrent readers are unaffected. - // Capacity has already been validated above, so this should never fail; - // a failure here would indicate a programming error and is logged. - if (new_entry_num > page_table_.entry_num()) { - if (!page_table_.extend(new_entry_num)) { - LOG_ERROR( - "extend_file: page_table_.extend(%zu) failed unexpectedly after " - "capacity pre-check (file=%s, new_size=%zu).", - new_entry_num, file_name_.c_str(), new_size); - return false; - } - } return true; } char *VecBufferPoolHandle::get_single_page(size_t file_offset, size_t len, size_t &out_page_id) { + if (file_offset >= pool_.file_size_ || len > pool_.file_size_ - file_offset) { + return nullptr; + } size_t first_page = file_offset / kVectorPageSize; - assert(len == 0 || (file_offset + len - 1) / kVectorPageSize == first_page); + assert(len == 0 || len <= kVectorPageSize - (file_offset % kVectorPageSize)); out_page_id = first_page; char *page = pool_.acquire_buffer(first_page, 50); if (!page) { - LOG_ERROR( + LOG_DEBUG( "VecBufferPoolHandle::get_single_page: acquire_buffer failed, " "file_offset=%zu, len=%zu, page=%zu, page_size=%zu", file_offset, len, first_page, kVectorPageSize); @@ -556,36 +2116,231 @@ char *VecBufferPoolHandle::get_single_page(size_t file_offset, size_t len, return page + (file_offset - first_page * kVectorPageSize); } +bool VecBufferPoolHandle::acquire_pages(const block_id_t *page_ids, + size_t count, char **pages) { + return pool_.acquire_pages(page_ids, count, pages); +} + +bool VecBufferPoolHandle::try_acquire_resident_pages(const block_id_t *page_ids, + size_t count, + char **pages) { + return pool_.try_acquire_resident_pages(page_ids, count, pages); +} + +void VecBufferPoolHandle::release_pages(const block_id_t *page_ids, + size_t count) { + pool_.release_pages(page_ids, count); +} + bool VecBufferPoolHandle::read_range(size_t file_offset, size_t len, char *out) { if (len == 0) { return true; } + if (out == nullptr || file_offset > pool_.file_size_ || + len > pool_.file_size_ - file_offset) { + return false; + } size_t first_page = file_offset / kVectorPageSize; size_t last_page = (file_offset + len - 1) / kVectorPageSize; size_t remaining = len; size_t dst_cursor = 0; + + // Protect payload copies only for writable pools. + if (pool_.writable_) { + for (size_t pg = first_page; pg <= last_page; ++pg) { + char *page = pool_.acquire_buffer(static_cast(pg), 50); + if (page == nullptr) { + return false; + } + std::shared_lock page_lock( + pool_.block_mutexes_[pg % pool_.block_mutex_count_]); + const size_t page_start = pg * kVectorPageSize; + const size_t intra_offset = + (pg == first_page) ? (file_offset - page_start) : 0; + const size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + page_lock.unlock(); + pool_.page_table_.release_block(static_cast(pg)); + dst_cursor += chunk; + remaining -= chunk; + } + return true; + } + + static constexpr size_t kMaxRunPages = 1024; // 4MB max per bulk read + for (size_t pg = first_page; pg <= last_page; ++pg) { - char *page = pool_.acquire_buffer(pg, 50); - if (!page) { + char *page = pool_.page_table_.acquire_block(pg); + if (page) { + size_t page_start = pg * kVectorPageSize; + size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + pool_.page_table_.release_block(pg); + dst_cursor += chunk; + remaining -= chunk; + continue; + } + + size_t run_start = pg; + size_t run_end = pg + 1; + while (run_end <= last_page && !pool_.page_table_.is_loaded(run_end) && + (run_end - run_start) < kMaxRunPages) { + ++run_end; + } + size_t run_pages = run_end - run_start; + + if (run_pages <= 3) { + for (size_t j = 0; j < run_pages; ++j) { + block_id_t pid = static_cast(run_start + j); + // Once a sparse resident set breaks a cold range into short holes, + // apply the same frequency admission policy as the bulk path. Without + // this check, one- to three-page first touches continuously evict the + // useful working set even though long first-touch runs are bypassed. + const bool use_admission = + MemoryLimitPool::get_instance().page_admission_reserve() != 0; + const bool admit = !use_admission || pool_.should_admit_page(pid); + if (admit) { + page = pool_.acquire_buffer(pid, 50); + } else { + pool_.miss_count_.fetch_add(1, std::memory_order_relaxed); + page = nullptr; + } + size_t page_start = pid * kVectorPageSize; + size_t intra_offset = + (pid == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + if (page != nullptr) { + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + pool_.page_table_.release_block(pid); + } else if (!pool_.read_range_bypass(page_start + intra_offset, chunk, + out + dst_cursor)) { + return false; + } + dst_cursor += chunk; + remaining -= chunk; + } + pg = run_end - 1; + continue; + } + + size_t run_bytes = run_pages * kVectorPageSize; + size_t run_file_off = run_start * kVectorPageSize; + + char *bulk_buf = + static_cast(ailego_aligned_malloc(run_bytes, 4096)); + if (!bulk_buf) { + page = pool_.acquire_buffer(static_cast(pg), 50); + size_t page_start = pg * kVectorPageSize; + size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + if (page != nullptr) { + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + pool_.page_table_.release_block(static_cast(pg)); + } else if (!pool_.read_range_bypass(page_start + intra_offset, chunk, + out + dst_cursor)) { + return false; + } + dst_cursor += chunk; + remaining -= chunk; + continue; + } + + ssize_t got = zvec_pread(pool_.fd_, bulk_buf, run_bytes, run_file_off); + // read_range validated file_offset + len against file_size_ above. + size_t needed_bytes = (file_offset + len) - run_file_off; + if (needed_bytes > run_bytes) needed_bytes = run_bytes; + if (got < 0 || static_cast(got) < needed_bytes) { + ailego_free(bulk_buf); LOG_ERROR( - "VecBufferPoolHandle::read_range: acquire_buffer failed, " - "file_offset=%zu, len=%zu, page=%zu, first_page=%zu, last_page=%zu, " - "page_size=%zu", - file_offset, len, pg, first_page, last_page, kVectorPageSize); + "read_range bulk pread failed: off=%zu len=%zu got=%zd needed=%zu", + run_file_off, run_bytes, got, needed_bytes); return false; } - size_t page_start = pg * kVectorPageSize; - size_t intra_offset = (pg == first_page) ? (file_offset - page_start) : 0; - size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + size_t actually_read = static_cast(got); + // Account for pages populated outside acquire_buffer(). + size_t pages_read = (actually_read + kVectorPageSize - 1) / kVectorPageSize; + pool_.miss_count_.fetch_add(pages_read, std::memory_order_relaxed); + + for (size_t j = 0; j < run_pages; ++j) { + block_id_t pid = static_cast(run_start + j); + size_t page_start = pid * kVectorPageSize; + size_t intra_offset = + (pid == first_page) ? (file_offset - page_start) : 0; + size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); + std::memcpy(out + dst_cursor, + bulk_buf + j * kVectorPageSize + intra_offset, chunk); + dst_cursor += chunk; + remaining -= chunk; + + size_t page_end_in_buf = (j + 1) * kVectorPageSize; + // Large sequential reads (for example IVF posting-list scans) must not + // populate every cold page once the shared cache is under pressure. + // Reuse the same compact frequency admission policy as batched random + // reads so first-touch scan pages bypass while repeated pages can enter. + if (page_end_in_buf <= actually_read && + !pool_.page_table_.is_loaded(pid) && pool_.should_admit_page(pid)) { + char *page_buf = nullptr; + bool found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, page_buf); + if (!found) { + BlockEvictionQueue::get_instance().recycle(); + found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, page_buf); + } + if (found) { + std::memcpy(page_buf, bulk_buf + j * kVectorPageSize, + kVectorPageSize); + char *installed = pool_.page_table_.set_block_acquired( + pid, page_buf, run_file_off + j * kVectorPageSize); + if (installed != nullptr) { + pool_.page_table_.release_block(pid); + } + } + } + } + ailego_free(bulk_buf); + pg = run_end - 1; + } + return true; +} + +bool VecBufferPoolHandle::read_range_immutable(size_t file_offset, size_t len, + char *out) { + if (len == 0) { + return true; + } + if (out == nullptr || file_offset > pool_.file_size_ || + len > pool_.file_size_ - file_offset) { + return false; + } + + const size_t first_page = file_offset / kVectorPageSize; + const size_t last_page = (file_offset + len - 1) / kVectorPageSize; + size_t remaining = len; + size_t dst_cursor = 0; + for (size_t pg = first_page; pg <= last_page; ++pg) { + char *page = pool_.acquire_buffer(static_cast(pg), 50); + if (page == nullptr) { + return false; + } + const size_t page_start = pg * kVectorPageSize; + const size_t intra_offset = pg == first_page ? file_offset - page_start : 0; + const size_t chunk = std::min(kVectorPageSize - intra_offset, remaining); std::memcpy(out + dst_cursor, page + intra_offset, chunk); - pool_.page_table_.release_block(pg); + pool_.page_table_.release_block(static_cast(pg)); dst_cursor += chunk; remaining -= chunk; } return true; } +bool VecBufferPoolHandle::read_range_bypass(size_t file_offset, size_t len, + char *out) { + return pool_.read_range_bypass(file_offset, len, out); +} + int VecBufferPoolHandle::get_meta(size_t offset, size_t length, char *buffer) { return pool_.get_meta(offset, length, buffer); } @@ -595,6 +2350,11 @@ int VecBufferPoolHandle::write_range(size_t file_offset, size_t len, return pool_.write_range(file_offset, len, src); } +int VecBufferPoolHandle::write_fragments( + const VecBufferWriteFragment *fragments, size_t count) { + return pool_.write_fragments(fragments, count); +} + int VecBufferPoolHandle::write_meta(size_t offset, size_t length, const char *buffer) { return pool_.write_meta(offset, length, buffer); @@ -613,11 +2373,605 @@ void VecBufferPoolHandle::release_one(block_id_t block_id) { } void VecBufferPoolHandle::acquire_one(block_id_t block_id) { - // The caller must guarantee the block is already loaded before calling - // acquire_one(). The return value of acquire_block() is intentionally - // ignored here, as a null return would indicate a contract violation. + // Caller guarantees the page is resident. pool_.page_table_.acquire_block(block_id); } +void VecBufferPool::warmup() { + const size_t total_pages = page_table_.entry_num(); + // Read sequentially in 4 MB chunks. + static constexpr size_t kChunkPages = 1024; + const size_t kChunkSize = kChunkPages * kVectorPageSize; + + // Aligned buffer for bulk read (O_DIRECT requires alignment). + char *chunk_buf = + static_cast(ailego_aligned_malloc(kChunkSize, 4096)); + if (!chunk_buf) return; + + size_t loaded = 0; + bool pool_full = false; + for (size_t base = 0; base < total_pages && !pool_full; base += kChunkPages) { + const size_t pages_in_chunk = std::min(kChunkPages, total_pages - base); + const size_t read_bytes = pages_in_chunk * kVectorPageSize; + const size_t file_offset = base * kVectorPageSize; + const size_t expected_bytes = + std::min(read_bytes, file_size_ - file_offset); + + // One large sequential pread instead of N individual ones. + ssize_t got = zvec_pread(fd_, chunk_buf, read_bytes, file_offset); + if (got != static_cast(expected_bytes)) break; + // The final page may extend past EOF. Keep its unread tail deterministic, + // matching the regular single-page load path. + if (expected_bytes < read_bytes) { + std::memset(chunk_buf + expected_bytes, 0, read_bytes - expected_bytes); + } + + // Distribute chunk data into individual page buffers. + for (size_t j = 0; j < pages_in_chunk; ++j) { + auto page_id = static_cast(base + j); + // Skip if already loaded. + char *existing = page_table_.acquire_block(page_id); + if (existing) { + page_table_.release_block(page_id); + ++loaded; + continue; + } + // Allocate page buffer from pool (no retry - stop if full). + char *buf = nullptr; + bool found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buf); + if (!found) { + pool_full = true; + break; + } + std::memcpy(buf, chunk_buf + j * kVectorPageSize, kVectorPageSize); + char *installed = page_table_.set_block_acquired( + page_id, buf, file_offset + j * kVectorPageSize); + if (installed != nullptr) { + page_table_.release_block(page_id); + ++loaded; + } + } + } + ailego_free(chunk_buf); + LOG_DEBUG("VecBufferPool::warmup: preloaded %zu/%zu pages for file[%s]", + loaded, total_pages, file_name_.c_str()); +} + +void VecBufferPool::prefetch_pages(block_id_t first_page, size_t page_count, + uint8_t priority) { + const size_t total_pages = page_table_.entry_num(); + if (priority > kHighPriority || page_count == 0 || + first_page >= total_pages) { + return; + } + page_count = std::min(page_count, total_pages - first_page); + + bool all_loaded = true; + for (size_t page = first_page; page < first_page + page_count; ++page) { + if (page_table_.is_loaded(page)) { + page_table_.promote_evict_priority(page, priority); + } else { + all_loaded = false; + } + } + if (all_loaded) { + return; + } + +#if defined(__linux__) + if (aio_enabled_) { + prefetch_pages_aio(first_page, page_count, priority); + return; + } +#endif + + prefetch_pages_sync(first_page, page_count, priority); +} + +void VecBufferPool::prefetch_pages_sync(block_id_t first_page, + size_t page_count, uint8_t priority) { + const size_t end_page = first_page + page_count; + + static constexpr size_t kChunkPages = 1024; + const size_t kChunkSize = kChunkPages * kVectorPageSize; + char *chunk_buf = + static_cast(ailego_aligned_malloc(kChunkSize, 4096)); + if (!chunk_buf) return; + + bool pool_full = false; + size_t pg = first_page; + while (pg < end_page && !pool_full) { + if (page_table_.is_loaded(pg)) { + page_table_.promote_evict_priority(pg, priority); + ++pg; + continue; + } + size_t run_start = pg; + size_t run_end = pg + 1; + while (run_end < end_page && !page_table_.is_loaded(run_end) && + (run_end - run_start) < kChunkPages) { + ++run_end; + } + + size_t run_pages = run_end - run_start; + size_t read_bytes = run_pages * kVectorPageSize; + size_t file_off = run_start * kVectorPageSize; + size_t expected_bytes = std::min(read_bytes, file_size_ - file_off); + ssize_t got = zvec_pread(fd_, chunk_buf, read_bytes, file_off); + if (got != static_cast(expected_bytes)) { + pg = run_end; + continue; + } + if (expected_bytes < read_bytes) { + std::memset(chunk_buf + expected_bytes, 0, read_bytes - expected_bytes); + } + + for (size_t j = 0; j < run_pages; ++j) { + block_id_t pid = static_cast(run_start + j); + if (page_table_.is_loaded(pid)) { + page_table_.promote_evict_priority(pid, priority); + continue; + } + char *buf = nullptr; + bool found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buf); + if (!found) { + BlockEvictionQueue::get_instance().recycle(); + found = MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buf); + if (!found) { + pool_full = true; + break; + } + } + std::memcpy(buf, chunk_buf + j * kVectorPageSize, kVectorPageSize); + page_table_.promote_evict_priority(pid, priority); + char *installed = page_table_.set_block_acquired( + pid, buf, file_off + j * kVectorPageSize); + if (installed != nullptr) { + page_table_.release_block(pid); + } + } + pg = run_end; + } + ailego_free(chunk_buf); +} + +void VecBufferPoolHandle::prefetch_range(size_t file_offset, size_t len, + uint8_t priority) { + if (len == 0 || file_offset >= pool_.file_size_) return; + len = std::min(len, pool_.file_size_ - file_offset); + size_t first_page = file_offset / kVectorPageSize; + size_t last_page = (file_offset + len - 1) / kVectorPageSize; + pool_.prefetch_pages(static_cast(first_page), + last_page - first_page + 1, priority); +} + +#if defined(__linux__) +namespace { +template +struct ThreadLocalIoUringContext { + IoUringRing ring{}; + bool inited{false}; + + bool ensure() { + if (!inited) { + inited = true; + ring.setup(QueueDepth); + } + return ring.is_valid(); + } +}; + +template +struct ThreadLocalAioContext { + io_context_t ctx{nullptr}; + bool inited{false}; + + bool ensure() { + if (inited) return ctx != nullptr; + inited = true; + if (!LibAioLoader::Instance().load() || + !LibAioLoader::Instance().is_available()) { + return false; + } + if (LibAioLoader::Instance().io_setup(QueueDepth, &ctx) == 0) { + return true; + } + ctx = nullptr; + return false; + } + + bool destroy_context(const char *context_name) { + if (!ctx) return true; + const int ret = LibAioLoader::Instance().io_destroy(ctx); + if (ret != 0) { + LOG_ERROR( + "%s: io_destroy failed, ret=%d; in-flight buffers remain " + "quarantined", + context_name, ret); + return false; + } + ctx = nullptr; + return true; + } +}; + +// One blocking AIO context per calling thread, shared by reads and prefetches. +struct ThreadLocalBlockingAioCtx + : ThreadLocalAioContext { + char *quarantined[kBlockingAioBatchSize]{}; + size_t quarantined_count{0}; + + void quarantine(char **buffers, size_t count) { + for (size_t i = 0; i < count; ++i) { + if (buffers[i]) { + assert(quarantined_count < kBlockingAioBatchSize); + quarantined[quarantined_count++] = buffers[i]; + buffers[i] = nullptr; + } + } + } + + void release_quarantined() { + for (size_t i = 0; i < quarantined_count; ++i) { + MemoryLimitPool::get_instance().release_buffer(quarantined[i], + kVectorPageSize); + } + quarantined_count = 0; + } + + ~ThreadLocalBlockingAioCtx() { + if (destroy_context("ThreadLocalBlockingAioCtx")) { + release_quarantined(); + } + } +}; +static thread_local ThreadLocalIoUringContext + tl_blocking_io_uring; +static thread_local ThreadLocalBlockingAioCtx tl_blocking_aio; +} // namespace +#endif + +void VecBufferPool::prefetch_pages_aio(block_id_t first_page, size_t page_count, + uint8_t priority) { + const size_t total_pages = page_table_.entry_num(); + if (priority > kHighPriority || page_count == 0 || + first_page >= total_pages) { + return; + } + page_count = std::min(page_count, total_pages - first_page); + std::array pages{}; + size_t offset = 0; + while (offset < page_count) { + const size_t batch = std::min(kBlockingAioBatchSize, page_count - offset); + for (size_t i = 0; i < batch; ++i) { + pages[i] = first_page + offset + i; + } + if (!load_pages_aio(pages.data(), batch, priority)) { + prefetch_pages_sync(first_page, page_count, priority); + return; + } + offset += batch; + } +} + +bool VecBufferPool::load_pages_aio(const block_id_t *page_ids, size_t count, + uint8_t priority) { + if (count == 0) return true; + if (page_ids == nullptr || priority > kHighPriority) return false; + +#if defined(__linux__) + if (!aio_enabled_) return false; + bool use_io_uring = io_backend_type_ == IOBackendType::kIoUring && + tl_blocking_io_uring.ensure(); + if (!use_io_uring && !tl_blocking_aio.ensure()) return false; + + size_t cursor = 0; + while (cursor < count) { + std::array candidate_pages{}; + size_t candidate_count = 0; + while (cursor < count && candidate_count < kBlockingAioBatchSize) { + const block_id_t pid = page_ids[cursor++]; + if (pid >= page_table_.entry_num()) return false; + bool duplicate = false; + for (size_t i = 0; i < candidate_count; ++i) { + if (candidate_pages[i] == pid) { + duplicate = true; + break; + } + } + if (!duplicate) candidate_pages[candidate_count++] = pid; + } + + std::array load_pages{}; + size_t load_count = 0; + for (size_t i = 0; i < candidate_count; ++i) { + const block_id_t pid = candidate_pages[i]; + const auto claim = page_table_.try_claim_block_load(pid); + switch (claim) { + case VectorPageTable::LoadClaimResult::kClaimed: + load_pages[load_count++] = pid; + break; + case VectorPageTable::LoadClaimResult::kResident: + page_table_.promote_evict_priority(pid, priority); + break; + case VectorPageTable::LoadClaimResult::kLoading: + page_table_.promote_evict_priority(pid, priority); + singleflight_waits_.fetch_add(1, std::memory_order_relaxed); + break; + case VectorPageTable::LoadClaimResult::kEvicting: + page_table_.promote_evict_priority(pid, priority); + break; + } + } + if (load_count == 0) continue; + + std::array buffers{}; + size_t allocated = MemoryLimitPool::get_instance().batch_acquire_buffers( + kVectorPageSize, buffers.data(), load_count); + + // Admit into unused capacity first. Reclaim only the actual shortage so a + // miss batch cannot evict an equal number of resident pages while the + // shared pool still has room to grow. + while (allocated < load_count) { + const size_t shortage = load_count - allocated; + if (BlockEvictionQueue::get_instance().batch_recycle(shortage) == 0) { + break; + } + const size_t acquired = + MemoryLimitPool::get_instance().batch_acquire_buffers( + kVectorPageSize, buffers.data() + allocated, shortage); + allocated += acquired; + if (acquired == 0) { + break; + } + } + bool batch_ok = allocated == load_count; + for (size_t i = allocated; i < load_count; ++i) { + if (!page_table_.cancel_block_load(load_pages[i])) { + LOG_ERROR( + "VecBufferPool::load_pages_aio: failed to cancel unallocated " + "load claim for page=%zu", + static_cast(load_pages[i])); + } + } + if (allocated == 0) return false; + + auto abandon_claims = [&](size_t abandon_count, bool release_buffers) { + for (size_t i = 0; i < abandon_count; ++i) { + if (buffers[i]) { + if (release_buffers) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + } + buffers[i] = nullptr; + } + if (!page_table_.cancel_block_load(load_pages[i])) { + LOG_ERROR( + "VecBufferPool::load_pages_aio: failed to abandon load claim " + "for page=%zu", + static_cast(load_pages[i])); + } + } + }; + + std::array read_ok{}; + if (use_io_uring) { + std::array requests{}; + for (size_t i = 0; i < allocated; ++i) { + const size_t offset = load_pages[i] * kVectorPageSize; + const size_t expected = std::min(kVectorPageSize, file_size_ - offset); + requests[i] = + IoUringRead(offset, kVectorPageSize, buffers[i], expected); + } + aio_pages_submitted_.fetch_add(allocated, std::memory_order_relaxed); + if (tl_blocking_io_uring.ring.execute(fd_, requests.data(), allocated) == + 0) { + std::fill_n(read_ok.begin(), allocated, true); + } else if (tl_blocking_aio.ensure()) { + // io_uring uses ring-owned staging, so libaio can safely reuse the + // destination buffers after execute() reports a drained failure. + use_io_uring = false; + } else { + abandon_claims(allocated, /*release_buffers=*/true); + return false; + } + } + + if (!use_io_uring) { + std::array cbs{}; + std::array cb_ptrs{}; + for (size_t i = 0; i < allocated; ++i) { + const size_t offset = load_pages[i] * kVectorPageSize; + io_prep_pread(&cbs[i], fd_, buffers[i], kVectorPageSize, + static_cast(offset)); + cbs[i].data = reinterpret_cast(i); + cb_ptrs[i] = &cbs[i]; + } + + const int submit_ret = LibAioLoader::Instance().io_submit( + tl_blocking_aio.ctx, static_cast(allocated), cb_ptrs.data()); + if (submit_ret <= 0 || static_cast(submit_ret) > allocated) { + abandon_claims(allocated, /*release_buffers=*/true); + return false; + } + + const size_t accepted = static_cast(submit_ret); + aio_pages_submitted_.fetch_add(accepted, std::memory_order_relaxed); + batch_ok = batch_ok && accepted == allocated; + for (size_t i = accepted; i < allocated; ++i) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + buffers[i] = nullptr; + if (!page_table_.cancel_block_load(load_pages[i])) { + LOG_ERROR( + "VecBufferPool::load_pages_aio: failed to cancel unsubmitted " + "load claim for page=%zu", + static_cast(load_pages[i])); + } + } + + std::array events{}; + size_t completed = 0; + bool wait_failed = false; + while (completed < accepted) { + const int get_ret = LibAioLoader::Instance().io_getevents( + tl_blocking_aio.ctx, static_cast(accepted - completed), + static_cast(accepted - completed), events.data() + completed, + nullptr); + if (get_ret == -EINTR) continue; + if (get_ret <= 0) { + LOG_ERROR( + "VecBufferPool::load_pages_aio: io_getevents failed, ret=%d", + get_ret); + wait_failed = true; + break; + } + completed += static_cast(get_ret); + } + + if (wait_failed) { + if (tl_blocking_aio.destroy_context("ThreadLocalBlockingAioCtx")) { + abandon_claims(accepted, /*release_buffers=*/true); + } else { + tl_blocking_aio.quarantine(buffers.data(), accepted); + abandon_claims(accepted, /*release_buffers=*/false); + } + return false; + } + + std::array seen{}; + for (size_t i = 0; i < completed; ++i) { + const size_t idx = reinterpret_cast(events[i].data); + if (idx >= accepted || seen[idx] || buffers[idx] == nullptr) { + batch_ok = false; + continue; + } + seen[idx] = true; + const size_t offset = load_pages[idx] * kVectorPageSize; + const size_t expected = std::min(kVectorPageSize, file_size_ - offset); + read_ok[idx] = static_cast(events[i].res) == + static_cast(expected) && + events[i].res2 == 0; + if (!read_ok[idx]) { + batch_ok = false; + continue; + } + if (expected < kVectorPageSize) { + std::memset(buffers[idx] + expected, 0, kVectorPageSize - expected); + } + } + } + + for (size_t i = 0; i < allocated; ++i) { + if (buffers[i] == nullptr) continue; + if (!read_ok[i]) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + buffers[i] = nullptr; + if (!page_table_.cancel_block_load(load_pages[i])) { + LOG_ERROR( + "VecBufferPool::load_pages_aio: failed to cancel failed load " + "claim for page=%zu", + static_cast(load_pages[i])); + } + batch_ok = false; + continue; + } + const block_id_t pid = load_pages[i]; + miss_count_.fetch_add(1, std::memory_order_relaxed); + page_table_.promote_evict_priority(pid, priority); + char *installed = page_table_.publish_claimed_block( + pid, buffers[i], pid * kVectorPageSize); + if (installed != nullptr) { + page_table_.release_block(pid); + } else { + batch_ok = false; + } + buffers[i] = nullptr; + } + if (!batch_ok) return false; + } + return true; +#else + (void)page_ids; + (void)count; + (void)priority; + return false; +#endif +} + +void VecBufferPool::log_stats() const { + Stats s = stats(); + const auto resident = page_table_.resident_pages_by_priority(); + const auto queue_stats = BlockEvictionQueue::get_instance().stats(); + LOG_INFO( + "VecBufferPool stats: file[%s] hit=%llu miss=%llu hit_rate=%.4f " + "evict=%llu second_chance=%llu dirty_flush=%llu " + "writeback_requests=%llu writeback_batches=%llu " + "writeback_pages=%llu writeback_failures=%llu " + "writeback_aio_batches=%llu writeback_aio_pages=%llu " + "writeback_aio_fallbacks=%llu " + "writeback_waits=%llu writeback_wait_us=%llu " + "writeback_pending=%llu writeback_peak_pending=%llu " + "bypass_reads=%llu " + "bypass_bytes=%llu bypass_io_requests=%llu bypass_rechecks=%llu " + "bypass_cache_joins=%llu singleflight_waits=%llu " + "aio_pages_submitted=%llu " + "admission_admitted=%llu admission_rejected=%llu " + "ghost_hot_marks=%llu ghost_hot_hits=%llu " + "page_table_metadata_bytes=%zu page_lock_metadata_bytes=%zu " + "writeback_staging_bytes=%zu writeback_io_staging_bytes=%zu " + "resident_by_priority=[%zu,%zu,%zu] " + "promotions=[%llu,%llu,%llu] demotions=[%llu,%llu,%llu] " + "evictions_by_priority=[%llu,%llu,%llu] " + "global_queue_approx=[%zu,%zu,%zu] protected_aging_dequeues=%llu", + file_name_.c_str(), static_cast(s.hit), + static_cast(s.miss), s.hit_rate(), + static_cast(s.evict), + static_cast(s.second_chance), + static_cast(s.dirty_flush), + static_cast(s.writeback_requests), + static_cast(s.writeback_batches), + static_cast(s.writeback_pages), + static_cast(s.writeback_failures), + static_cast(s.writeback_aio_batches), + static_cast(s.writeback_aio_pages), + static_cast(s.writeback_aio_fallbacks), + static_cast(s.writeback_waits), + static_cast(s.writeback_wait_us), + static_cast(s.writeback_pending), + static_cast(s.writeback_peak_pending), + static_cast(s.bypass_reads), + static_cast(s.bypass_bytes), + static_cast(s.bypass_io_requests), + static_cast(s.bypass_rechecks), + static_cast(s.bypass_cache_joins), + static_cast(s.singleflight_waits), + static_cast(s.aio_pages_submitted), + static_cast(s.admission_admitted), + static_cast(s.admission_rejected), + static_cast(s.ghost_hot_marks), + static_cast(s.ghost_hot_hits), + s.page_table_metadata_bytes, s.page_lock_metadata_bytes, + s.writeback_staging_bytes, s.writeback_io_staging_bytes, resident[0], + resident[1], resident[2], + static_cast(s.priority_promotions[0]), + static_cast(s.priority_promotions[1]), + static_cast(s.priority_promotions[2]), + static_cast(s.priority_demotions[0]), + static_cast(s.priority_demotions[1]), + static_cast(s.priority_demotions[2]), + static_cast(s.evictions_by_priority[0]), + static_cast(s.evictions_by_priority[1]), + static_cast(s.evictions_by_priority[2]), + queue_stats.approximate_queue_sizes[0], + queue_stats.approximate_queue_sizes[1], + queue_stats.approximate_queue_sizes[2], + static_cast(queue_stats.protected_aging_dequeues)); +} + } // namespace ailego } // namespace zvec diff --git a/src/ailego/io/iouring_def.h b/src/ailego/io/iouring_def.h index 5ff099c18..a5923e596 100644 --- a/src/ailego/io/iouring_def.h +++ b/src/ailego/io/iouring_def.h @@ -193,6 +193,22 @@ static inline void io_uring_prep_read(struct io_uring_sqe *sqe, int fd, sqe->buf.personality = 0; } +static inline void io_uring_prep_write(struct io_uring_sqe *sqe, int fd, + const void *buf, uint32_t nbytes, + uint64_t offset) { + sqe->opcode = IORING_OP_WRITE; + sqe->flags = 0; + sqe->ioprio = 0; + sqe->fd = fd; + sqe->off = offset; + sqe->addr = reinterpret_cast(buf); + sqe->len = nbytes; + sqe->rw_flags = 0; + sqe->user_data = 0; + sqe->buf.buf_index = 0; + sqe->buf.personality = 0; +} + // --------------------------------------------------------------------------- // End: struct and constant definitions from // --------------------------------------------------------------------------- diff --git a/src/ailego/io/iouring_loader.cc b/src/ailego/io/iouring_loader.cc index ee25932b7..633f7fb55 100644 --- a/src/ailego/io/iouring_loader.cc +++ b/src/ailego/io/iouring_loader.cc @@ -16,13 +16,22 @@ #include // syscall(), __NR_io_uring_setup #include // close() +#include +#include #include +#include #include +#include +#include #include #include namespace zvec { -namespace core { +namespace ailego { + +// Retry budget for draining in-flight requests when the kernel keeps +// returning EAGAIN/EBUSY (100 us sleep per retry, about one second total). +static constexpr size_t kIoUringDrainRetries = 10000; bool IoUringRing::setup(uint32_t entries) { struct io_uring_params params; @@ -58,10 +67,10 @@ bool IoUringRing::setup(uint32_t entries) { // --- mmap the three shared regions --- // 1. SQ ring (includes head, tail, mask, entries, flags, dropped, array). - size_t sq_ring_sz = + sq_ring_size_ = static_cast(params.sq_off.array) + sq_entries_ * sizeof(uint32_t); - sq_ring_ptr_ = ::mmap(nullptr, sq_ring_sz, PROT_READ | PROT_WRITE, MAP_SHARED, - ring_fd_, IORING_OFF_SQ_RING); + sq_ring_ptr_ = ::mmap(nullptr, sq_ring_size_, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_SQ_RING); if (sq_ring_ptr_ == MAP_FAILED) { LOG_ERROR("mmap SQ ring failed: %s", ::strerror(errno)); sq_ring_ptr_ = nullptr; @@ -70,9 +79,9 @@ bool IoUringRing::setup(uint32_t entries) { } // 2. SQE array. - size_t sqes_sz = sq_entries_ * sizeof(struct io_uring_sqe); + sqes_size_ = sq_entries_ * sizeof(struct io_uring_sqe); sqes_ptr_ = reinterpret_cast( - ::mmap(nullptr, sqes_sz, PROT_READ | PROT_WRITE, MAP_SHARED, ring_fd_, + ::mmap(nullptr, sqes_size_, PROT_READ | PROT_WRITE, MAP_SHARED, ring_fd_, IORING_OFF_SQES)); if (sqes_ptr_ == MAP_FAILED) { LOG_ERROR("mmap SQEs failed: %s", ::strerror(errno)); @@ -82,10 +91,10 @@ bool IoUringRing::setup(uint32_t entries) { } // 3. CQ ring (includes head, tail, mask, entries, overflow, cqes[]). - size_t cq_ring_sz = static_cast(params.cq_off.cqes) + - cq_entries_ * sizeof(struct io_uring_cqe); - cq_ring_ptr_ = ::mmap(nullptr, cq_ring_sz, PROT_READ | PROT_WRITE, MAP_SHARED, - ring_fd_, IORING_OFF_CQ_RING); + cq_ring_size_ = static_cast(params.cq_off.cqes) + + cq_entries_ * sizeof(struct io_uring_cqe); + cq_ring_ptr_ = ::mmap(nullptr, cq_ring_size_, PROT_READ | PROT_WRITE, + MAP_SHARED, ring_fd_, IORING_OFF_CQ_RING); if (cq_ring_ptr_ == MAP_FAILED) { LOG_ERROR("mmap CQ ring failed: %s", ::strerror(errno)); cq_ring_ptr_ = nullptr; @@ -139,25 +148,21 @@ bool IoUringRing::setup(uint32_t entries) { void IoUringRing::teardown() { if (sq_ring_ptr_ && sq_ring_ptr_ != MAP_FAILED) { - // We don't track the exact mmap size; munmap with a large enough size - // is safe because the kernel only unmaps what was actually mapped. - // However, to be correct we use the page-aligned size. - size_t sz = static_cast(sq_entries_) * sizeof(uint32_t) + 4096; - ::munmap(sq_ring_ptr_, sz); + ::munmap(sq_ring_ptr_, sq_ring_size_); } if (sqes_ptr_ && sqes_ptr_ != MAP_FAILED) { - size_t sz = static_cast(sq_entries_) * sizeof(struct io_uring_sqe); - ::munmap(sqes_ptr_, sz); + ::munmap(sqes_ptr_, sqes_size_); } if (cq_ring_ptr_ && cq_ring_ptr_ != MAP_FAILED) { - size_t sz = - static_cast(cq_entries_) * sizeof(struct io_uring_cqe) + 4096; - ::munmap(cq_ring_ptr_, sz); + ::munmap(cq_ring_ptr_, cq_ring_size_); } sq_ring_ptr_ = nullptr; sqes_ptr_ = nullptr; cq_ring_ptr_ = nullptr; + sq_ring_size_ = 0; + sqes_size_ = 0; + cq_ring_size_ = 0; sqes_ = nullptr; cqes_ = nullptr; sq_head_ = sq_tail_ = sq_ring_mask_ = sq_ring_entries_ = nullptr; @@ -195,7 +200,206 @@ bool IoUringRing::ensure_staging(size_t bytes) { return true; } -} // namespace core +int IoUringRing::execute(int fd, const IoUringRead *read_reqs, size_t count) { + return execute_impl(fd, read_reqs, nullptr, count); +} + +int IoUringRing::execute_writes(int fd, const IoUringWrite *write_reqs, + size_t count) { + return execute_impl(fd, nullptr, write_reqs, count); +} + +int IoUringRing::execute_impl(int fd, const IoUringRead *read_reqs, + const IoUringWrite *write_reqs, size_t count) { + const bool is_write = write_reqs != nullptr; + if (!is_valid() || + (count != 0 && ((read_reqs == nullptr) == (write_reqs == nullptr)))) { + return -1; + } + if (count == 0) { + return 0; + } + + const size_t batch_size = std::min(sq_entries_, kIoUringMaxBatch); + if (batch_size == 0) { + return -1; + } + + for (size_t batch_start = 0; batch_start < count; batch_start += batch_size) { + const size_t n_ops = std::min(batch_size, count - batch_start); + std::array slot_offsets{}; + size_t staging_bytes = 0; + for (size_t j = 0; j < n_ops; ++j) { + const size_t req_idx = batch_start + j; + const uint64_t offset = + is_write ? write_reqs[req_idx].offset : read_reqs[req_idx].offset; + const uint64_t len = + is_write ? write_reqs[req_idx].len : read_reqs[req_idx].len; + const uint64_t expected_len = + is_write ? len + : (read_reqs[req_idx].expected_len == 0 + ? len + : read_reqs[req_idx].expected_len); + const void *buf = + is_write ? write_reqs[req_idx].buf : read_reqs[req_idx].buf; + if (buf == nullptr || len == 0 || expected_len > len || + len > std::numeric_limits::max() || offset % 512 != 0 || + len % 512 != 0 || + (!is_write && reinterpret_cast(buf) % 512 != 0)) { + return -1; + } + const size_t aligned_len = + (static_cast(len) + kIoUringStagingAlign - 1) & + ~(kIoUringStagingAlign - 1); + if (aligned_len < len || + staging_bytes > std::numeric_limits::max() - aligned_len) { + return -1; + } + slot_offsets[j] = staging_bytes; + staging_bytes += aligned_len; + } + if (!ensure_staging(staging_bytes)) { + return -1; + } + + const unsigned tail = __atomic_load_n(sq_tail_, __ATOMIC_ACQUIRE); + const unsigned mask = *sq_ring_mask_; + for (size_t j = 0; j < n_ops; ++j) { + const unsigned idx = (tail + static_cast(j)) & mask; + const unsigned sqe_idx = sq_array_[idx]; + struct io_uring_sqe *sqe = &sqes_[sqe_idx]; + const size_t req_idx = batch_start + j; + if (is_write) { + const IoUringWrite &req = write_reqs[req_idx]; + std::memcpy(staging_ + slot_offsets[j], req.buf, req.len); + io_uring_prep_write(sqe, fd, staging_ + slot_offsets[j], + static_cast(req.len), req.offset); + } else { + const IoUringRead &req = read_reqs[req_idx]; + io_uring_prep_read(sqe, fd, staging_ + slot_offsets[j], + static_cast(req.len), req.offset); + } + sqe->user_data = req_idx; + } + + __sync_synchronize(); + __atomic_store_n(sq_tail_, tail + static_cast(n_ops), + __ATOMIC_RELEASE); + + size_t submitted = 0; + size_t completed = 0; + bool all_ok = true; + auto reap_available = [&]() { + unsigned chead = *cq_head_; + const unsigned ctail = __atomic_load_n(cq_tail_, __ATOMIC_ACQUIRE); + const unsigned cq_mask = *cq_ring_mask_; + while (chead != ctail) { + struct io_uring_cqe *cqe = &cqes_[chead & cq_mask]; + const size_t req_idx = static_cast(cqe->user_data); + if (req_idx < batch_start || req_idx >= batch_start + n_ops) { + LOG_WARN("io_uring completion referenced unknown request: %zu", + req_idx); + all_ok = false; + } else { + const uint64_t offset = + is_write ? write_reqs[req_idx].offset : read_reqs[req_idx].offset; + const uint64_t len = + is_write ? write_reqs[req_idx].len : read_reqs[req_idx].len; + const uint64_t expected_len = + is_write ? len + : (read_reqs[req_idx].expected_len == 0 + ? len + : read_reqs[req_idx].expected_len); + const char *operation = is_write ? "write" : "read"; + if (cqe->res < 0) { + LOG_WARN("io_uring %s failed: req=%zu, res=%d, offset=%lu", + operation, req_idx, cqe->res, + static_cast(offset)); + all_ok = false; + } else if (static_cast(cqe->res) != expected_len) { + LOG_WARN("io_uring short %s: req=%zu, got=%d, expected=%lu", + operation, req_idx, cqe->res, + static_cast(expected_len)); + all_ok = false; + } else if (!is_write) { + const IoUringRead &req = read_reqs[req_idx]; + const size_t slot = req_idx - batch_start; + std::memcpy(req.buf, staging_ + slot_offsets[slot], expected_len); + if (expected_len < len) { + std::memset(static_cast(req.buf) + expected_len, 0, + len - expected_len); + } + } + } + ++chead; + ++completed; + } + __atomic_store_n(cq_head_, chead, __ATOMIC_RELEASE); + }; + + while (completed < n_ops) { + reap_available(); + if (completed >= n_ops) { + break; + } + + const unsigned to_submit = static_cast(n_ops - submitted); + const int ret = static_cast(syscall( + __NR_io_uring_enter, ring_fd_, to_submit, 1u, IORING_ENTER_GETEVENTS, + static_cast(nullptr), static_cast(0))); + if (ret >= 0) { + submitted += static_cast(ret); + continue; + } + if (errno == EINTR || + ((errno == EAGAIN || errno == EBUSY) && completed < submitted)) { + continue; + } + + LOG_WARN( + "io_uring_enter failed; errno=%d, %s, submitted=%zu/%zu, " + "completed=%zu. draining before falling back to p%s", + errno, ::strerror(errno), submitted, n_ops, completed, + is_write ? "write" : "read"); + __atomic_store_n(sq_tail_, tail + static_cast(submitted), + __ATOMIC_RELEASE); + + size_t drain_retries = 0; + while (completed < submitted) { + reap_available(); + if (completed >= submitted) { + break; + } + const int wait_ret = static_cast(syscall( + __NR_io_uring_enter, ring_fd_, 0u, 1u, IORING_ENTER_GETEVENTS, + static_cast(nullptr), static_cast(0))); + if (wait_ret >= 0 || errno == EINTR) { + continue; + } + if ((errno == EAGAIN || errno == EBUSY) && + drain_retries++ < kIoUringDrainRetries) { + std::this_thread::sleep_for(std::chrono::microseconds(100)); + continue; + } + LOG_ERROR( + "io_uring drain failed; errno=%d, %s. leaking the staging pool " + "and disabling io_uring for this context", + errno, ::strerror(errno)); + abandon_staging(); + teardown(); + return -1; + } + return -1; + } + + if (!all_ok) { + return -1; + } + } + return 0; +} + +} // namespace ailego } // namespace zvec #endif // __linux__ diff --git a/src/ailego/io/iouring_loader.h b/src/ailego/io/iouring_loader.h index a229de2f8..8521070f1 100644 --- a/src/ailego/io/iouring_loader.h +++ b/src/ailego/io/iouring_loader.h @@ -27,16 +27,57 @@ #include #include +#include +#include +#include #include #include #include namespace zvec { -namespace core { +namespace ailego { + +// Generic aligned read request shared by DiskANN and Buffer Storage. +// expected_len may be smaller than len for the final page of a file: the +// kernel still receives an O_DIRECT-aligned length, while execute() accepts +// the expected short read and zero-fills the rest of the destination. +struct IoUringRead { + uint64_t offset{0}; + uint64_t len{0}; + void *buf{nullptr}; + uint64_t expected_len{0}; + + IoUringRead() = default; + + IoUringRead(uint64_t offset, uint64_t len, void *buf, + uint64_t expected_len = 0) + : offset(offset), + len(len), + buf(buf), + expected_len(expected_len == 0 ? len : expected_len) { + assert(static_cast(offset) % 512 == 0); + assert(static_cast(len) % 512 == 0); + assert(reinterpret_cast(buf) % 512 == 0); + assert(this->expected_len <= len); + } +}; + +// Generic aligned write request. execute_writes() copies every source into +// ring-owned staging before submission, so caller buffers are no longer +// referenced once the method returns, including on a drained failure. +struct IoUringWrite { + uint64_t offset{0}; + uint64_t len{0}; + const void *buf{nullptr}; -// AlignedRead lives in diskann_file_reader.h; a forward declaration -// suffices since execute() takes it by reference. -struct AlignedRead; + IoUringWrite() = default; + + IoUringWrite(uint64_t offset, uint64_t len, const void *buf) + : offset(offset), len(len), buf(buf) { + assert(static_cast(offset) % 512 == 0); + assert(static_cast(len) % 512 == 0); + } +}; // Max SQEs submitted per io_uring_enter() call. static constexpr uint32_t kIoUringMaxBatch = 128; @@ -80,19 +121,37 @@ class IoUringRing { return ring_fd_ >= 0; } - // Execute a batch of aligned reads via io_uring. Returns 0 on success, - // -1 on failure — the caller may always fall back to pread, since the - // kernel only writes into the staging pool. In diskann_file_reader.cc - // (AlignedRead is defined there). - int execute(int fd, std::vector &read_reqs); + // Execute a batch of aligned reads via io_uring. Returns 0 on success and + // -1 on failure. The caller may always fall back to pread because the + // kernel only writes into the ring-owned staging pool. + int execute(int fd, const IoUringRead *read_reqs, size_t count); + + int execute(int fd, const std::vector &read_reqs) { + return execute(fd, read_reqs.data(), read_reqs.size()); + } + + // Execute a batch of aligned writes via io_uring. The call blocks until all + // accepted requests have completed, but the requests execute concurrently + // in the kernel. Returns 0 on success and -1 on a safely drained failure. + int execute_writes(int fd, const IoUringWrite *write_reqs, size_t count); + + int execute_writes(int fd, const std::vector &write_reqs) { + return execute_writes(fd, write_reqs.data(), write_reqs.size()); + } private: + int execute_impl(int fd, const IoUringRead *read_reqs, + const IoUringWrite *write_reqs, size_t count); + int ring_fd_{-1}; // mmap'd region bases (needed for munmap). void *sq_ring_ptr_{nullptr}; struct io_uring_sqe *sqes_ptr_{nullptr}; void *cq_ring_ptr_{nullptr}; + size_t sq_ring_size_{0}; + size_t sqes_size_{0}; + size_t cq_ring_size_{0}; // SQ ring field pointers (into sq_ring_ptr_). unsigned *sq_head_{nullptr}; @@ -123,7 +182,7 @@ class IoUringRing { unsigned cq_entries_{0}; }; -} // namespace core +} // namespace ailego } // namespace zvec #endif // __linux__ diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 681901e34..d32c98295 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -70,9 +69,6 @@ static void log_diskann_io_backend(ailego::IOBackendType type) { typedef struct io_event io_event_t; typedef struct iocb iocb_t; -// Retry budget for draining in-flight io_uring requests when the kernel -// keeps returning EAGAIN/EBUSY (100 us sleep per retry, ~1 s total). -static constexpr size_t kIoUringDrainRetries = 10000; #endif void log_diskann_io_backend() { @@ -369,218 +365,6 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, #endif } -// --------------------------------------------------------------------------- -// IoUringRing::execute — defined here (not in iouring_loader.h) because it -// accesses AlignedRead members, and AlignedRead is defined in -// diskann_file_reader.h after iouring_loader.h is included. -// --------------------------------------------------------------------------- -#if (defined(__linux) || defined(__linux__)) -int IoUringRing::execute(int fd, std::vector &read_reqs) { - if (!is_valid()) { - return -1; - } - if (read_reqs.empty()) { - return 0; - } - - // Process in batches limited by the SQ ring size. - uint32_t batch_size = - std::min(sq_entries_, static_cast(kIoUringMaxBatch)); - uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), batch_size); - - for (uint64_t iter = 0; iter < iters; iter++) { - uint64_t n_ops = - std::min(static_cast(read_reqs.size()) - iter * batch_size, - static_cast(batch_size)); - - // --- Phase 1: Fill SQEs --- - // - // Reads land in the ring-owned staging pool, never in the caller's - // buffers. io_uring teardown is asynchronous — closing the ring fd - // only initiates cancellation — so the kernel may still write into - // request buffers after execute() has returned an error. Staging - // memory can simply be leaked in that case (abandon_staging()), while - // the caller's buffers stay safe to reuse or free. The copy-out below - // costs one sector-scale memcpy per read, negligible next to the I/O. - std::vector slot_off(n_ops); - size_t staging_bytes = 0; - for (uint64_t j = 0; j < n_ops; j++) { - slot_off[j] = staging_bytes; - size_t len = read_reqs[j + iter * batch_size].len; - // Round every slot up so each staging pointer stays O_DIRECT-legal. - staging_bytes += - (len + kIoUringStagingAlign - 1) & ~(kIoUringStagingAlign - 1); - } - // Safe: the previous batch is fully drained before we get here, so no - // in-flight request can reference the old pool being freed on growth. - if (!ensure_staging(staging_bytes)) { - return -1; // nothing submitted; pread fallback is safe - } - - unsigned tail = __atomic_load_n(sq_tail_, __ATOMIC_ACQUIRE); - unsigned mask = *sq_ring_mask_; - - for (uint64_t j = 0; j < n_ops; j++) { - unsigned idx = (tail + static_cast(j)) & mask; - unsigned sqe_idx = sq_array_[idx]; - struct io_uring_sqe *sqe = &sqes_[sqe_idx]; - - uint64_t req_idx = j + iter * batch_size; - io_uring_prep_read(sqe, fd, staging_ + slot_off[j], - static_cast(read_reqs[req_idx].len), - read_reqs[req_idx].offset); - // Store the request index so we can verify the completion. - sqe->user_data = req_idx; - } - - // Memory barrier: ensure SQE contents are visible before tail update. - __sync_synchronize(); - __atomic_store_n(sq_tail_, tail + static_cast(n_ops), - __ATOMIC_RELEASE); - - // --- Phase 2: Submit and reap completions --- - // - // io_uring_enter() returns the number of SQEs consumed, not the number - // of CQEs available. A partial submission returns before the wait - // phase, and a signal can interrupt the wait while preserving a - // positive submission count, so IORING_ENTER_GETEVENTS guarantees - // min_complete completions only when the call finishes normally. - // Completions must therefore be counted against cq_tail instead of - // assuming n_ops CQEs are ready. - uint64_t submitted = 0; - uint64_t completed = 0; - bool all_ok = true; - - // Consume every CQE the kernel has published so far and verify it. - // Completion order is unspecified, so use cqe->user_data to find the - // request instead of assuming submission order. - auto reap_available = [&]() { - unsigned chead = *cq_head_; // single consumer — plain load is enough - unsigned ctail = __atomic_load_n(cq_tail_, __ATOMIC_ACQUIRE); - unsigned cq_mask = *cq_ring_mask_; - if (chead == ctail) { - return; - } - while (chead != ctail) { - struct io_uring_cqe *cqe = &cqes_[chead & cq_mask]; - uint64_t req_idx = cqe->user_data; - - if (req_idx < iter * batch_size || - req_idx >= iter * batch_size + n_ops) { - LOG_WARN("io_uring completion referenced unknown request: %lu", - (unsigned long)req_idx); - all_ok = false; - } else if (cqe->res < 0) { - LOG_WARN("io_uring read failed: req=%lu, res=%d, offset=%lu", - (unsigned long)req_idx, cqe->res, - (unsigned long)read_reqs[req_idx].offset); - all_ok = false; - } else if (static_cast(cqe->res) != read_reqs[req_idx].len) { - LOG_WARN("io_uring short read: req=%lu, got=%d, expected=%lu", - (unsigned long)req_idx, cqe->res, - (unsigned long)read_reqs[req_idx].len); - all_ok = false; - } else { - // Verified completion — copy from staging into the caller's - // buffer. This is the only place caller memory is written. - std::memcpy(read_reqs[req_idx].buf, - staging_ + slot_off[req_idx - iter * batch_size], - read_reqs[req_idx].len); - } - chead++; - completed++; - } - // Release: CQE reads must complete before the kernel may reuse slots. - __atomic_store_n(cq_head_, chead, __ATOMIC_RELEASE); - }; - - while (completed < n_ops) { - reap_available(); - if (completed >= n_ops) { - break; - } - - unsigned to_submit = static_cast(n_ops - submitted); - int ret = static_cast(syscall( - __NR_io_uring_enter, ring_fd_, to_submit, 1u, IORING_ENTER_GETEVENTS, - static_cast(nullptr), static_cast(0))); - if (ret >= 0) { - submitted += static_cast(ret); - continue; - } - if (errno == EINTR) { - // Interrupted during submit or wait; the SQEs already consumed are - // tracked in `submitted`, so simply retry. - continue; - } - if ((errno == EAGAIN || errno == EBUSY) && completed < submitted) { - // Kernel resources are exhausted, but in-flight requests will free - // them as they complete; keep reaping and retrying. - continue; - } - - // Unrecoverable failure (or EAGAIN with nothing in flight). - LOG_WARN( - "io_uring_enter failed; errno=%d, %s, submitted=%lu/%lu, " - "completed=%lu. draining before falling back to pread", - errno, ::strerror(errno), (unsigned long)submitted, - (unsigned long)n_ops, (unsigned long)completed); - - // Un-publish the SQEs the kernel never consumed so a later batch - // cannot submit them against stale buffers. - __atomic_store_n(sq_tail_, tail + static_cast(submitted), - __ATOMIC_RELEASE); - - // Drain every in-flight request before the staging pool may be - // freed or reused by a later batch. CQEs are posted to the shared - // ring by the kernel on its own, so completions can still be reaped - // here even when io_uring_enter() keeps failing. - size_t drain_retries = 0; - while (completed < submitted) { - reap_available(); - if (completed >= submitted) { - break; - } - int wret = static_cast(syscall( - __NR_io_uring_enter, ring_fd_, 0u, 1u, IORING_ENTER_GETEVENTS, - static_cast(nullptr), static_cast(0))); - if (wret >= 0 || errno == EINTR) { - continue; - } - if ((errno == EAGAIN || errno == EBUSY) && - drain_retries++ < kIoUringDrainRetries) { - // Give in-flight requests time to complete; entering the kernel - // via the sleep also lets pending completion task-work run. - std::this_thread::sleep_for(std::chrono::microseconds(100)); - continue; - } - // The ring cannot be drained. Leak the staging pool — the kernel - // may keep writing into it through the asynchronous teardown — and - // disable io_uring for this context. The caller's buffers were - // never exposed to the kernel, so the pread fallback stays safe. - LOG_ERROR( - "io_uring drain failed; errno=%d, %s. leaking the staging pool " - "and disabling io_uring for this context", - errno, ::strerror(errno)); - abandon_staging(); - teardown(); - return -1; - } - return -1; - } - - if (!all_ok) { - // Every request completed and the staging pool is quiesced, but at - // least one read failed or was short — let the caller retry with - // pread. - return -1; - } - } - - return 0; -} -#endif // __linux__ - LinuxAlignedFileReader::LinuxAlignedFileReader(int file_desc) { this->file_desc = file_desc; } diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 809b68c7f..d26f83e48 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -50,7 +50,7 @@ struct IoBackend { ailego::IOBackendType type{ailego::IOBackendType::kPread}; #if (defined(__linux) || defined(__linux__)) - IoUringRing ring{}; + ailego::IoUringRing ring{}; io_context_t aio_ctx{nullptr}; #endif }; @@ -64,6 +64,9 @@ int destroy_io_ctx(IOContext &ctx); // backend on first call. No-op outside Linux and macOS. void log_diskann_io_backend(); +#if (defined(__linux) || defined(__linux__)) +using AlignedRead = ailego::IoUringRead; +#else struct AlignedRead { uint64_t offset; uint64_t len; @@ -81,6 +84,7 @@ struct AlignedRead { #endif } }; +#endif struct PendingBatch { #if (defined(__linux) || defined(__linux__)) diff --git a/src/include/zvec/ailego/buffer/block_eviction_queue.h b/src/include/zvec/ailego/buffer/block_eviction_queue.h index b93a62bd6..1e87b807d 100644 --- a/src/include/zvec/ailego/buffer/block_eviction_queue.h +++ b/src/include/zvec/ailego/buffer/block_eviction_queue.h @@ -15,31 +15,21 @@ #pragma once -#include -#include +#include #include -#include -#include -#include -#include -#include -#include -#include +#include +#include +#include +#include #include -#include #include -#include -#include -#include +#include #include +#include #include #include #include "concurrentqueue.h" -#if defined(_MSC_VER) -#include -#endif - namespace zvec { namespace ailego { @@ -53,40 +43,77 @@ class ZVEC_AILEGO_API EvictableBlockOwner { virtual bool is_dead_block(eviction_key_t owner_key, version_t version) = 0; - virtual void evict_block(eviction_key_t owner_key) = 0; + //! Evict a block; return true only when memory was reclaimed. + virtual bool evict_block(eviction_key_t owner_key) = 0; + + //! Current eviction-queue priority for an item. Owners that do not support + //! priority keep the default low-priority queue. + virtual uint8_t eviction_priority(eviction_key_t /*owner_key*/) const { + return 0; + } + + //! Clear persistent membership after a failed requeue. + virtual void eviction_requeue_failed(eviction_key_t /*owner_key*/, + version_t /*version*/) {} + + //! Recover missing queue entries after reclamation finds an empty queue. + virtual size_t recover_eviction_queue() { + return 0; + } }; class BlockEvictionQueue { public: + static constexpr size_t kQueueCount = 3; + static constexpr uint8_t kProbationPriority = 0; + static constexpr uint8_t kProtectedPriority = 1; + static constexpr uint8_t kExplicitHotPriority = 2; + struct BlockType { eviction_key_t owner_key{0}; version_t version{0}; EvictableBlockOwner *owner{nullptr}; + uint8_t priority{0}; }; typedef moodycamel::ConcurrentQueue ConcurrentQueue; static BlockEvictionQueue &get_instance() { - static BlockEvictionQueue instance; - return instance; + // Weak COMDAT storage keeps the singleton shared across loaded images. + static std::atomic instance{nullptr}; + BlockEvictionQueue *current = instance.load(std::memory_order_acquire); + if (current != nullptr) { + return *current; + } + + BlockEvictionQueue *created = new BlockEvictionQueue(); + if (instance.compare_exchange_strong(current, created, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return *created; + } + delete created; + return *current; } BlockEvictionQueue(const BlockEvictionQueue &) = delete; BlockEvictionQueue &operator=(const BlockEvictionQueue &) = delete; BlockEvictionQueue(BlockEvictionQueue &&) = delete; BlockEvictionQueue &operator=(BlockEvictionQueue &&) = delete; - int init(); - bool evict_single_block(BlockType &item); bool evict_block(BlockType &item); bool add_single_block(const BlockType &block, int queue_index); - // void clear_dead_node(); - - bool is_valid(EvictableBlockOwner *owner) { - std::shared_lock lock(valid_owners_mutex_); - return valid_owners_.find(owner) != valid_owners_.end(); + //! Return a non-zero generation that prevents owner-address ABA. + version_t next_version() { + version_t version = + version_sequence_.fetch_add(1, std::memory_order_relaxed); + // Zero is reserved for non-versioned owners. + if (ailego_unlikely(version == 0)) { + version = version_sequence_.fetch_add(1, std::memory_order_relaxed); + } + return version; } void set_valid(EvictableBlockOwner *owner) { @@ -99,31 +126,71 @@ class BlockEvictionQueue { valid_owners_.erase(owner); } - // Atomically checks under the shared lock that the owner is still valid AND - // the block version has not been superseded, preventing TOCTOU races when an - // owner is concurrently destroyed. - bool is_valid_and_alive(const BlockType &item); - void recycle(); + size_t batch_recycle(size_t count); + + struct Stats { + std::array approximate_queue_sizes{}; + uint64_t protected_aging_dequeues{0}; + }; + + Stats stats() const { + Stats result; + for (size_t i = 0; i < kQueueCount; ++i) { + result.approximate_queue_sizes[i] = + approximate_queue_sizes_[i].load(std::memory_order_relaxed); + } + result.protected_aging_dequeues = + protected_aging_dequeues_.load(std::memory_order_relaxed); + return result; + } + private: + bool evict_single_block(BlockType &item, bool age_protected); + + bool evict_block(BlockType &item, size_t &attempts, size_t max_attempts, + bool &age_protected); + + size_t recover_owner_queues(); + BlockEvictionQueue() { - init(); + for (size_t i = 0; i < kQueueCount; ++i) { + evict_queues_.push_back(ConcurrentQueue(kEvictQueueCapacity)); + } } - private: - constexpr static size_t CACHE_QUEUE_NUM = 3; - size_t evict_batch_size_{0}; + // Foreground one-page reclaim should not pay for protected aging. Larger + // reclaim batches inspect the protected queue at most once. + static constexpr size_t kEvictQueueCapacity = 512 * 200; + static constexpr size_t kProtectedAgingMinBatch = 8; + static constexpr size_t kProtectedDominanceRatio = 3; std::vector evict_queues_; std::unordered_set valid_owners_; std::shared_mutex valid_owners_mutex_; + std::atomic version_sequence_{1}; + std::array, kQueueCount> approximate_queue_sizes_{}; + std::atomic protected_aging_dequeues_{0}; }; class MemoryLimitPool { public: static MemoryLimitPool &get_instance() { - static MemoryLimitPool instance; - return instance; + // Retain process-wide state while any loaded image may reference it. + static std::atomic instance{nullptr}; + MemoryLimitPool *current = instance.load(std::memory_order_acquire); + if (current != nullptr) { + return *current; + } + + MemoryLimitPool *created = new MemoryLimitPool(); + if (instance.compare_exchange_strong(current, created, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + return *created; + } + delete created; + return *current; } MemoryLimitPool(const MemoryLimitPool &) = delete; MemoryLimitPool &operator=(const MemoryLimitPool &) = delete; @@ -134,20 +201,257 @@ class MemoryLimitPool { bool try_acquire_buffer(const size_t buffer_size, char *&buffer); - void charge_external(const size_t buffer_size); + //! Wait briefly for another cache owner to release enough logical budget. + //! This is used as writable-cache backpressure after eviction has been + //! requested; it never reserves the returned capacity. + bool wait_for_available(const size_t buffer_size, + std::chrono::milliseconds timeout); + + //! Reserve bounded capacity outside the page cache, evicting if needed. + bool try_charge_external(const size_t buffer_size); + + //! Reserve non-evictable buffer-pool metadata capacity. + bool try_charge_metadata(const size_t buffer_size); void release_buffer(char *buffer, const size_t buffer_size); void release_external(const size_t buffer_size); + void release_metadata(const size_t buffer_size); + bool is_full(); + //! Whether another cache page would exceed the page-specific admission + //! limit. This can become true below the process-wide hard cap because page + //! storage preserves headroom for external cache consumers. + bool is_page_full() const { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + const size_t used = used_size_.load(std::memory_order_relaxed); + if (capacity == 0) { + return used != 0; + } + const size_t reserve = page_admission_reserve(); + const size_t external = external_used_size_.load(std::memory_order_relaxed); + const size_t remaining_reserve = + reserve > external ? reserve - external : 0; + const size_t limit = capacity - remaining_reserve; + return used >= limit || limit - used < page_buffer_size(); + } + + //! Whether page admission should protect the current resident set. Uses the + //! background-reclaim low watermark so admission remains active after a + //! reclaim pass, and clears once meaningful headroom returns. + bool under_cache_pressure() const { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + if (capacity == 0) { + return false; + } + // Page admission can reach its reserved-headroom limit before the shared + // pool reaches the general low watermark. Admission control must already + // be active at that point, otherwise foreground misses churn the resident + // set while is_full() still reports false. + if (page_admission_reserve() != 0 && is_page_full()) { + return true; + } + const size_t used = used_size_.load(std::memory_order_relaxed); + if (used >= capacity || capacity - used < page_buffer_size()) { + return true; + } + if (used <= fixed_used()) { + return false; + } + const size_t low = low_watermark(); + return used >= low || low - used <= page_buffer_size(); + } + + //! Lock-free estimate of currently available bytes. + size_t available() const { + size_t used = used_size_.load(std::memory_order_relaxed); + size_t capacity = pool_size_.load(std::memory_order_relaxed); + return (used >= capacity) ? 0 : (capacity - used); + } + + //! Headroom kept for shared non-page consumers (for example decoded + //! Parquet columns). Tiny test/application pools retain their historical + //! ability to use every page-sized byte. + size_t page_admission_reserve() const { + const size_t capacity = pool_size_.load(std::memory_order_relaxed); + if (capacity < (256UL << 20)) { + return 0; + } + return std::min(32UL << 20, capacity / 16); + } + + size_t batch_acquire_buffers(size_t buffer_size, char **out, size_t count); + + //! Current bytes in use (atomic, lock-free). + size_t used() const { + return used_size_.load(std::memory_order_relaxed); + } + + //! Bytes physically retained by page buffers plus external reservations. + size_t committed() const { + return committed_size_.load(std::memory_order_relaxed); + } + + //! Bytes reserved by shared-cache consumers outside VecBufferPool pages, + //! such as decoded or application-level cache entries. + size_t external_used() const { + return external_used_size_.load(std::memory_order_relaxed); + } + + size_t metadata_used() const { + return metadata_used_size_.load(std::memory_order_relaxed); + } + + //! Current configured capacity in bytes. + size_t capacity() const { + return pool_size_.load(std::memory_order_relaxed); + } + + //! Whether init() published a capacity, including zero. + bool initialized() const { + return initialized_.load(std::memory_order_acquire); + } + + //! VecBufferPool pages are carved from 4 MiB-aligned virtual mappings. + static size_t page_buffer_size(); + static constexpr size_t slab_size() { + return 4UL << 20; + } + static constexpr size_t slab_alignment() { + return slab_size(); + } + + //! Snapshot of pool-level counters for monitoring / export. + struct PoolStats { + size_t pool_size{0}; + size_t used{0}; + size_t committed{0}; + size_t page_used{0}; + size_t external_used{0}; + size_t metadata_used{0}; + size_t free_buffers{0}; // buffers cached across all shards + size_t slab_count{0}; // live 4 MiB virtual mappings + size_t slab_mapped_bytes{0}; // virtual address space held by slabs + size_t slab_header_bytes{0}; // resident-capable allocator metadata + uint64_t alloc_from_freelist{0}; // acquisitions served from a shard + uint64_t alloc_from_slab{0}; // cold page acquisitions from slabs + uint64_t slab_reclaimed_pages{0}; // pages decommitted under pressure + uint64_t bg_evict_rounds{0}; // background reclaim passes + uint64_t bg_evicted_buffers{0}; // buffers reclaimed by background thread + uint64_t bg_no_progress_sleeps{0}; // backoffs after zero-page reclaim + uint64_t high_watermark_hits{0}; // foreground acquire hit the capacity cap + uint64_t capacity_waits{0}; // waits for a released budget slot + uint64_t capacity_wait_timeouts{0}; // waits that observed no release + }; + PoolStats stats() const; + void log_stats() const; + private: + struct ReclaimableSlab; + MemoryLimitPool() = default; + ~MemoryLimitPool(); + + void drain_free_list(); + bool try_reserve_used(size_t bytes); + bool try_reserve_page_used(size_t bytes); + bool try_reserve_committed(size_t bytes); + bool try_charge_fixed(size_t bytes, std::atomic *counter); + void release_fixed(size_t bytes, std::atomic *counter); + bool is_cacheable_buffer_size(size_t buffer_size); + char *pop_free_buffer(size_t start_shard); + void push_free_buffer(char *buffer, size_t shard); + char *acquire_slab_buffer(); + bool reclaim_slab_buffer(char *buffer); + void release_all_slabs_locked(); + size_t trim_free_buffers(size_t bytes_needed); + size_t pick_shard(); + + // Reclaim in the background to keep eviction off foreground allocations. + void start_background_evictor(); + void stop_background_evictor(); + void background_evict_loop(); + // Keep a small absolute reserve without shrinking a fitting working set. + size_t reserve_margin(size_t capacity) const { + size_t m = capacity / 64; // ~1.5% of the pool + const size_t lo = 8UL << 20; // but at least 8 MB + const size_t hi = 64UL << 20; // and at most 64 MB + if (m < lo) m = lo; + if (m > hi) m = hi; + if (m * 2 >= capacity) m = capacity / 8; // tiny pools: fall back + return m; + } + size_t high_watermark() const { + size_t capacity = pool_size_.load(std::memory_order_relaxed); + const size_t fixed = fixed_used(); + const size_t page_capacity = capacity > fixed ? capacity - fixed : 0; + return fixed + page_capacity - reserve_margin(page_capacity); + } + size_t low_watermark() const { + size_t capacity = pool_size_.load(std::memory_order_relaxed); + const size_t fixed = fixed_used(); + const size_t page_capacity = capacity > fixed ? capacity - fixed : 0; + return fixed + page_capacity - reserve_margin(page_capacity) * 2; + } + size_t fixed_used() const { + return external_used_size_.load(std::memory_order_relaxed) + + metadata_used_size_.load(std::memory_order_relaxed); + } + bool should_background_reclaim() const { + const size_t used = used_size_.load(std::memory_order_relaxed); + const size_t fixed = fixed_used(); + return used > fixed && used >= high_watermark(); + } - private: - size_t pool_size_{0}; + // Shard the aligned free list to reduce allocation-path contention. + static constexpr size_t kNumFreeShards = 64; + struct alignas(64) FreeShard { + std::mutex mutex; + char *head{nullptr}; + std::atomic count{0}; + }; + + // init() may publish capacity while monitoring reads it. + std::atomic pool_size_{0}; + std::atomic initialized_{false}; + // Serialize reinitialization with reservations, but not releases. + mutable std::shared_mutex lifecycle_mutex_; std::atomic used_size_{0}; + // Includes free-list buffers so all consumers share the same hard cap. + std::atomic committed_size_{0}; + std::atomic external_used_size_{0}; + std::atomic metadata_used_size_{0}; + + FreeShard free_shards_[kNumFreeShards]; + std::atomic shard_seq_{0}; + + // Cold allocations are serialized; the page reuse hot path remains sharded. + std::mutex slab_mutex_; + ReclaimableSlab *slabs_{nullptr}; + ReclaimableSlab *allocation_slab_{nullptr}; + std::atomic slab_count_{0}; + std::atomic slab_mapped_bytes_{0}; + + // Observability counters (relaxed atomics; statistics only). + std::atomic alloc_from_freelist_{0}; + std::atomic alloc_from_slab_{0}; + std::atomic slab_reclaimed_pages_{0}; + std::atomic bg_evict_rounds_{0}; + std::atomic bg_evicted_buffers_{0}; + std::atomic bg_no_progress_sleeps_{0}; + std::atomic high_watermark_hits_{0}; + std::atomic capacity_waits_{0}; + std::atomic capacity_wait_timeouts_{0}; + + std::mutex capacity_mutex_; + std::condition_variable capacity_cv_; + + std::thread bg_thread_; + std::atomic bg_running_{false}; + std::mutex bg_mutex_; + std::condition_variable bg_cv_; }; } // namespace ailego diff --git a/src/include/zvec/ailego/buffer/external_cache.h b/src/include/zvec/ailego/buffer/external_cache.h index 671ef4564..a30167e4b 100644 --- a/src/include/zvec/ailego/buffer/external_cache.h +++ b/src/include/zvec/ailego/buffer/external_cache.h @@ -14,11 +14,16 @@ #pragma once +#include #include +#include #include #include +#include #include +#include #include +#include #include #include #include @@ -31,17 +36,32 @@ template (1, max_concurrent_loads)) { BlockEvictionQueue::get_instance().set_valid(this); } - explicit ExternalCache(Loader loader) : loader_(std::move(loader)) { + explicit ExternalCache( + Loader loader, size_t max_concurrent_loads = kDefaultMaxConcurrentLoads) + : loader_(std::move(loader)), + max_concurrent_loads_(std::max(1, max_concurrent_loads)) { BlockEvictionQueue::get_instance().set_valid(this); } ~ExternalCache() { BlockEvictionQueue::get_instance().set_invalid(this); + std::unique_lock lock(mutex_); + for (auto &item : table_) { + Entry &entry = item.second; + if (entry.size != 0) { + MemoryLimitPool::get_instance().release_external(entry.size); + entry.size = 0; + } + clear_noexcept(entry.payload); + } } ExternalCache(const ExternalCache &) = delete; @@ -61,10 +81,7 @@ class ExternalCache : public EvictableBlockOwner { } } - if (!ensure_capacity()) { - return Value{}; - } - + // Join an in-flight load before attempting reclamation. std::unique_lock lock(mutex_); auto iter = table_.find(key); if (iter != table_.end()) { @@ -72,24 +89,137 @@ class ExternalCache : public EvictableBlockOwner { if (value) { return value; } - } else { + if (iter->second.loading) { + return wait_for_loading(key, lock); + } + } + + // Bound concurrent loaders because payload size is unknown up front. + while (true) { + iter = table_.find(key); + if (iter != table_.end()) { + Value value = acquire_loaded(iter->second); + if (value) { + return value; + } + if (iter->second.loading) { + return wait_for_loading(key, lock); + } + } + if (active_loads_ < max_concurrent_loads_) { + break; + } + loader_slot_cv_.wait(lock); + } + + if (iter == table_.end()) { auto inserted = table_.try_emplace(key); iter = inserted.first; - iter->second.owner_key = next_owner_key_++; - owner_keys_.emplace(iter->second.owner_key, key); + iter->second.owner_key = next_owner_key_; + try { + owner_keys_.emplace(iter->second.owner_key, key); + iter->second.load_state = std::make_shared(); + } catch (...) { + owner_keys_.erase(iter->second.owner_key); + table_.erase(iter); + throw; + } + ++next_owner_key_; + } else { + iter->second.load_state = std::make_shared(); } + iter->second.loading = true; + ++active_loads_; + std::shared_ptr claimed_load_state = iter->second.load_state; - Entry &entry = iter->second; + // Load outside the cache lock, then reserve before publication. + lock.unlock(); + // Recheck capacity after waiting for a loader slot. + if (!ensure_capacity()) { + finish_loading(key); + return Value{}; + } + + std::optional loaded_payload; size_t size = 0; - if (!loader_.load(key, entry.payload, size)) { + bool loaded = false; + try { + // Keep construction inside the guarded completion path. + loaded_payload.emplace(); + loaded = loader_.load(key, *loaded_payload, size); + } catch (...) { + if (loaded_payload) { + clear_noexcept(*loaded_payload); + } + finish_loading(key); + throw; + } + if (!loaded) { + clear_noexcept(*loaded_payload); + finish_loading(key); + return Value{}; + } + + if (!MemoryLimitPool::get_instance().try_charge_external(size)) { + clear_noexcept(*loaded_payload); + finish_loading(key); return Value{}; } - entry.size = size; - MemoryLimitPool::get_instance().charge_external(entry.size); - entry.generation.fetch_add(1, std::memory_order_relaxed); - entry.ref_count.store(1, std::memory_order_release); - return loader_.value(entry.payload); + lock.lock(); + iter = table_.find(key); + if (iter == table_.end()) { + MemoryLimitPool::get_instance().release_external(size); + clear_noexcept(*loaded_payload); + claimed_load_state->complete = true; + if (active_loads_ != 0) { + --active_loads_; + } + lock.unlock(); + claimed_load_state->cv.notify_all(); + loader_slot_cv_.notify_one(); + return Value{}; + } + Entry &entry = iter->second; + assert(entry.loading && entry.load_state == claimed_load_state); + + Value value; + try { + entry.payload = std::move(*loaded_payload); + entry.size = size; + entry.generation.store(BlockEvictionQueue::get_instance().next_version(), + std::memory_order_relaxed); + entry.in_evict_queue.store(false, std::memory_order_relaxed); + entry.ref_count.store(1, std::memory_order_release); + value = loader_.value(entry.payload); + if (!value) { + throw std::runtime_error( + "ExternalCache loader returned an empty value after a " + "successful load"); + } + } catch (...) { + entry.ref_count.store(std::numeric_limits::min(), + std::memory_order_release); + MemoryLimitPool::get_instance().release_external(size); + entry.size = 0; + clear_noexcept(entry.payload); + clear_noexcept(*loaded_payload); + finish_loading_locked(iter, lock); + throw; + } + entry.loading = false; + claimed_load_state->complete = true; + --active_loads_; + lock.unlock(); + claimed_load_state->cv.notify_all(); + loader_slot_cv_.notify_one(); + return value; + } + + //! Number of live or in-flight cache entries. + size_t entry_count() { + std::shared_lock lock(mutex_); + return table_.size(); } Value retain(const Key &key) { @@ -110,12 +240,24 @@ class ExternalCache : public EvictableBlockOwner { Entry &entry = iter->second; if (entry.ref_count.fetch_sub(1, std::memory_order_release) == 1) { + bool expected = false; + if (!entry.in_evict_queue.compare_exchange_strong( + expected, true, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + return; + } std::atomic_thread_fence(std::memory_order_acquire); BlockEvictionQueue::BlockType block; block.owner = this; block.owner_key = entry.owner_key; block.version = entry.generation.load(std::memory_order_relaxed); - BlockEvictionQueue::get_instance().add_single_block(block, 0); + // Do not call the failure callback while holding mutex_: reclaiming an + // unqueued entry needs the exclusive side of the same mutex. + lock.unlock(); + if (!BlockEvictionQueue::get_instance().add_single_block( + block, BlockEvictionQueue::kExplicitHotPriority)) { + eviction_requeue_failed(block.owner_key, block.version); + } } } @@ -130,40 +272,116 @@ class ExternalCache : public EvictableBlockOwner { if (iter == table_.end()) { return true; } - return iter->second.generation.load(std::memory_order_relaxed) != version; + const Entry &entry = iter->second; + return entry.generation.load(std::memory_order_relaxed) != version || + !entry.in_evict_queue.load(std::memory_order_relaxed); + } + + uint8_t eviction_priority(eviction_key_t /*owner_key*/) const override { + // Reloading and decoding an Arrow/Parquet column is far more expensive + // than a vector-page read. Keep published payloads behind both ordinary + // and protected pages; the global queue can still reclaim them when no + // lower-priority resident page remains. + return BlockEvictionQueue::kExplicitHotPriority; } - void evict_block(eviction_key_t owner_key) override { + bool evict_block(eviction_key_t owner_key) override { std::unique_lock lock(mutex_); auto key_iter = owner_keys_.find(owner_key); if (key_iter == owner_keys_.end()) { - return; + return false; + } + + auto iter = table_.find(key_iter->second); + if (iter == table_.end()) { + return false; + } + + Entry &entry = iter->second; + int expected = 0; + if (entry.ref_count.compare_exchange_strong( + expected, std::numeric_limits::min())) { + entry.in_evict_queue.store(false, std::memory_order_relaxed); + MemoryLimitPool::get_instance().release_external(entry.size); + entry.size = 0; + clear_noexcept(entry.payload); + if (!entry.loading) { + owner_keys_.erase(key_iter); + table_.erase(iter); + } + return true; } + // Move a pinned entry to the tail without duplicating membership. + if (expected >= 0) { + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = entry.owner_key; + block.version = entry.generation.load(std::memory_order_relaxed); + if (!BlockEvictionQueue::get_instance().add_single_block( + block, BlockEvictionQueue::kExplicitHotPriority)) { + entry.in_evict_queue.store(false, std::memory_order_relaxed); + } + } else { + entry.in_evict_queue.store(false, std::memory_order_relaxed); + } + return false; + } + + void eviction_requeue_failed(eviction_key_t owner_key, + version_t version) override { + std::unique_lock lock(mutex_); + auto key_iter = owner_keys_.find(owner_key); + if (key_iter == owner_keys_.end()) { + return; + } auto iter = table_.find(key_iter->second); if (iter == table_.end()) { return; } Entry &entry = iter->second; + if (entry.generation.load(std::memory_order_relaxed) != version) { + return; + } + bool queued = true; + if (!entry.in_evict_queue.compare_exchange_strong( + queued, false, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + return; + } int expected = 0; if (entry.ref_count.compare_exchange_strong( expected, std::numeric_limits::min())) { MemoryLimitPool::get_instance().release_external(entry.size); entry.size = 0; - loader_.clear(entry.payload); + clear_noexcept(entry.payload); + if (!entry.loading) { + owner_keys_.erase(key_iter); + table_.erase(iter); + } } } private: + struct LoadState { + std::condition_variable_any cv; + bool complete{false}; + }; + struct Entry { Payload payload{}; size_t size{0}; eviction_key_t owner_key{0}; + bool loading{false}; + std::shared_ptr load_state{}; alignas(64) std::atomic ref_count{std::numeric_limits::min()}; alignas(64) std::atomic generation{0}; + std::atomic in_evict_queue{false}; }; + using Table = std::unordered_map; + Value acquire_loaded(Entry &entry) { while (true) { int current_count = entry.ref_count.load(std::memory_order_acquire); @@ -173,14 +391,83 @@ class ExternalCache : public EvictableBlockOwner { if (entry.ref_count.compare_exchange_weak( current_count, current_count + 1, std::memory_order_acq_rel, std::memory_order_acquire)) { - if (current_count == 0) { - entry.generation.fetch_add(1, std::memory_order_relaxed); + bool pin_rolled_back = false; + try { + Value value = loader_.value(entry.payload); + if (!value) { + entry.ref_count.fetch_sub(1, std::memory_order_release); + pin_rolled_back = true; + throw std::runtime_error( + "ExternalCache loader returned an empty cached value"); + } + return value; + } catch (...) { + if (!pin_rolled_back) { + entry.ref_count.fetch_sub(1, std::memory_order_release); + } + throw; } - return loader_.value(entry.payload); } } } + Value wait_for_loading(const Key &key, + std::unique_lock &lock) { + auto iter = table_.find(key); + assert(iter != table_.end() && iter->second.loading); + std::shared_ptr load_state = iter->second.load_state; + load_state->cv.wait(lock, [&load_state] { return load_state->complete; }); + iter = table_.find(key); + if (iter == table_.end()) { + return Value{}; + } + return acquire_loaded(iter->second); + } + + //! Complete a failed load and remove its unpublished placeholder. + void finish_loading(const Key &key) { + std::unique_lock lock(mutex_); + auto iter = table_.find(key); + if (iter == table_.end()) { + if (active_loads_ != 0) { + --active_loads_; + } + lock.unlock(); + loader_slot_cv_.notify_one(); + return; + } + finish_loading_locked(iter, lock); + } + + void finish_loading_locked(typename Table::iterator iter, + std::unique_lock &lock) { + Entry &entry = iter->second; + assert(entry.loading); + std::shared_ptr load_state = entry.load_state; + entry.loading = false; + load_state->complete = true; + if (active_loads_ != 0) { + --active_loads_; + } + if (entry.size == 0 && + entry.ref_count.load(std::memory_order_relaxed) < 0 && + !entry.in_evict_queue.load(std::memory_order_relaxed)) { + owner_keys_.erase(entry.owner_key); + table_.erase(iter); + } + lock.unlock(); + load_state->cv.notify_all(); + loader_slot_cv_.notify_one(); + } + + void clear_noexcept(Payload &payload) noexcept { + try { + loader_.clear(payload); + } catch (...) { + // Cache cleanup must always publish completion to single-flight waiters. + } + } + bool ensure_capacity() { bool found = !MemoryLimitPool::get_instance().is_full(); if (found) { @@ -200,13 +487,14 @@ class ExternalCache : public EvictableBlockOwner { private: static constexpr int kRecycleAttempts = 5; - using Table = std::unordered_map; - Loader loader_{}; Table table_; std::unordered_map owner_keys_; eviction_key_t next_owner_key_{1}; + size_t active_loads_{0}; + const size_t max_concurrent_loads_; std::shared_mutex mutex_; + std::condition_variable_any loader_slot_cv_; }; } // namespace ailego diff --git a/src/include/zvec/ailego/buffer/vector_page_table.h b/src/include/zvec/ailego/buffer/vector_page_table.h index 3e1372fcb..e123377f5 100644 --- a/src/include/zvec/ailego/buffer/vector_page_table.h +++ b/src/include/zvec/ailego/buffer/vector_page_table.h @@ -15,63 +15,103 @@ #pragma once -#include -#include +#include #include #include -#include -#include -#include +#include +#include +#include +#include #include -#include #include -#include #include #include -#include +#include #include #include -#include -#include +#include +#include +#include +#include #include #include "block_eviction_queue.h" -#include "concurrentqueue.h" - -#if defined(_MSC_VER) -#include -#endif namespace zvec { namespace ailego { extern const size_t kVectorPageSize; +#if defined(__linux__) +class IoUringRing; +#endif + class ZVEC_AILEGO_API VectorPageTable : public EvictableBlockOwner { - struct Entry { + // Keep every field used by the resident hit/release path in one compact + // object. Four entries fit in a 64-byte cache line and no entry straddles a + // line, so pinning a page also brings its resident pointer into cache. + struct alignas(16) Entry { + std::atomic buffer{nullptr}; std::atomic ref_count; std::atomic in_evict_queue; - std::atomic is_dirty; - char *buffer; + std::atomic referenced; + std::atomic evict_priority{0}; + std::atomic ghost_state{0}; + }; + static_assert(sizeof(Entry) == 16, + "VectorPageTable::Entry must stay hot and compact"); + + // Metadata that is not needed by a resident cache hit. Field order keeps + // the combined hot+cold cost at the previous 40 bytes per page. + struct MetadataEntry { + size_t next_loaded; size_t file_offset; + // High 24 bits: aging epoch; low 8 bits: recent rejected-miss count. + std::atomic admission_state{0}; + std::atomic is_dirty; + std::atomic writeback_pending{false}; + std::atomic ever_loaded{ + false}; // true once the page has been loaded at least once }; + static_assert(sizeof(MetadataEntry) == 24, + "VectorPageTable::MetadataEntry must stay compact"); + static_assert(sizeof(Entry) + sizeof(MetadataEntry) == 40, + "VectorPageTable metadata must remain 40 bytes per page"); public: + static constexpr uint8_t kLowPriority = + BlockEvictionQueue::kProbationPriority; + static constexpr uint8_t kNormalPriority = + BlockEvictionQueue::kProtectedPriority; + static constexpr uint8_t kHighPriority = + BlockEvictionQueue::kExplicitHotPriority; + static constexpr size_t kPriorityCount = BlockEvictionQueue::kQueueCount; + + enum class LoadClaimResult : uint8_t { + kClaimed, + kResident, + kLoading, + kEvicting, + }; + // Callback invoked by evict_block() to persist a dirty block before its // memory is released. Signature: (block_id, buffer, size, file_offset). using FlushCallback = std::function; + // Writable pools enqueue dirty candidates here instead of performing disk + // I/O on the global eviction thread. + using WritebackCallback = std::function; - VectorPageTable() { + VectorPageTable() : owner_version_(next_owner_version()) { BlockEvictionQueue::get_instance().set_valid(this); } ~VectorPageTable() { BlockEvictionQueue::get_instance().set_invalid(this); - // Destructor runs without concurrent readers/writers (callers guarantee - // no live handles by the time the page table is destroyed), so a relaxed - // load is sufficient here. + // No readers remain during destruction. size_t cnt = segment_count_.load(std::memory_order_relaxed); for (size_t i = 0; i < cnt; ++i) { delete[] segments_[i]; + delete[] metadata_segments_[i]; } + MemoryLimitPool::get_instance().release_metadata(metadata_bytes()); } VectorPageTable(const VectorPageTable &) = delete; @@ -79,40 +119,151 @@ class ZVEC_AILEGO_API VectorPageTable : public EvictableBlockOwner { VectorPageTable(VectorPageTable &&) = delete; VectorPageTable &operator=(VectorPageTable &&) = delete; - //! Initialize the page table to cover `entry_num` entries. - //! Returns false (without modifying state) if `entry_num` exceeds the - //! statically allocated segment table capacity (kMaxEntries). + //! Initialize up to kMaxEntries entries without partial publication. bool init(size_t entry_num); - //! Extend the page table to cover at least `new_entry_num` entries. - //! Existing entries stay at their original addresses (no invalidation). - //! Safe to call while readers operate on existing pages. - //! Returns false (without modifying state) if `new_entry_num` exceeds - //! the statically allocated segment table capacity (kMaxEntries). + //! Extend without moving existing entries or partial publication. bool extend(size_t new_entry_num); - char *acquire_block(block_id_t block_id); + //! Roll back an extension not yet exposed to page users. + bool rollback_extend(size_t old_entry_num); + + char *acquire_block(block_id_t block_id, bool record_reuse = true); + + //! Claim an unloaded page before issuing I/O. Exactly one concurrent loader + //! may receive kClaimed for a residency cycle. + LoadClaimResult try_claim_block_load(block_id_t block_id); + + //! Wait until an in-flight load or eviction publishes a stable state. + bool wait_for_block_transition(block_id_t block_id) const; + + //! Publish a page whose loading state is owned by the caller. + [[nodiscard]] char *publish_claimed_block(block_id_t block_id, char *buffer, + size_t file_offset); + + //! Roll an owned loading claim back to the unloaded state after I/O failure. + bool cancel_block_load(block_id_t block_id); void release_block(block_id_t block_id); - void evict_block(block_id_t block_id) override; + bool evict_block(block_id_t block_id) override; + + void eviction_requeue_failed(eviction_key_t owner_key, + version_t version) override { + if (version == owner_version_ && + owner_key < entry_num_.load(std::memory_order_acquire)) { + eviction_recovery_needed_.store(true, std::memory_order_release); + Entry &e = entry_at(owner_key); + e.in_evict_queue.store(false, std::memory_order_relaxed); + // Reclaim released pages whose queue slot was consumed. + if (e.ref_count.load(std::memory_order_acquire) == 0) { + (void)do_evict_block(owner_key, /*force=*/false); + } + } + } + + size_t recover_eviction_queue() override; + + uint8_t eviction_priority(eviction_key_t owner_key) const override { + if (owner_key >= entry_num_.load(std::memory_order_acquire)) { + return 0; + } + return entry_at(owner_key).evict_priority.load(std::memory_order_relaxed); + } + + //! Reclaim a block without CLOCK second chance. + bool force_evict_block(block_id_t block_id); + + //! Reclaim only if the released block is already clean. Unlike teardown + //! eviction this never persists or discards a dirty page. + bool reclaim_clean_block(block_id_t block_id); - char *set_block_acquired(block_id_t block_id, char *buffer, - size_t file_offset); + //! Reclaim loaded entries without scanning untouched file pages. + void force_evict_all_loaded(); + + void set_evict_priority(block_id_t block_id, uint8_t priority) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &e = entry_at(block_id); + e.evict_priority.store(priority, std::memory_order_relaxed); + // Existing queue membership adopts the priority on its next requeue. + } + + //! Raise an admission hint without allowing a colder caller to demote an + //! already protected page. + bool promote_evict_priority(block_id_t block_id, uint8_t priority) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + Entry &e = entry_at(block_id); + uint8_t current = e.evict_priority.load(std::memory_order_relaxed); + while (current < priority && + !e.evict_priority.compare_exchange_weak(current, priority, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + } + if (current >= priority) { + return false; + } + e.referenced.store(true, std::memory_order_relaxed); + inc_priority_promotion(priority); + // Existing queue membership adopts the priority on its next requeue. + return true; + } + + void set_adaptive_priority(bool enabled) { + adaptive_priority_enabled_ = enabled; + } + + [[nodiscard]] char *set_block_acquired(block_id_t block_id, char *buffer, + size_t file_offset); void set_flush_callback(FlushCallback cb) { flush_callback_ = std::move(cb); } + void set_writeback_callback(WritebackCallback cb) { + writeback_callback_ = std::move(cb); + } + //! Mark a loaded block as dirty so that it is persisted on eviction. void mark_dirty(block_id_t block_id) { assert(block_id < entry_num_.load(std::memory_order_acquire)); - entry_at(block_id).is_dirty.store(true, std::memory_order_relaxed); + auto &dirty = metadata_entry_at(block_id).is_dirty; + // Page writers hold the page's exclusive latch, while flush keeps its + // shared latch through dirty clearing. Avoid repeatedly taking ownership + // of the cold metadata cache line once a build page is already dirty. + if (!dirty.load(std::memory_order_relaxed)) { + dirty.store(true, std::memory_order_relaxed); + } } bool is_block_dirty(block_id_t block_id) const { assert(block_id < entry_num_.load(std::memory_order_acquire)); - return entry_at(block_id).is_dirty.load(std::memory_order_relaxed); + return metadata_entry_at(block_id).is_dirty.load(std::memory_order_relaxed); + } + + bool try_mark_writeback_pending(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + bool expected = false; + return metadata_entry_at(block_id) + .writeback_pending.compare_exchange_strong(expected, true, + std::memory_order_acq_rel, + std::memory_order_relaxed); + } + + void clear_writeback_pending(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + metadata_entry_at(block_id).writeback_pending.store( + false, std::memory_order_release); + } + + void record_dirty_flush(size_t count) { + inc_dirty_flush(count); + } + + //! Clear the dirty flag after a successful batched flush. + void clear_dirty(block_id_t block_id) { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + metadata_entry_at(block_id).is_dirty.store(false, + std::memory_order_relaxed); } //! Flush a single dirty block without evicting it. Caller guarantees the @@ -120,69 +271,153 @@ class ZVEC_AILEGO_API VectorPageTable : public EvictableBlockOwner { int flush_block(block_id_t block_id) { assert(block_id < entry_num_.load(std::memory_order_acquire)); Entry &e = entry_at(block_id); - char *buffer = e.buffer; + MetadataEntry &metadata = metadata_entry_at(block_id); + char *buffer = e.buffer.load(std::memory_order_acquire); if (!buffer || !flush_callback_) { return 0; } - if (!e.is_dirty.load(std::memory_order_relaxed)) { + if (!metadata.is_dirty.load(std::memory_order_relaxed)) { return 0; } - int rc = flush_callback_(block_id, buffer, kVectorPageSize, e.file_offset); + int rc = flush_callback_(block_id, buffer, kVectorPageSize, + metadata.file_offset); if (rc == 0) { - e.is_dirty.store(false, std::memory_order_relaxed); + metadata.is_dirty.store(false, std::memory_order_relaxed); } return rc; } - //! Returns the current number of entries. Uses acquire ordering so that - //! callers iterating over [0, entry_num()) are guaranteed to see all - //! segments_[s] writes performed by a concurrent extend()/init(). + //! Return the published entry count with initialized segments visible. size_t entry_num() const { return entry_num_.load(std::memory_order_acquire); } + size_t metadata_bytes() const { + return metadata_bytes_for_entries(entry_num()); + } + + //! Cache observability counters (monotonic, relaxed atomics). + struct Stats { + uint64_t hit{0}; // estimated cache hits (1/64 sampling) + uint64_t evict{0}; // pages actually reclaimed + uint64_t second_chance{0}; // pages spared by the CLOCK bit + uint64_t dirty_flush{0}; // dirty pages written back on eviction + uint64_t ghost_hot_marks{0}; // protected residencies remembered on aging + uint64_t ghost_hot_hits{0}; // remembered pages recognized on reload + std::array priority_promotions{}; + std::array priority_demotions{}; + std::array evictions_by_priority{}; + }; + Stats stats() const { + Stats s; + for (size_t i = 0; i < kCounterShards; ++i) { + const CounterShard &c = counters_[i]; + s.hit += c.hit.load(std::memory_order_relaxed); + s.evict += c.evict.load(std::memory_order_relaxed); + s.second_chance += c.second_chance.load(std::memory_order_relaxed); + s.dirty_flush += c.dirty_flush.load(std::memory_order_relaxed); + } + s.ghost_hot_marks = ghost_hot_marks_.load(std::memory_order_relaxed); + s.ghost_hot_hits = ghost_hot_hits_.load(std::memory_order_relaxed); + for (size_t priority = 0; priority < kPriorityCount; ++priority) { + s.priority_promotions[priority] = + priority_promotions_[priority].load(std::memory_order_relaxed); + s.priority_demotions[priority] = + priority_demotions_[priority].load(std::memory_order_relaxed); + s.evictions_by_priority[priority] = + evictions_by_priority_[priority].load(std::memory_order_relaxed); + } + return s; + } + + //! Logical resident pages in each priority tier. Intended for infrequent + //! diagnostics; this scans pages that have been loaded at least once. + std::array resident_pages_by_priority() const; + bool is_released(block_id_t block_id) const { assert(block_id < entry_num_.load(std::memory_order_acquire)); return entry_at(block_id).ref_count.load(std::memory_order_relaxed) <= 0; } - inline bool is_dead_block(block_id_t block_id, - version_t /*version*/) override { + inline bool is_dead_block(block_id_t block_id, version_t version) override { + // Reject stale entries after owner-address reuse. + if (version != owner_version_ || + block_id >= entry_num_.load(std::memory_order_acquire)) { + return true; + } const Entry &e = entry_at(block_id); return !e.in_evict_queue.load(std::memory_order_relaxed); } + //! Check if a page is loaded (has a non-null buffer). + bool is_loaded(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + return entry_at(block_id).buffer.load(std::memory_order_acquire) != nullptr; + } + + //! Check whether reload requires I/O rather than initial zero-fill. + bool is_ever_loaded(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + return metadata_entry_at(block_id).ever_loaded.load( + std::memory_order_relaxed); + } + + bool has_evicted() const { + return has_evicted_.load(std::memory_order_relaxed); + } + + //! Whether a page is resident or currently changing residency. Used to + //! rejoin cache single-flight without recording another admission miss. + bool has_residency_activity(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + return entry_at(block_id).ref_count.load(std::memory_order_acquire) != + kUnloadedRefCount; + } + + //! Return true when an unloaded demand page is worth admitting under + //! pressure. Resident/in-flight, protected, and ghost-hot pages always use + //! the cache path; cold pages require a second recent miss. + bool should_admit_miss(block_id_t block_id, uint32_t epoch); + private: // Segmented page table: entries are split across fixed-size segments so // that extend() can grow the table without moving existing entries. - static constexpr size_t kSegmentShift = 16; // 65536 entries per segment + static constexpr size_t kSegmentShift = 14; // 16384 entries per segment static constexpr size_t kSegmentSize = size_t{1} << kSegmentShift; static constexpr size_t kSegmentMask = kSegmentSize - 1; public: static constexpr size_t kMaxSegments = - 2048; // up to 128M entries (512GB @ 4K) - // Maximum number of entries the segment table can ever hold. Callers - // (e.g. VecBufferPool::extend_file) can use this to pre-validate a target - // file size before mutating any on-disk state. + 8192; // up to 128M entries (512GB @ 4K) + // Capacity used to validate growth before changing the file. static constexpr size_t kMaxEntries = kMaxSegments * kSegmentSize; + //! Heap bytes allocated by the segmented page table for `entry_num`. + static size_t metadata_bytes_for_entries(size_t entry_num); + private: - // entry_num_ and segment_count_ are mutated by writers in init()/extend() - // and observed by readers in entry_num() and the hot-path methods. They - // are atomic to establish a release/acquire synchronization edge with the - // (non-atomic) writes to segments_[s] performed prior to the store: any - // reader that observes the new entry_num_ is guaranteed to see the - // fully-initialized Entry slots in the corresponding segment. + // Release/acquire publication makes initialized segment slots visible. std::atomic entry_num_{0}; std::atomic segment_count_{0}; - Entry *segments_[kMaxSegments]{}; - - // Pair with the release-store on segment_count_ in init()/extend() so - // that any reader observing the published segment table also sees the - // fully-initialized segments_[s] pointer and Entry slots. Without this - // acquire load, segments_[s] can be re-read as nullptr or a torn - // pointer on weak memory models (and even reordered on x86 under -O2). + static constexpr size_t kSegmentMetadataBytes = + kSegmentSize * (sizeof(Entry) + sizeof(MetadataEntry)); + static constexpr size_t kSegmentDirectoryBytes = + kMaxSegments * (sizeof(Entry *) + sizeof(MetadataEntry *)); + std::unique_ptr segments_{}; + std::unique_ptr metadata_segments_{}; + static constexpr size_t kInvalidLoadedBlock = + std::numeric_limits::max(); + static constexpr int kUnloadedRefCount = std::numeric_limits::min(); + static constexpr int kLoadingRefCount = kUnloadedRefCount + 1; + static constexpr int kEvictingRefCount = std::numeric_limits::min() / 2; + static constexpr uint8_t kNoGhostHistory = 0; + static constexpr uint8_t kEvictedHot = 1; + static constexpr uint8_t kGhostAdmitted = 2; + std::atomic loaded_head_{kInvalidLoadedBlock}; + std::atomic has_evicted_{false}; + bool adaptive_priority_enabled_{true}; + + // Pair with segment_count_ publication before dereferencing a segment. Entry &entry_at(size_t idx) { (void)segment_count_.load(std::memory_order_acquire); return segments_[idx >> kSegmentShift][idx & kSegmentMask]; @@ -191,86 +426,425 @@ class ZVEC_AILEGO_API VectorPageTable : public EvictableBlockOwner { (void)segment_count_.load(std::memory_order_acquire); return segments_[idx >> kSegmentShift][idx & kSegmentMask]; } + MetadataEntry &metadata_entry_at(size_t idx) { + (void)segment_count_.load(std::memory_order_acquire); + return metadata_segments_[idx >> kSegmentShift][idx & kSegmentMask]; + } + const MetadataEntry &metadata_entry_at(size_t idx) const { + (void)segment_count_.load(std::memory_order_acquire); + return metadata_segments_[idx >> kSegmentShift][idx & kSegmentMask]; + } + + // `force` bypasses CLOCK second chance. + static void initialize_segment(Entry *entries, + MetadataEntry *metadata_entries); + bool do_evict_block(block_id_t block_id, bool force); + static version_t next_owner_version(); + + // Prevent stale queue entries from targeting a reused owner address. + const version_t owner_version_; FlushCallback flush_callback_{}; + WritebackCallback writeback_callback_{}; + // Scan loaded entries only after a queue insertion failure. + std::atomic eviction_recovery_needed_{false}; + + // Shard relaxed statistics to avoid hot-path cache-line contention. + static constexpr size_t kCounterShards = 64; // power of two for masking + struct alignas(64) CounterShard { + std::atomic hit{0}; + std::atomic evict{0}; + std::atomic second_chance{0}; + std::atomic dirty_flush{0}; + }; + CounterShard counters_[kCounterShards]; + // Priority transitions happen at most once per residency phase, so they do + // not need the per-hit counter sharding above. + std::array, kPriorityCount> priority_promotions_{}; + std::array, kPriorityCount> priority_demotions_{}; + std::array, kPriorityCount> evictions_by_priority_{}; + std::atomic ghost_hot_marks_{0}; + std::atomic ghost_hot_hits_{0}; + + // Keep each thread on one counter shard. + static size_t counter_shard() { + static std::atomic seq{0}; + thread_local size_t idx = seq.fetch_add(1, std::memory_order_relaxed); + return idx & (kCounterShards - 1); + } + // Sample and scale hits to avoid an atomic RMW on every acquisition. + static constexpr uint32_t kHitSampleRate = 64; + static constexpr uint32_t kReusePolicySampleRate = 8; + static uint32_t next_hit_sample() { + thread_local uint32_t sample_cursor = 0; + return sample_cursor++; + } + void inc_sampled_hit() { + counters_[counter_shard()].hit.fetch_add(kHitSampleRate, + std::memory_order_relaxed); + } + void inc_evict(uint8_t priority) { + has_evicted_.store(true, std::memory_order_relaxed); + CounterShard &counter = counters_[counter_shard()]; + counter.evict.fetch_add(1, std::memory_order_relaxed); + if (priority < kPriorityCount) { + evictions_by_priority_[priority].fetch_add(1, std::memory_order_relaxed); + } + } + void inc_second_chance() { + counters_[counter_shard()].second_chance.fetch_add( + 1, std::memory_order_relaxed); + } + void inc_dirty_flush() { + counters_[counter_shard()].dirty_flush.fetch_add(1, + std::memory_order_relaxed); + } + void inc_dirty_flush(size_t count) { + counters_[counter_shard()].dirty_flush.fetch_add(count, + std::memory_order_relaxed); + } + void inc_priority_promotion(uint8_t priority) { + if (priority < kPriorityCount) { + priority_promotions_[priority].fetch_add(1, std::memory_order_relaxed); + } + } + void inc_priority_demotion(uint8_t priority) { + if (priority < kPriorityCount) { + priority_demotions_[priority].fetch_add(1, std::memory_order_relaxed); + } + } }; +class VecBufferPool; class VecBufferPoolHandle; +struct VecBufferWriteFragment { + size_t file_offset; + size_t length; + const char *src; +}; + class ZVEC_AILEGO_API VecBufferPool { public: typedef std::shared_ptr Pointer; - static constexpr size_t kMutexBucketCount = 64UL * 1024UL; + static constexpr size_t kMutexBucketCount = 4UL * 1024UL; + static constexpr size_t kWritebackBatchPages = 128; + static constexpr uint8_t kLowPriority = VectorPageTable::kLowPriority; + static constexpr uint8_t kNormalPriority = VectorPageTable::kNormalPriority; + static constexpr uint8_t kHighPriority = VectorPageTable::kHighPriority; + + //! Non-evictable page-table and striped-lock memory required for a pool + //! covering `page_count` pages. + static size_t metadata_bytes_for_page_count(size_t page_count, + bool writable = false); VecBufferPool(const std::string &filename, bool writable = false); - ~VecBufferPool() { - // Flush any remaining dirty blocks before tearing down memory/fd so that - // writes are not silently lost. Safe to call even in read-only mode. - (void)this->flush_all(); - for (size_t i = 0; i < page_table_.entry_num(); ++i) { - assert(page_table_.is_released(i)); - page_table_.evict_block(i); + ~VecBufferPool(); + + int init(); + + //! Aggregated cache statistics for this pool. + struct Stats { + uint64_t hit{0}; + uint64_t miss{0}; + uint64_t evict{0}; + uint64_t second_chance{0}; + uint64_t dirty_flush{0}; + uint64_t writeback_requests{0}; + uint64_t writeback_batches{0}; + uint64_t writeback_pages{0}; + uint64_t writeback_failures{0}; + uint64_t writeback_aio_batches{0}; + uint64_t writeback_aio_pages{0}; + uint64_t writeback_aio_fallbacks{0}; + uint64_t writeback_waits{0}; + uint64_t writeback_wait_us{0}; + uint64_t writeback_pending{0}; + uint64_t writeback_peak_pending{0}; + uint64_t bypass_reads{0}; + uint64_t bypass_bytes{0}; + uint64_t bypass_io_requests{0}; + uint64_t bypass_rechecks{0}; + uint64_t bypass_cache_joins{0}; + uint64_t singleflight_waits{0}; + uint64_t aio_pages_submitted{0}; + uint64_t admission_admitted{0}; + uint64_t admission_rejected{0}; + uint64_t ghost_hot_marks{0}; + uint64_t ghost_hot_hits{0}; + size_t page_table_metadata_bytes{0}; + size_t page_lock_metadata_bytes{0}; + size_t writeback_staging_bytes{0}; + size_t writeback_io_staging_bytes{0}; + std::array priority_promotions{}; + std::array priority_demotions{}; + std::array + evictions_by_priority{}; + double hit_rate() const { + uint64_t total = hit + miss; + return total ? static_cast(hit) / static_cast(total) + : 0.0; } -#if defined(_MSC_VER) - _close(fd_); -#else - close(fd_); -#endif + }; + Stats stats() const { + VectorPageTable::Stats p = page_table_.stats(); + Stats s; + s.hit = p.hit; + s.evict = p.evict; + s.second_chance = p.second_chance; + s.dirty_flush = p.dirty_flush; + s.writeback_requests = writeback_requests_.load(std::memory_order_relaxed); + s.writeback_batches = writeback_batches_.load(std::memory_order_relaxed); + s.writeback_pages = writeback_pages_.load(std::memory_order_relaxed); + s.writeback_failures = writeback_failures_.load(std::memory_order_relaxed); + s.writeback_aio_batches = + writeback_aio_batches_.load(std::memory_order_relaxed); + s.writeback_aio_pages = + writeback_aio_pages_.load(std::memory_order_relaxed); + s.writeback_aio_fallbacks = + writeback_aio_fallbacks_.load(std::memory_order_relaxed); + s.writeback_waits = writeback_waits_.load(std::memory_order_relaxed); + s.writeback_wait_us = writeback_wait_us_.load(std::memory_order_relaxed); + s.writeback_pending = writeback_pending_.load(std::memory_order_relaxed); + s.writeback_peak_pending = + writeback_peak_pending_.load(std::memory_order_relaxed); + s.priority_promotions = p.priority_promotions; + s.priority_demotions = p.priority_demotions; + s.evictions_by_priority = p.evictions_by_priority; + s.miss = miss_count_.load(std::memory_order_relaxed); + s.bypass_reads = bypass_reads_.load(std::memory_order_relaxed); + s.bypass_bytes = bypass_bytes_.load(std::memory_order_relaxed); + s.bypass_io_requests = bypass_io_requests_.load(std::memory_order_relaxed); + s.bypass_rechecks = bypass_rechecks_.load(std::memory_order_relaxed); + s.bypass_cache_joins = bypass_cache_joins_.load(std::memory_order_relaxed); + s.singleflight_waits = singleflight_waits_.load(std::memory_order_relaxed); + s.aio_pages_submitted = + aio_pages_submitted_.load(std::memory_order_relaxed); + s.admission_admitted = admission_admitted_.load(std::memory_order_relaxed); + s.admission_rejected = admission_rejected_.load(std::memory_order_relaxed); + s.ghost_hot_marks = p.ghost_hot_marks; + s.ghost_hot_hits = p.ghost_hot_hits; + s.page_table_metadata_bytes = page_table_.metadata_bytes(); + s.page_lock_metadata_bytes = block_mutex_metadata_bytes(); + s.writeback_staging_bytes = writeback_staging_size_; + s.writeback_io_staging_bytes = writeback_io_staging_charge_; + return s; } - int init(); + //! Log the current cache statistics at INFO level. + void log_stats() const; VecBufferPoolHandle get_handle(); - char *acquire_buffer(block_id_t page_id, int retry = 0); + char *acquire_buffer(block_id_t page_id, int retry = 0, + bool record_reuse = true); + + //! Pin scattered pages only when every page is already resident. Rolls back + //! all pins on the first miss and never triggers I/O or cache admission. + bool try_acquire_resident_pages(const block_id_t *page_ids, size_t count, + char **pages); + + //! Pin scattered pages; roll back all pins on failure. + bool acquire_pages(const block_id_t *page_ids, size_t count, char **pages); + + //! Release one pin per page id acquired by acquire_pages(). + void release_pages(const block_id_t *page_ids, size_t count); + + //! Observe residency without changing hit or CLOCK state. + bool is_page_resident(block_id_t page_id) const { + return page_id < page_table_.entry_num() && page_table_.is_loaded(page_id); + } + + //! Decide whether a demand miss should enter the cache. Admission control + //! activates when the process-wide shared pool is under pressure. + bool should_admit_page(block_id_t page_id); + + //! Recheck a rejected page without adding another frequency observation. + bool should_join_cache_path(block_id_t page_id) const { + return page_id < page_table_.entry_num() && + page_table_.has_residency_activity(page_id); + } + + void record_bypass_recheck(size_t checked, size_t joined) { + if (checked != 0) { + bypass_rechecks_.fetch_add(checked, std::memory_order_relaxed); + } + if (joined != 0) { + bypass_cache_joins_.fetch_add(joined, std::memory_order_relaxed); + } + } + + //! Account for one successful direct read that bypassed cache admission. + void record_bypass_read(size_t length, size_t io_requests = 1) { + bypass_reads_.fetch_add(1, std::memory_order_relaxed); + bypass_bytes_.fetch_add(length, std::memory_order_relaxed); + bypass_io_requests_.fetch_add(io_requests, std::memory_order_relaxed); + } int get_meta(size_t offset, size_t length, char *buffer); + //! Read without cache admission. + bool read_range_bypass(size_t file_offset, size_t length, char *buffer); + //! Write a contiguous range via the page cache; marks touched pages dirty. //! Returns 0 on success, -1 on failure (e.g. read-only pool or I/O error). int write_range(size_t file_offset, size_t length, const char *src); - //! Write raw bytes directly via pwrite, bypassing the page cache. Used for - //! metadata regions (header/footer/segments_meta) which are only read via - //! get_meta() and never cached. + //! Apply ordered fragments that all lie in one page under one pin/latch. + int write_fragments(const VecBufferWriteFragment *fragments, size_t count); + + //! Write metadata without cache admission. int write_meta(size_t offset, size_t length, const char *buffer); //! Iterate all entries and persist any dirty blocks to disk. Safe to call //! repeatedly; no-op in read-only mode. int flush_all(); - //! Extend the backing file to `new_size` bytes via ftruncate (no-op if - //! already >= new_size), refresh the cached file_size_, and extend the - //! page_table to cover the new range. Returns true on success, false on - //! a read-only pool or I/O failure. + //! Extend the backing file and page table to `new_size`. bool extend_file(size_t new_size); bool writable() const { return writable_; } + bool has_evicted() const { + return page_table_.has_evicted(); + } + size_t file_size() const { return file_size_; } + //! Sequentially preload pages into the pool until pool is full. + void warmup(); + + void prefetch_pages(block_id_t first_page, size_t page_count, + uint8_t priority = kLowPriority); + + void prefetch_pages_aio(block_id_t first_page, size_t page_count, + uint8_t priority = kLowPriority); + + bool set_page_priority(block_id_t page_id, uint8_t priority) { + if (page_id >= page_table_.entry_num() || priority > kHighPriority) { + return false; + } + page_table_.set_evict_priority(page_id, priority); + return true; + } + + bool aio_enabled() const { +#if defined(__linux__) + // Backend contexts are created lazily per calling thread. + return aio_enabled_; +#else + return false; +#endif + } + + IOBackendType io_backend_type() const { +#if defined(__linux__) + return io_backend_type_; +#else + return IOBackendType::kPread; +#endif + } + + //! Acquire a resident page without triggering I/O. + char *try_acquire_buffer(block_id_t page_id) { + assert(page_id < page_table_.entry_num()); + return page_table_.acquire_block(page_id); + } + private: - int fd_; + friend class VecBufferPoolHandle; + void prefetch_pages_sync(block_id_t first_page, size_t page_count, + uint8_t priority); + bool load_pages_aio(const block_id_t *page_ids, size_t count, + uint8_t priority); + bool enqueue_writeback(block_id_t page_id); + void start_writeback(); + void stop_writeback(); + void drain_writeback(); + void writeback_loop(); + bool flush_writeback_batch(std::vector &page_ids, char *staging); + int writeback_error() const { + return writeback_error_.load(std::memory_order_acquire); + } + + int fd_; // page-data channel: O_DIRECT or F_NOCACHE when supported + int meta_fd_; // metadata channel: always buffered IO size_t file_size_; + size_t initial_file_size_; // file size at open time; pages beyond this + // are created by extend_file and can skip + // pread on first load (content is zeros). std::string file_name_; bool writable_{false}; + bool direct_io_enabled_{false}; + bool initialized_{false}; + // One miss per page populated on the cold path. + std::atomic miss_count_{0}; + std::atomic bypass_reads_{0}; + std::atomic bypass_bytes_{0}; + std::atomic bypass_io_requests_{0}; + std::atomic bypass_rechecks_{0}; + std::atomic bypass_cache_joins_{0}; + std::atomic singleflight_waits_{0}; + std::atomic aio_pages_submitted_{0}; + std::atomic admission_observations_{0}; + std::atomic admission_admitted_{0}; + std::atomic admission_rejected_{0}; + std::atomic writeback_requests_{0}; + std::atomic writeback_batches_{0}; + std::atomic writeback_pages_{0}; + std::atomic writeback_failures_{0}; + std::atomic writeback_aio_batches_{0}; + std::atomic writeback_aio_pages_{0}; + std::atomic writeback_aio_fallbacks_{0}; + std::atomic writeback_waits_{0}; + std::atomic writeback_wait_us_{0}; + std::atomic writeback_pending_{0}; + std::atomic writeback_peak_pending_{0}; + std::atomic writeback_error_{0}; +#if defined(__linux__) + IOBackendType io_backend_type_{IOBackendType::kPread}; + bool aio_enabled_{false}; +#endif public: VectorPageTable page_table_; private: - std::unique_ptr block_mutexes_{}; + // Serialize writable in-place payload access. Single-flight owns loading. + size_t block_mutex_metadata_bytes() const { + return block_mutex_count_ * sizeof(std::shared_mutex); + } + std::unique_ptr block_mutexes_{}; + size_t block_mutex_count_{0}; + + std::thread writeback_thread_{}; + char *writeback_staging_{nullptr}; + size_t writeback_staging_size_{0}; + size_t writeback_io_staging_charge_{0}; +#if defined(__linux__) + std::unique_ptr writeback_io_uring_{}; +#endif + std::mutex writeback_mutex_{}; + std::condition_variable writeback_cv_{}; + std::condition_variable writeback_drained_cv_{}; + std::deque writeback_queue_{}; + std::mutex writeback_flush_mutex_{}; + size_t writeback_inflight_{0}; + bool writeback_stopping_{false}; }; class ZVEC_AILEGO_API VecBufferPoolHandle { public: VecBufferPoolHandle(VecBufferPool &pool) : pool_(pool) {} - VecBufferPoolHandle(VecBufferPoolHandle &&other) : pool_(other.pool_) {} + explicit VecBufferPoolHandle(std::shared_ptr pool) + : pool_owner_(std::move(pool)), pool_(checked_pool(pool_owner_)) {} + VecBufferPoolHandle(VecBufferPoolHandle &&other) + : pool_owner_(std::move(other.pool_owner_)), pool_(other.pool_) {} ~VecBufferPoolHandle() = default; @@ -278,12 +852,31 @@ class ZVEC_AILEGO_API VecBufferPoolHandle { char *get_single_page(size_t file_offset, size_t len, size_t &out_page_id); + bool try_acquire_resident_pages(const block_id_t *page_ids, size_t count, + char **pages); + + bool acquire_pages(const block_id_t *page_ids, size_t count, char **pages); + + void release_pages(const block_id_t *page_ids, size_t count); + bool read_range(size_t file_offset, size_t len, char *out); + // Copy a range whose bytes are immutable after publication. Writers may + // still update disjoint records in the same page, so pages stay pinned for + // the copy but do not need the writable payload latch. + bool read_range_immutable(size_t file_offset, size_t len, char *out); + + bool read_range_bypass(size_t file_offset, size_t len, char *out); + + void prefetch_range(size_t file_offset, size_t len, + uint8_t priority = VecBufferPool::kLowPriority); + int get_meta(size_t offset, size_t length, char *buffer); int write_range(size_t file_offset, size_t len, const char *src); + int write_fragments(const VecBufferWriteFragment *fragments, size_t count); + int write_meta(size_t offset, size_t length, const char *buffer); int flush_all(); @@ -295,6 +888,17 @@ class ZVEC_AILEGO_API VecBufferPoolHandle { void acquire_one(block_id_t block_id); private: + static VecBufferPool &checked_pool( + const std::shared_ptr &pool) { + if (!pool) { + throw std::invalid_argument( + "VecBufferPoolHandle requires a non-null owning pool"); + } + return *pool; + } + + // Storage-backed handles own the pool; stack handles remain non-owning. + std::shared_ptr pool_owner_{}; VecBufferPool &pool_; }; diff --git a/tests/ailego/CMakeLists.txt b/tests/ailego/CMakeLists.txt index 9a52af94f..15ded49d6 100644 --- a/tests/ailego/CMakeLists.txt +++ b/tests/ailego/CMakeLists.txt @@ -4,11 +4,29 @@ include(${PROJECT_ROOT_DIR}/cmake/option.cmake) file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) +# Exercise the process-wide cache singleton contract when both a host and a +# component DSO embed zvec_ailego's static archive. The storage layer relies +# on shared capacity/accounting state across those images. +if(NOT WIN32 AND NOT IOS AND NOT ANDROID) + cc_library( + NAME cross_dso_buffer_pool_probe SHARED EXCLUDE STRICT + SRCS buffer/cross_dso_buffer_pool_probe.cc + LIBS zvec_ailego + ) +endif() + foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + set(CC_TEST_LIBS zvec_ailego) + if(CC_TARGET STREQUAL "cross_dso_buffer_pool_test") + if(WIN32 OR IOS OR ANDROID) + continue() + endif() + list(APPEND CC_TEST_LIBS cross_dso_buffer_pool_probe) + endif() cc_gtest( NAME ${CC_TARGET} STRICT - LIBS zvec_ailego + LIBS ${CC_TEST_LIBS} SRCS ${CC_SRCS} ) cc_test_suite(zvec_ailego ${CC_TARGET}) diff --git a/tests/ailego/buffer/cross_dso_buffer_pool_probe.cc b/tests/ailego/buffer/cross_dso_buffer_pool_probe.cc new file mode 100644 index 000000000..1a37a4ca1 --- /dev/null +++ b/tests/ailego/buffer/cross_dso_buffer_pool_probe.cc @@ -0,0 +1,38 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +extern "C" { + +ZVEC_HELPER_DLL_EXPORT void *zvec_test_dso_memory_limit_pool() { + return &zvec::ailego::MemoryLimitPool::get_instance(); +} + +ZVEC_HELPER_DLL_EXPORT void *zvec_test_dso_block_eviction_queue() { + return &zvec::ailego::BlockEvictionQueue::get_instance(); +} + +ZVEC_HELPER_DLL_EXPORT bool zvec_test_dso_charge_external(size_t bytes) { + return zvec::ailego::MemoryLimitPool::get_instance().try_charge_external( + bytes); +} + +ZVEC_HELPER_DLL_EXPORT void zvec_test_dso_release_external(size_t bytes) { + zvec::ailego::MemoryLimitPool::get_instance().release_external(bytes); +} + +} // extern "C" diff --git a/tests/ailego/buffer/cross_dso_buffer_pool_test.cc b/tests/ailego/buffer/cross_dso_buffer_pool_test.cc new file mode 100644 index 000000000..4047302f1 --- /dev/null +++ b/tests/ailego/buffer/cross_dso_buffer_pool_test.cc @@ -0,0 +1,46 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +extern "C" { +void *zvec_test_dso_memory_limit_pool(); +void *zvec_test_dso_block_eviction_queue(); +bool zvec_test_dso_charge_external(size_t bytes); +void zvec_test_dso_release_external(size_t bytes); +} + +namespace zvec { +namespace ailego { + +TEST(CrossDsoBufferPoolTest, SharesAddressesAndAccountingState) { + auto &pool = MemoryLimitPool::get_instance(); + auto &eviction_queue = BlockEvictionQueue::get_instance(); + + EXPECT_EQ(&pool, zvec_test_dso_memory_limit_pool()); + EXPECT_EQ(&eviction_queue, zvec_test_dso_block_eviction_queue()); + + const size_t bytes = 4 * kVectorPageSize; + ASSERT_EQ(0, pool.init(bytes)); + ASSERT_TRUE(zvec_test_dso_charge_external(kVectorPageSize)); + EXPECT_EQ(kVectorPageSize, pool.external_used()); + zvec_test_dso_release_external(kVectorPageSize); + EXPECT_EQ(0u, pool.external_used()); +} + +} // namespace ailego +} // namespace zvec diff --git a/tests/ailego/buffer/vector_page_table_test.cc b/tests/ailego/buffer/vector_page_table_test.cc new file mode 100644 index 000000000..9b97a89b4 --- /dev/null +++ b/tests/ailego/buffer/vector_page_table_test.cc @@ -0,0 +1,2153 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Tests for the buffer-pool optimizations: +// 1. CLOCK second-chance eviction (access-aware, data-correct under pressure) +// 2. Background evictor (proactive reclaim down to the low watermark) +// 3. Sharded free-list correctness under concurrent access +// 4. Reclaimable 4 MiB-aligned slab allocation +// 5. Observability counters (hit / miss / evict / second_chance / stats) + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::ailego; + +namespace { + +// Create a backing file of `num_pages` pages, page p filled with byte (p & +// 0xff) so page content can be verified after arbitrary eviction/reload. +std::string MakeBackingFile(size_t num_pages) { + static std::atomic seq{0}; + const size_t ps = kVectorPageSize; + std::string path = "vpt_test_" + std::to_string(seq.fetch_add(1)) + ".bin"; + std::remove(path.c_str()); + FILE *f = std::fopen(path.c_str(), "wb"); + EXPECT_NE(f, nullptr); + std::vector page(ps); + for (size_t p = 0; p < num_pages; ++p) { + std::memset(page.data(), static_cast(p & 0xff), ps); + EXPECT_EQ(std::fwrite(page.data(), 1, ps, f), ps); + } + std::fclose(f); + return path; +} + +// Verify that a page-sized buffer holds the expected fill byte. +void ExpectPageContent(const char *buf, size_t page_id) { + const size_t ps = kVectorPageSize; + char expected = static_cast(page_id & 0xff); + ASSERT_EQ(buf[0], expected) << "page " << page_id << " head mismatch"; + ASSERT_EQ(buf[ps - 1], expected) << "page " << page_id << " tail mismatch"; +} + +class BufferPoolTest : public ::testing::Test { + protected: + void InitPool(size_t capacity_pages) { + ASSERT_EQ(0, MemoryLimitPool::get_instance().init(capacity_pages * + kVectorPageSize)); + } + void InitVecPool(size_t capacity_pages, size_t file_pages, + bool writable = false) { + ASSERT_EQ(0, MemoryLimitPool::get_instance().init( + capacity_pages * kVectorPageSize + + VecBufferPool::metadata_bytes_for_page_count(file_pages, + writable))); + } + void InitTablePool(size_t capacity_pages, size_t entry_num) { + ASSERT_EQ(0, MemoryLimitPool::get_instance().init( + capacity_pages * kVectorPageSize + + VectorPageTable::metadata_bytes_for_entries(entry_num))); + } + void TearDown() override { + for (const auto &p : files_) std::remove(p.c_str()); + files_.clear(); + } + std::string NewFile(size_t num_pages) { + files_.push_back(MakeBackingFile(num_pages)); + return files_.back(); + } + std::vector files_; +}; + +struct SizedCachePayload { + std::shared_ptr> data; +}; + +struct SizedCacheLoader { + using Value = std::shared_ptr>; + + bool load(size_t bytes, SizedCachePayload &payload, size_t &size) { + payload.data = std::make_shared>(bytes); + size = bytes; + return true; + } + + Value value(const SizedCachePayload &payload) const { + return payload.data; + } + + void clear(SizedCachePayload &payload) const { + payload.data.reset(); + } +}; + +using SizedExternalCache = + ExternalCache, std::equal_to>; + +struct EmptyValueLoader { + using Value = std::shared_ptr>; + + bool load(size_t bytes, SizedCachePayload &payload, size_t &size) { + payload.data = std::make_shared>(bytes); + size = bytes; + return true; + } + + Value value(const SizedCachePayload &) const { + return nullptr; + } + + void clear(SizedCachePayload &payload) const { + payload.data.reset(); + } +}; + +using EmptyValueExternalCache = + ExternalCache, std::equal_to>; + +struct BlockingLoadState { + std::atomic load_calls{0}; + std::atomic active_loads{0}; + std::atomic max_active_loads{0}; + std::atomic finish{false}; +}; + +struct BlockingLoader { + using Value = std::shared_ptr>; + + bool load(size_t bytes, SizedCachePayload &payload, size_t &size) { + state->load_calls.fetch_add(1, std::memory_order_release); + const size_t active = + state->active_loads.fetch_add(1, std::memory_order_acq_rel) + 1; + size_t observed = state->max_active_loads.load(std::memory_order_relaxed); + while (observed < active && + !state->max_active_loads.compare_exchange_weak( + observed, active, std::memory_order_relaxed)) { + } + while (!state->finish.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + payload.data = std::make_shared>(bytes); + size = bytes; + state->active_loads.fetch_sub(1, std::memory_order_release); + return true; + } + + Value value(const SizedCachePayload &payload) const { + return payload.data; + } + + void clear(SizedCachePayload &payload) const { + payload.data.reset(); + } + + std::shared_ptr state; +}; + +using BlockingExternalCache = + ExternalCache, + std::equal_to>; + +struct ThrowingCachePayload { + ThrowingCachePayload() { + const size_t current = ++construction_count; + if (throw_on_construction != 0 && current == throw_on_construction) { + throw std::runtime_error("injected payload construction failure"); + } + } + + std::shared_ptr> data; + static size_t construction_count; + static size_t throw_on_construction; +}; + +size_t ThrowingCachePayload::construction_count = 0; +size_t ThrowingCachePayload::throw_on_construction = 0; + +struct ThrowingCacheLoader { + using Value = std::shared_ptr>; + + bool load(size_t bytes, ThrowingCachePayload &payload, size_t &size) { + payload.data = std::make_shared>(bytes); + size = bytes; + return true; + } + + Value value(const ThrowingCachePayload &payload) const { + return payload.data; + } + + void clear(ThrowingCachePayload &payload) const { + payload.data.reset(); + } +}; + +using ThrowingExternalCache = + ExternalCache, std::equal_to>; + +class AlwaysDeadOwner : public EvictableBlockOwner { + public: + AlwaysDeadOwner() { + BlockEvictionQueue::get_instance().set_valid(this); + } + + ~AlwaysDeadOwner() override { + BlockEvictionQueue::get_instance().set_invalid(this); + } + + bool is_dead_block(eviction_key_t, version_t) override { + ++dead_checks; + return true; + } + + bool evict_block(eviction_key_t) override { + return false; + } + + size_t dead_checks{0}; +}; + +version_t FindLiveVersion(SizedExternalCache &cache, eviction_key_t owner_key) { + constexpr version_t kMaxProbe = 1UL << 20; + for (version_t version = 1; version < kMaxProbe; ++version) { + if (!cache.is_dead_block(owner_key, version)) { + return version; + } + } + return 0; +} + +} // namespace + +TEST_F(BufferPoolTest, AdmissionControlRejectsFirstColdMissAfterPressure) { + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + + // Fill-before-pressure remains unchanged. + EXPECT_TRUE(pool.should_admit_page(1)); + EXPECT_EQ(pool.stats().admission_rejected, 0u); + + char *page = pool.acquire_buffer(0, 10); + ASSERT_NE(page, nullptr); + + EXPECT_FALSE(pool.should_admit_page(1)); + EXPECT_TRUE(pool.should_admit_page(1)); + EXPECT_FALSE(pool.is_page_resident(1)); + pool.page_table_.release_block(0); + + const auto stats = pool.stats(); + EXPECT_EQ(stats.admission_rejected, 1u); + EXPECT_EQ(stats.admission_admitted, 1u); +} + +TEST_F(BufferPoolTest, BulkReadDoesNotAdmitFirstTouchScanUnderPressure) { + constexpr size_t kFilePages = 5; + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/kFilePages); + std::string file = NewFile(kFilePages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + // Hold the only cache page so the four-page range takes the bulk cold-read + // path while the pool is under pressure. + char *pinned = pool.acquire_buffer(/*block_id=*/0, 10); + ASSERT_NE(nullptr, pinned); + + std::vector data(4 * kVectorPageSize); + ASSERT_TRUE(handle.read_range(kVectorPageSize, data.size(), data.data())); + for (size_t page = 1; page < kFilePages; ++page) { + ExpectPageContent(data.data() + (page - 1) * kVectorPageSize, page); + EXPECT_FALSE(pool.is_page_resident(page)); + } + + const auto stats = pool.stats(); + EXPECT_EQ(4u, stats.admission_rejected); + pool.page_table_.release_block(/*block_id=*/0); +} + +TEST_F(BufferPoolTest, ShortReadDoesNotEvictHotPageOnFirstTouch) { + constexpr size_t kFilePages = 3; + constexpr size_t kCapacity = 256UL * 1024UL * 1024UL; + auto &memory_pool = MemoryLimitPool::get_instance(); + ASSERT_EQ(0, memory_pool.init(kCapacity)); + std::string file = NewFile(kFilePages); + + size_t charged_metadata = 0; + { + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + const size_t reserve = memory_pool.page_admission_reserve(); + ASSERT_EQ(16UL * 1024UL * 1024UL, reserve); + ASSERT_LT(memory_pool.used() + kVectorPageSize, kCapacity - reserve); + charged_metadata = + kCapacity - reserve - memory_pool.used() - kVectorPageSize; + ASSERT_TRUE(memory_pool.try_charge_metadata(charged_metadata)); + + char *hot = pool.acquire_buffer(/*block_id=*/0, 10); + ASSERT_NE(nullptr, hot); + pool.page_table_.release_block(/*block_id=*/0); + + std::vector data(2 * kVectorPageSize); + ASSERT_TRUE(handle.read_range(kVectorPageSize, data.size(), data.data())); + ExpectPageContent(data.data(), /*page=*/1); + ExpectPageContent(data.data() + kVectorPageSize, /*page=*/2); + + EXPECT_TRUE(pool.is_page_resident(0)); + EXPECT_FALSE(pool.is_page_resident(1)); + EXPECT_FALSE(pool.is_page_resident(2)); + const auto stats = pool.stats(); + EXPECT_EQ(2u, stats.admission_rejected); + EXPECT_EQ(2u, stats.bypass_reads); + } + memory_pool.release_metadata(charged_metadata); +} + +TEST_F(BufferPoolTest, BypassRecheckRecognizesResidencyActivity) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/2); + std::string file = NewFile(/*num_pages=*/2); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + EXPECT_FALSE(pool.should_join_cache_path(1)); + + ASSERT_EQ(VectorPageTable::LoadClaimResult::kClaimed, + pool.page_table_.try_claim_block_load(1)); + EXPECT_TRUE(pool.should_join_cache_path(1)); + ASSERT_TRUE(pool.page_table_.cancel_block_load(1)); + EXPECT_FALSE(pool.should_join_cache_path(1)); + + char *page = pool.acquire_buffer(1, 10); + ASSERT_NE(page, nullptr); + EXPECT_TRUE(pool.should_join_cache_path(1)); + pool.page_table_.release_block(1); +} + +TEST_F(BufferPoolTest, ResidentOnlyAcquirePreservesTransitionStates) { + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/1); + std::string file = NewFile(/*num_pages=*/1); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + + // The resident-only fast path must not claim or pin an unloaded page. + EXPECT_EQ(nullptr, pool.try_acquire_buffer(/*page_id=*/0)); + ASSERT_EQ(VectorPageTable::LoadClaimResult::kClaimed, + pool.page_table_.try_claim_block_load(/*block_id=*/0)); + + // Observing an in-flight page must leave ownership with the loader. + EXPECT_EQ(nullptr, pool.try_acquire_buffer(/*page_id=*/0)); + EXPECT_EQ(VectorPageTable::LoadClaimResult::kLoading, + pool.page_table_.try_claim_block_load(/*block_id=*/0)); + ASSERT_TRUE(pool.page_table_.cancel_block_load(/*block_id=*/0)); + + char *loaded = pool.acquire_buffer(/*page_id=*/0, /*retry=*/10); + ASSERT_NE(nullptr, loaded); + char *resident = pool.try_acquire_buffer(/*page_id=*/0); + EXPECT_EQ(loaded, resident); + if (resident != nullptr) { + pool.page_table_.release_block(/*block_id=*/0); + } + pool.page_table_.release_block(/*block_id=*/0); +} + +// --------------------------------------------------------------------------- +// 1. Data stays correct when the working set far exceeds pool capacity, which +// forces the CLOCK evictor to run repeatedly. Also asserts the observability +// counters get populated (hits, misses, evictions). +// --------------------------------------------------------------------------- +TEST_F(BufferPoolTest, DataCorrectUnderEviction) { + const size_t num_pages = 64; + InitVecPool(/*capacity_pages=*/16, + /*file_pages=*/num_pages); // 4x smaller than working set + std::string file = NewFile(num_pages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + std::vector buf(kVectorPageSize); + for (int iter = 0; iter < 3; ++iter) { + for (size_t p = 0; p < num_pages; ++p) { + ASSERT_TRUE( + handle.read_range(p * kVectorPageSize, kVectorPageSize, buf.data())); + ExpectPageContent(buf.data(), p); + } + } + + VecBufferPool::Stats s = pool.stats(); + EXPECT_GT(s.hit + s.miss, 0u); + EXPECT_GT(s.miss, 0u); // capacity < working set => guaranteed misses +} + +// A page encountered by the evictor while pinned stays registered with the +// queue and becomes reclaimable after its final release. This exercises the +// install-time queue registration used to keep release_block() free of the +// steady-state in_evict_queue CAS. +TEST_F(BufferPoolTest, PinnedEvictionBecomesReclaimableAfterRelease) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/2); + std::string file = NewFile(/*num_pages=*/2); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + size_t page_id = 0; + char *page = handle.get_single_page(/*file_offset=*/0, /*len=*/1, page_id); + ASSERT_NE(page, nullptr); + EXPECT_EQ(page_id, 0u); + + // The active pin prevents eviction, but the failed attempt must not make + // the page depend on a release-side CAS to become eligible again. + EXPECT_FALSE(pool.page_table_.evict_block(page_id)); + handle.release_one(page_id); + EXPECT_TRUE(pool.page_table_.evict_block(page_id)); + EXPECT_FALSE(pool.page_table_.is_loaded(page_id)); +} + +// A stale eviction item must not become valid again when a later page table +// reuses the same owner address. Version zero represents an entry issued by a +// different/legacy owner generation; the current resident page must survive. +TEST_F(BufferPoolTest, StaleOwnerGenerationIsDead) { + InitTablePool(/*capacity_pages=*/2, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0), + buffer); + table.release_block(/*block_id=*/0); + + EXPECT_TRUE(table.is_dead_block(/*block_id=*/0, /*stale version=*/0)); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, ForceEvictUnloadedPageDoesNotEnqueueDeadItem) { + InitTablePool(/*capacity_pages=*/0, /*entry_num=*/1); + auto &queue = BlockEvictionQueue::get_instance(); + BlockEvictionQueue::BlockType item; + while (queue.evict_single_block(item)) { + } + + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + EXPECT_FALSE(table.force_evict_block(/*block_id=*/0)); + EXPECT_FALSE(queue.evict_single_block(item)); +} + +TEST_F(BufferPoolTest, ConcurrentInstallPublishesOneResidentBuffer) { + constexpr size_t kThreadCount = 16; + InitTablePool(/*capacity_pages=*/kThreadCount, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + std::array input{}; + std::array result{}; + for (char *&buffer : input) { + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_NE(nullptr, buffer); + } + + std::atomic ready{0}; + std::atomic start{false}; + std::vector workers; + workers.reserve(kThreadCount); + for (size_t i = 0; i < kThreadCount; ++i) { + workers.emplace_back([&, i] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + result[i] = table.set_block_acquired(/*block_id=*/0, input[i], + /*file_offset=*/0); + }); + } + while (ready.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + + ASSERT_NE(nullptr, result[0]); + for (char *buffer : result) { + EXPECT_EQ(result[0], buffer); + table.release_block(/*block_id=*/0); + } + EXPECT_EQ(kVectorPageSize, MemoryLimitPool::get_instance().stats().page_used); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().stats().page_used); +} + +TEST_F(BufferPoolTest, PageLoadClaimCoalescesConcurrentWaiters) { + constexpr size_t kThreadCount = 16; + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + ASSERT_EQ(VectorPageTable::LoadClaimResult::kClaimed, + table.try_claim_block_load(/*block_id=*/0)); + + std::atomic observed_loading{0}; + std::atomic completed{0}; + std::atomic succeeded{0}; + std::atomic release{false}; + std::array workers; + for (auto &worker : workers) { + worker = std::thread([&] { + EXPECT_EQ(VectorPageTable::LoadClaimResult::kLoading, + table.try_claim_block_load(/*block_id=*/0)); + observed_loading.fetch_add(1, std::memory_order_release); + const bool stable = table.wait_for_block_transition(/*block_id=*/0); + EXPECT_TRUE(stable); + char *page = stable ? table.acquire_block(/*block_id=*/0) : nullptr; + EXPECT_NE(nullptr, page); + if (page != nullptr && page[0] == static_cast(0x5a)) { + succeeded.fetch_add(1, std::memory_order_release); + } + completed.fetch_add(1, std::memory_order_release); + while (!release.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + if (page != nullptr) { + table.release_block(/*block_id=*/0); + } + }); + } + + while (observed_loading.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + std::memset(buffer, 0x5a, kVectorPageSize); + ASSERT_EQ(buffer, table.publish_claimed_block(/*block_id=*/0, buffer, + /*file_offset=*/0)); + + while (completed.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + EXPECT_EQ(kThreadCount, succeeded.load(std::memory_order_acquire)); + table.release_block(/*block_id=*/0); + release.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + // The background reclaimer may win the eviction claim but spend the + // page's CLOCK second chance instead of reclaiming it. Retry after each + // transition while keeping the final unloaded-state checks strict. + for (size_t attempt = 0; + attempt < 32 && table.has_residency_activity(/*block_id=*/0); + ++attempt) { + (void)table.force_evict_block(/*block_id=*/0); + ASSERT_TRUE(table.wait_for_block_transition(/*block_id=*/0)); + } + EXPECT_FALSE(table.has_residency_activity(/*block_id=*/0)); + EXPECT_FALSE(table.is_loaded(/*block_id=*/0)); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().stats().page_used); +} + +TEST_F(BufferPoolTest, FailedPageLoadClaimCanBeRetried) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + ASSERT_EQ(VectorPageTable::LoadClaimResult::kClaimed, + table.try_claim_block_load(/*block_id=*/0)); + EXPECT_TRUE(table.cancel_block_load(/*block_id=*/0)); + EXPECT_EQ(VectorPageTable::LoadClaimResult::kClaimed, + table.try_claim_block_load(/*block_id=*/0)); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(buffer, table.publish_claimed_block(/*block_id=*/0, buffer, + /*file_offset=*/0)); + table.release_block(/*block_id=*/0); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, DirtyFlushFailureKeepsPageResident) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0), + buffer); + table.mark_dirty(/*block_id=*/0); + table.release_block(/*block_id=*/0); + + size_t flush_attempts = 0; + table.set_flush_callback([&](block_id_t, char *, size_t, size_t) { + ++flush_attempts; + return -1; + }); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_EQ(1u, flush_attempts); + EXPECT_TRUE(table.is_loaded(/*block_id=*/0)); + EXPECT_TRUE(table.is_block_dirty(/*block_id=*/0)); + EXPECT_EQ(kVectorPageSize, MemoryLimitPool::get_instance().stats().page_used); + + table.set_flush_callback([&](block_id_t, char *, size_t, size_t) { + ++flush_attempts; + return 0; + }); + EXPECT_TRUE(table.evict_block(/*block_id=*/0)); + EXPECT_EQ(2u, flush_attempts); + EXPECT_FALSE(table.is_loaded(/*block_id=*/0)); + EXPECT_FALSE(table.is_block_dirty(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, ConcurrentWritablePressureUsesBackgroundWriteback) { + constexpr size_t kCapacityPages = 4; + constexpr size_t kFilePages = 64; + constexpr size_t kThreadCount = 8; + InitVecPool(kCapacityPages, kFilePages, /*writable=*/true); + // BufferStorage creates a small metadata-only file and grows it as segments + // are appended. Exercise that path instead of opening a pre-sized file. + std::string file = NewFile(/*num_pages=*/1); + + VecBufferPool::Stats final_stats; + { + VecBufferPool pool(file, /*writable=*/true); + ASSERT_EQ(0, pool.init()); + ASSERT_TRUE(pool.extend_file(kFilePages * kVectorPageSize)); + + std::atomic next_page{0}; + std::atomic failures{0}; + std::vector writers; + writers.reserve(kThreadCount); + for (size_t thread_id = 0; thread_id < kThreadCount; ++thread_id) { + writers.emplace_back([&, thread_id] { + std::vector payload(kVectorPageSize, + static_cast(thread_id + 1)); + while (true) { + const size_t page_id = + next_page.fetch_add(1, std::memory_order_relaxed); + if (page_id >= kFilePages) { + break; + } + std::fill(payload.begin(), payload.end(), + static_cast(page_id + 1)); + if (pool.write_range(page_id * kVectorPageSize, kVectorPageSize, + payload.data()) != 0) { + failures.fetch_add(1, std::memory_order_relaxed); + break; + } + } + }); + } + for (auto &writer : writers) { + writer.join(); + } + + EXPECT_EQ(0u, failures.load(std::memory_order_relaxed)); + EXPECT_EQ(0, pool.flush_all()); + final_stats = pool.stats(); + EXPECT_GT(final_stats.writeback_requests, 0u); + EXPECT_GT(final_stats.writeback_batches, 0u); + EXPECT_GT(final_stats.writeback_pages, 0u); + EXPECT_EQ(0u, final_stats.writeback_failures); + EXPECT_EQ(0u, final_stats.writeback_pending); + EXPECT_GT(final_stats.evict, 0u); +#if defined(__linux__) + if (current_io_backend_type() == IOBackendType::kIoUring) { + EXPECT_GT(final_stats.writeback_aio_batches, 0u); + EXPECT_GT(final_stats.writeback_aio_pages, 0u); + EXPECT_EQ(0u, final_stats.writeback_aio_fallbacks); + } +#endif + } + + FILE *input = std::fopen(file.c_str(), "rb"); + ASSERT_NE(nullptr, input); + std::vector page(kVectorPageSize); + for (size_t page_id = 0; page_id < kFilePages; ++page_id) { + ASSERT_EQ(0, std::fseek(input, static_cast(page_id * kVectorPageSize), + SEEK_SET)); + ASSERT_EQ(kVectorPageSize, + std::fread(page.data(), 1, kVectorPageSize, input)); + EXPECT_EQ(static_cast(page_id + 1), page.front()); + EXPECT_EQ(static_cast(page_id + 1), page.back()); + } + std::fclose(input); +} + +TEST_F(BufferPoolTest, RecoversDirtyPageAfterQueueRegistrationFailure) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + size_t flush_attempts = 0; + table.set_flush_callback([&](block_id_t, char *, size_t, size_t) { + ++flush_attempts; + return -1; + }); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0), + buffer); + table.mark_dirty(/*block_id=*/0); + // Force the queue's priority-rewrite path to reject registration. The + // failed flush then leaves a released resident page for recovery to find. + table.set_evict_priority(/*block_id=*/0, std::numeric_limits::max()); + table.release_block(/*block_id=*/0); + EXPECT_EQ(0u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(1u, flush_attempts); + EXPECT_TRUE(table.is_loaded(/*block_id=*/0)); + EXPECT_TRUE(table.is_block_dirty(/*block_id=*/0)); + + table.set_evict_priority(/*block_id=*/0, 0); + EXPECT_EQ(1u, table.recover_eviction_queue()); + table.set_flush_callback([&](block_id_t, char *, size_t, size_t) { + ++flush_attempts; + return 0; + }); + EXPECT_EQ(1u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(2u, flush_attempts); + EXPECT_FALSE(table.is_loaded(/*block_id=*/0)); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().stats().page_used); +} + +TEST_F(BufferPoolTest, MetadataIsCountedAndReleasedWithPool) { + constexpr size_t kPageCount = 2; + const size_t expected_metadata = + VecBufferPool::metadata_bytes_for_page_count(kPageCount); + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/kPageCount); + std::string file = NewFile(kPageCount); + + { + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + const auto stats = MemoryLimitPool::get_instance().stats(); + EXPECT_EQ(expected_metadata, stats.metadata_used); + EXPECT_EQ(expected_metadata, stats.used); + EXPECT_EQ(0u, stats.page_used); + } + + const auto stats = MemoryLimitPool::get_instance().stats(); + EXPECT_EQ(0u, stats.metadata_used); + EXPECT_EQ(0u, stats.used); +} + +TEST_F(BufferPoolTest, FixedMetadataStaysCompactAndReadOnlyAvoidsPageLocks) { + constexpr size_t kOneSegmentPages = 16UL * 1024UL; + const size_t page_table_bytes = + VectorPageTable::metadata_bytes_for_entries(kOneSegmentPages); + const size_t read_only_bytes = + VecBufferPool::metadata_bytes_for_page_count(kOneSegmentPages); + const size_t writable_bytes = VecBufferPool::metadata_bytes_for_page_count( + kOneSegmentPages, /*writable=*/true); + + EXPECT_EQ(page_table_bytes, read_only_bytes); + EXPECT_LT(read_only_bytes, 1UL * 1024UL * 1024UL); + size_t expected_writable_bytes = + VecBufferPool::kMutexBucketCount * sizeof(std::shared_mutex) + + VecBufferPool::kWritebackBatchPages * kVectorPageSize; +#if defined(__linux__) + if (current_io_backend_type() == IOBackendType::kIoUring) { + expected_writable_bytes += + VecBufferPool::kWritebackBatchPages * kVectorPageSize; + } +#endif + EXPECT_EQ(expected_writable_bytes, writable_bytes - read_only_bytes); + + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/1); + std::string file = NewFile(/*num_pages=*/1); + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + const auto stats = pool.stats(); + EXPECT_EQ(VectorPageTable::metadata_bytes_for_entries(1), + stats.page_table_metadata_bytes); + EXPECT_EQ(0u, stats.page_lock_metadata_bytes); +} + +TEST_F(BufferPoolTest, EmptyPageTableChargesDirectoryOnFirstExtend) { + const size_t first_segment_bytes = + VectorPageTable::metadata_bytes_for_entries(1); + ASSERT_EQ(0, MemoryLimitPool::get_instance().init(first_segment_bytes)); + + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/0)); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().metadata_used()); + ASSERT_TRUE(table.extend(/*new_entry_num=*/1)); + EXPECT_EQ(first_segment_bytes, + MemoryLimitPool::get_instance().metadata_used()); + ASSERT_TRUE(table.rollback_extend(/*old_entry_num=*/0)); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().metadata_used()); +} + +TEST_F(BufferPoolTest, WritablePoolReportsPageLockMetadata) { + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/1, /*writable=*/true); + std::string file = NewFile(/*num_pages=*/1); + VecBufferPool pool(file, /*writable=*/true); + ASSERT_EQ(pool.init(), 0); + + const auto stats = pool.stats(); + EXPECT_EQ(VecBufferPool::kMutexBucketCount * sizeof(std::shared_mutex), + stats.page_lock_metadata_bytes); + EXPECT_EQ(VecBufferPool::kWritebackBatchPages * kVectorPageSize, + stats.writeback_staging_bytes); + EXPECT_EQ(VecBufferPool::metadata_bytes_for_page_count( + /*page_count=*/1, /*writable=*/true), + stats.page_table_metadata_bytes + stats.page_lock_metadata_bytes + + stats.writeback_staging_bytes + + stats.writeback_io_staging_bytes); +} + +TEST_F(BufferPoolTest, FailedPageTableExtendLeavesStateUnchanged) { + constexpr size_t kSecondSegmentEntry = 16UL * 1024UL + 1; + InitTablePool(/*capacity_pages=*/0, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + const size_t metadata_before = + MemoryLimitPool::get_instance().metadata_used(); + + EXPECT_FALSE(table.extend(kSecondSegmentEntry)); + EXPECT_EQ(1u, table.entry_num()); + EXPECT_EQ(metadata_before, MemoryLimitPool::get_instance().metadata_used()); +} + +TEST_F(BufferPoolTest, FailedFileExtendDoesNotGrowBackingFile) { + constexpr size_t kSecondSegmentEntry = 16UL * 1024UL + 1; + InitVecPool(/*capacity_pages=*/1, /*file_pages=*/1, /*writable=*/true); + std::string file = NewFile(/*num_pages=*/1); + + VecBufferPool pool(file, /*writable=*/true); + ASSERT_EQ(pool.init(), 0); + const size_t old_size = pool.file_size(); + const size_t old_entries = pool.page_table_.entry_num(); + EXPECT_FALSE(pool.extend_file(kSecondSegmentEntry * kVectorPageSize)); + EXPECT_EQ(old_size, pool.file_size()); + EXPECT_EQ(old_entries, pool.page_table_.entry_num()); + + FILE *backing = std::fopen(file.c_str(), "rb"); + ASSERT_NE(nullptr, backing); + ASSERT_EQ(0, std::fseek(backing, 0, SEEK_END)); + EXPECT_EQ(static_cast(old_size), std::ftell(backing)); + std::fclose(backing); +} + +TEST_F(BufferPoolTest, ExternalReservationSharesThePageBudget) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/4); + + ASSERT_TRUE(memory_pool.try_charge_external(3 * kVectorPageSize)); + EXPECT_EQ(3 * kVectorPageSize, memory_pool.used()); + EXPECT_EQ(3 * kVectorPageSize, memory_pool.external_used()); + EXPECT_EQ(0u, memory_pool.stats().page_used); + + char *page = nullptr; + ASSERT_TRUE(memory_pool.try_acquire_buffer(kVectorPageSize, page)); + ASSERT_NE(nullptr, page); + auto shared = memory_pool.stats(); + EXPECT_EQ(kVectorPageSize, shared.page_used); + EXPECT_EQ(3 * kVectorPageSize, shared.external_used); + EXPECT_FALSE(memory_pool.try_charge_external(1)); + + memory_pool.release_buffer(page, kVectorPageSize); + memory_pool.release_external(3 * kVectorPageSize); + EXPECT_EQ(0u, memory_pool.used()); + EXPECT_EQ(0u, memory_pool.external_used()); +} + +TEST_F(BufferPoolTest, PageAdmissionLeavesRoomForExternalCache) { + auto &memory_pool = MemoryLimitPool::get_instance(); + constexpr size_t kCapacity = 512UL * 1024UL * 1024UL; + ASSERT_EQ(0, memory_pool.init(kCapacity)); + const size_t reserve = memory_pool.page_admission_reserve(); + ASSERT_EQ(32UL * 1024UL * 1024UL, reserve); + + ASSERT_TRUE(memory_pool.try_charge_metadata(kCapacity - reserve)); + char *page = nullptr; + EXPECT_FALSE(memory_pool.try_acquire_buffer(kVectorPageSize, page)); + EXPECT_EQ(nullptr, page); + EXPECT_TRUE(memory_pool.try_charge_external(reserve)); + + memory_pool.release_external(reserve); + memory_pool.release_metadata(kCapacity - reserve); + EXPECT_EQ(0u, memory_pool.used()); +} + +TEST_F(BufferPoolTest, ReadOnlyMissEvictsAtPageAdmissionLimit) { + auto &memory_pool = MemoryLimitPool::get_instance(); + constexpr size_t kCapacity = 256UL * 1024UL * 1024UL; + ASSERT_EQ(0, memory_pool.init(kCapacity)); + const size_t reserve = memory_pool.page_admission_reserve(); + ASSERT_EQ(16UL * 1024UL * 1024UL, reserve); + + std::string file = NewFile(/*num_pages=*/2); + const size_t pool_metadata = + VecBufferPool::metadata_bytes_for_page_count(/*page_count=*/2, + /*writable=*/false); + ASSERT_LT(pool_metadata + kVectorPageSize, kCapacity - reserve); + const size_t charged_metadata = + kCapacity - reserve - pool_metadata - kVectorPageSize; + ASSERT_TRUE(memory_pool.try_charge_metadata(charged_metadata)); + + { + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(0, pool.init()); + + char *first = pool.acquire_buffer(/*page_id=*/0); + ASSERT_NE(nullptr, first); + ExpectPageContent(first, /*page_id=*/0); + pool.page_table_.release_block(/*block_id=*/0); + + // The reserved headroom means the process-wide pool is not full, but the + // next page allocation has reached its page-specific admission limit. + EXPECT_FALSE(memory_pool.is_full()); + EXPECT_TRUE(memory_pool.under_cache_pressure()); + char *second = pool.acquire_buffer(/*page_id=*/1, /*retry=*/50); + ASSERT_NE(nullptr, second); + ExpectPageContent(second, /*page_id=*/1); + EXPECT_GT(pool.stats().evict, 0u); + pool.page_table_.release_block(/*block_id=*/1); + } + + memory_pool.release_metadata(charged_metadata); + EXPECT_EQ(0u, memory_pool.used()); +} + +TEST_F(BufferPoolTest, TinyBufferDoesNotPoisonThePageFreeList) { + auto &memory_pool = MemoryLimitPool::get_instance(); + ASSERT_EQ(0, memory_pool.init(3 * kVectorPageSize + 123)); + + char *tiny = nullptr; + ASSERT_TRUE(memory_pool.try_acquire_buffer(1, tiny)); + ASSERT_NE(nullptr, tiny); + memory_pool.release_buffer(tiny, 1); + EXPECT_EQ(0u, memory_pool.committed()); + + char *page = nullptr; + ASSERT_TRUE(memory_pool.try_acquire_buffer(kVectorPageSize, page)); + ASSERT_NE(nullptr, page); + memory_pool.release_buffer(page, kVectorPageSize); + EXPECT_EQ(1u, memory_pool.stats().free_buffers); +} + +TEST_F(BufferPoolTest, RejectsReinitializationWhileMemoryIsActive) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/4); + const size_t original_capacity = memory_pool.capacity(); + + ASSERT_TRUE(memory_pool.try_charge_external(kVectorPageSize)); + EXPECT_EQ(0, memory_pool.init(original_capacity)); + EXPECT_EQ(kVectorPageSize, memory_pool.external_used()); + EXPECT_NE(0, memory_pool.init(8 * kVectorPageSize)); + EXPECT_EQ(original_capacity, memory_pool.capacity()); + EXPECT_EQ(kVectorPageSize, memory_pool.used()); + EXPECT_EQ(kVectorPageSize, memory_pool.external_used()); + + memory_pool.release_external(kVectorPageSize); + ASSERT_EQ(0, memory_pool.init(8 * kVectorPageSize)); + EXPECT_EQ(8 * kVectorPageSize, memory_pool.capacity()); +} + +TEST_F(BufferPoolTest, ExternalCacheRejectsOversizedEntryAndReleasesOnDestroy) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/2); + + { + SizedExternalCache cache; + EXPECT_EQ(nullptr, cache.acquire(3 * kVectorPageSize)); + EXPECT_EQ(0u, cache.entry_count()); + EXPECT_EQ(0u, memory_pool.used()); + + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + EXPECT_EQ(kVectorPageSize, memory_pool.used()); + cache.release(kVectorPageSize); + } + + EXPECT_EQ(0u, memory_pool.used()); + EXPECT_EQ(0u, memory_pool.committed()); +} + +TEST_F(BufferPoolTest, + ExternalCacheReclaimsEntryAfterQueueRegistrationFailure) { + auto &memory_pool = MemoryLimitPool::get_instance(); + // Keep usage below the background high watermark so only the simulated + // enqueue-failure callback can reclaim this entry during the assertion. + InitPool(/*capacity_pages=*/2); + + SizedExternalCache cache; + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + cache.release(kVectorPageSize); + EXPECT_EQ(kVectorPageSize, memory_pool.external_used()); + + constexpr eviction_key_t kOwnerKey = 1; + version_t version = FindLiveVersion(cache, kOwnerKey); + ASSERT_NE(0u, version); + cache.eviction_requeue_failed(kOwnerKey, version); + EXPECT_EQ(0u, memory_pool.external_used()); + EXPECT_EQ(0u, cache.entry_count()); + EXPECT_EQ(nullptr, cache.retain(kVectorPageSize)); +} + +TEST_F(BufferPoolTest, ExternalCacheReusesOneQueueMembershipAcrossHits) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/4); + + SizedExternalCache cache; + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + cache.release(kVectorPageSize); + + constexpr eviction_key_t kOwnerKey = 1; + const version_t version = FindLiveVersion(cache, kOwnerKey); + ASSERT_NE(0u, version); + + // Repeated 0 -> 1 -> 0 transitions must keep using the existing logical + // queue item. Generating a new version/item for every hit lets stale queue + // nodes grow without bound while usage remains below the low watermark. + for (size_t i = 0; i < 10000; ++i) { + value = cache.retain(kVectorPageSize); + ASSERT_NE(nullptr, value); + cache.release(kVectorPageSize); + } + EXPECT_FALSE(cache.is_dead_block(kOwnerKey, version)); + + EXPECT_EQ(1u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(0u, memory_pool.external_used()); + EXPECT_EQ(0u, cache.entry_count()); + EXPECT_EQ(nullptr, cache.retain(kVectorPageSize)); +} + +TEST_F(BufferPoolTest, PinnedExternalCacheEntryStaysQueuedForLaterEviction) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/4); + + SizedExternalCache cache; + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + cache.release(kVectorPageSize); + + value = cache.retain(kVectorPageSize); + ASSERT_NE(nullptr, value); + EXPECT_EQ(0u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(kVectorPageSize, memory_pool.external_used()); + + cache.release(kVectorPageSize); + EXPECT_EQ(1u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(0u, memory_pool.external_used()); + EXPECT_EQ(0u, cache.entry_count()); +} + +TEST_F(BufferPoolTest, ConcurrentLoadsUseSingleFlight) { + constexpr size_t kThreadCount = 16; + InitPool(/*capacity_pages=*/4); + auto state = std::make_shared(); + BlockingExternalCache cache(BlockingLoader{state}); + std::atomic acquired{0}; + std::atomic release{false}; + + std::array workers; + for (auto &worker : workers) { + worker = std::thread([&] { + auto value = cache.acquire(kVectorPageSize); + EXPECT_NE(nullptr, value); + acquired.fetch_add(1, std::memory_order_release); + while (!release.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + cache.release(kVectorPageSize); + }); + } + while (state->load_calls.load(std::memory_order_acquire) == 0) { + std::this_thread::yield(); + } + EXPECT_EQ(1u, cache.entry_count()); + state->finish.store(true, std::memory_order_release); + while (acquired.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + EXPECT_EQ(1u, state->load_calls.load(std::memory_order_acquire)); + EXPECT_EQ(kVectorPageSize, MemoryLimitPool::get_instance().external_used()); + release.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + EXPECT_EQ(1u, BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(0u, cache.entry_count()); +} + +TEST_F(BufferPoolTest, DistinctExternalCacheLoadsAreConcurrentByDefault) { + constexpr size_t kThreadCount = 2; + InitPool(/*capacity_pages=*/4); + auto state = std::make_shared(); + BlockingExternalCache cache(BlockingLoader{state}); + std::array>, kThreadCount> values; + const std::array keys = {kVectorPageSize, + kVectorPageSize + 1}; + + std::array workers; + for (size_t i = 0; i < kThreadCount; ++i) { + workers[i] = std::thread([&, i] { values[i] = cache.acquire(keys[i]); }); + } + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (state->load_calls.load(std::memory_order_acquire) != kThreadCount && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + const size_t concurrent_loads = + state->load_calls.load(std::memory_order_acquire); + state->finish.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + + EXPECT_EQ(kThreadCount, concurrent_loads); + EXPECT_EQ(kThreadCount, + state->max_active_loads.load(std::memory_order_acquire)); + for (size_t i = 0; i < kThreadCount; ++i) { + ASSERT_NE(nullptr, values[i]); + cache.release(keys[i]); + } +} + +TEST_F(BufferPoolTest, DistinctExternalCacheLoadsRespectInflightLimit) { + constexpr size_t kThreadCount = 2; + InitPool(/*capacity_pages=*/4); + auto state = std::make_shared(); + BlockingExternalCache cache(BlockingLoader{state}, + /*max_concurrent_loads=*/1); + std::atomic ready{0}; + std::atomic start{false}; + std::array>, kThreadCount> values; + const std::array keys = {kVectorPageSize, + kVectorPageSize + 1}; + + std::array workers; + for (size_t i = 0; i < kThreadCount; ++i) { + workers[i] = std::thread([&, i] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + values[i] = cache.acquire(keys[i]); + }); + } + while (ready.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + while (state->load_calls.load(std::memory_order_acquire) == 0) { + std::this_thread::yield(); + } + state->finish.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + + EXPECT_EQ(kThreadCount, state->load_calls.load(std::memory_order_acquire)); + EXPECT_EQ(1u, state->max_active_loads.load(std::memory_order_acquire)); + for (size_t i = 0; i < kThreadCount; ++i) { + ASSERT_NE(nullptr, values[i]); + cache.release(keys[i]); + } +} + +TEST_F(BufferPoolTest, WaitingExternalLoadRechecksCapacityBeforeLoading) { + InitPool(/*capacity_pages=*/1); + auto state = std::make_shared(); + BlockingExternalCache cache(BlockingLoader{state}, + /*max_concurrent_loads=*/1); + std::shared_ptr> first_value; + std::shared_ptr> second_value; + + std::thread first([&] { first_value = cache.acquire(kVectorPageSize); }); + while (state->load_calls.load(std::memory_order_acquire) == 0) { + std::this_thread::yield(); + } + std::thread second( + [&] { second_value = cache.acquire(kVectorPageSize + 1); }); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + state->finish.store(true, std::memory_order_release); + first.join(); + second.join(); + + ASSERT_NE(nullptr, first_value); + EXPECT_EQ(nullptr, second_value); + EXPECT_EQ(1u, state->load_calls.load(std::memory_order_acquire)); + cache.release(kVectorPageSize); +} + +TEST_F(BufferPoolTest, ThrowingPayloadConstructorDoesNotClaimLoaderSlot) { + InitPool(/*capacity_pages=*/2); + ThrowingCachePayload::construction_count = 0; + ThrowingCachePayload::throw_on_construction = 2; + ThrowingExternalCache cache(/*max_concurrent_loads=*/1); + + EXPECT_THROW(cache.acquire(kVectorPageSize), std::runtime_error); + ASSERT_EQ(0u, cache.entry_count()); + + ThrowingCachePayload::throw_on_construction = 0; + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + cache.release(kVectorPageSize); +} + +TEST_F(BufferPoolTest, EmptyLoaderValueRollsBackChargeAndPlaceholder) { + InitPool(/*capacity_pages=*/2); + EmptyValueExternalCache cache; + + EXPECT_THROW(cache.acquire(kVectorPageSize), std::runtime_error); + EXPECT_EQ(0u, cache.entry_count()); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().external_used()); + + // The failed placeholder and single-flight state must not poison retries. + EXPECT_THROW(cache.acquire(kVectorPageSize), std::runtime_error); + EXPECT_EQ(0u, cache.entry_count()); + EXPECT_EQ(0u, MemoryLimitPool::get_instance().external_used()); +} + +TEST_F(BufferPoolTest, BatchRecycleBoundsStaleQueueScanning) { + auto &queue = BlockEvictionQueue::get_instance(); + BlockEvictionQueue::BlockType discarded; + while (queue.evict_single_block(discarded)) { + } + + AlwaysDeadOwner owner; + BlockEvictionQueue::BlockType stale; + stale.owner = &owner; + stale.version = 1; + for (size_t i = 0; i < 64; ++i) { + stale.owner_key = i; + ASSERT_TRUE(queue.add_single_block(stale, 0)); + } + + EXPECT_EQ(0u, queue.batch_recycle(1)); + EXPECT_EQ(20u, owner.dead_checks); + + queue.set_invalid(&owner); + while (queue.evict_single_block(discarded)) { + } +} + +TEST_F(BufferPoolTest, HighCardinalityEvictionRemovesKeyMetadata) { + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(/*capacity_pages=*/8); + SizedExternalCache cache; + + for (size_t key = 1; key <= 256; ++key) { + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value) << "key=" << key; + cache.release(kVectorPageSize); + } + for (size_t attempt = 0; memory_pool.external_used() != 0 && attempt < 256; + ++attempt) { + BlockEvictionQueue::get_instance().batch_recycle(64); + } + EXPECT_EQ(0u, memory_pool.external_used()); + EXPECT_EQ(0u, cache.entry_count()); +} + +TEST_F(BufferPoolTest, OwningHandleRejectsNullPool) { + EXPECT_THROW( + { + VecBufferPoolHandle handle(std::shared_ptr{}); + (void)handle; + }, + std::invalid_argument); +} + +TEST_F(BufferPoolTest, ExternalCacheRejectsStaleItemAfterAddressReuse) { + InitPool(/*capacity_pages=*/4); + alignas(SizedExternalCache) unsigned char storage[sizeof(SizedExternalCache)]; + constexpr eviction_key_t kOwnerKey = 1; + + auto *first = new (storage) SizedExternalCache(); + auto first_value = first->acquire(kVectorPageSize); + if (first_value == nullptr) { + first->~SizedExternalCache(); + FAIL() << "failed to populate first cache"; + } + first->release(kVectorPageSize); + version_t stale_version = FindLiveVersion(*first, kOwnerKey); + first->~SizedExternalCache(); + ASSERT_NE(0u, stale_version); + + auto *second = new (storage) SizedExternalCache(); + auto second_value = second->acquire(kVectorPageSize); + if (second_value == nullptr) { + second->~SizedExternalCache(); + FAIL() << "failed to populate replacement cache"; + } + second->release(kVectorPageSize); + version_t current_version = FindLiveVersion(*second, kOwnerKey); + + EXPECT_NE(0u, current_version); + EXPECT_NE(stale_version, current_version); + EXPECT_TRUE(second->is_dead_block(kOwnerKey, stale_version)); + second->~SizedExternalCache(); +} + +TEST_F(BufferPoolTest, ExternalCacheUsesExplicitHotEvictionPriority) { + InitPool(/*capacity_pages=*/2); + SizedExternalCache cache; + auto value = cache.acquire(kVectorPageSize); + ASSERT_NE(nullptr, value); + EXPECT_EQ(BlockEvictionQueue::kExplicitHotPriority, + cache.eviction_priority(/*owner_key=*/1)); + cache.release(kVectorPageSize); +} + +TEST_F(BufferPoolTest, ExternalReservationTrimsRetainedPageBuffers) { + constexpr size_t kCapacityPages = 4; + auto &memory_pool = MemoryLimitPool::get_instance(); + InitPool(kCapacityPages); + + std::vector pages; + for (size_t i = 0; i < kCapacityPages; ++i) { + char *page = nullptr; + ASSERT_TRUE(memory_pool.try_acquire_buffer(kVectorPageSize, page)); + pages.push_back(page); + } + for (char *page : pages) { + memory_pool.release_buffer(page, kVectorPageSize); + } + + MemoryLimitPool::PoolStats cached = memory_pool.stats(); + EXPECT_EQ(0u, cached.used); + EXPECT_EQ(kCapacityPages * kVectorPageSize, cached.committed); + EXPECT_EQ(kCapacityPages, cached.free_buffers); + EXPECT_EQ(1u, cached.slab_count); + const uint64_t reclaimed_before = cached.slab_reclaimed_pages; + + ASSERT_TRUE( + memory_pool.try_charge_external(kCapacityPages * kVectorPageSize)); + MemoryLimitPool::PoolStats charged = memory_pool.stats(); + EXPECT_EQ(kCapacityPages * kVectorPageSize, charged.used); + EXPECT_EQ(kCapacityPages * kVectorPageSize, charged.committed); + EXPECT_EQ(0u, charged.page_used); + EXPECT_EQ(kCapacityPages * kVectorPageSize, charged.external_used); + EXPECT_EQ(0u, charged.free_buffers); + EXPECT_EQ(1u, charged.slab_count); + EXPECT_GE(charged.slab_reclaimed_pages, reclaimed_before + kCapacityPages); + + memory_pool.release_external(kCapacityPages * kVectorPageSize); + EXPECT_EQ(0u, memory_pool.used()); + EXPECT_EQ(0u, memory_pool.committed()); + + pages.clear(); + for (size_t i = 0; i < kCapacityPages; ++i) { + char *page = nullptr; + ASSERT_TRUE(memory_pool.try_acquire_buffer(kVectorPageSize, page)); + ASSERT_NE(nullptr, page); + std::memset(page, static_cast(i + 1), kVectorPageSize); + EXPECT_EQ(static_cast(i + 1), page[0]); + EXPECT_EQ(static_cast(i + 1), page[kVectorPageSize - 1]); + pages.push_back(page); + } + EXPECT_EQ(kCapacityPages * kVectorPageSize, memory_pool.committed()); + for (char *page : pages) { + memory_pool.release_buffer(page, kVectorPageSize); + } +} + +TEST_F(BufferPoolTest, LargeExternalReservationReclaimsMultipleBatches) { + constexpr size_t kCapacityPages = 512; + constexpr size_t kExternalPages = 400; + auto &memory_pool = MemoryLimitPool::get_instance(); + InitVecPool(kCapacityPages, /*file_pages=*/kCapacityPages); + std::string file = NewFile(kCapacityPages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + for (size_t page = 0; page < kCapacityPages; ++page) { + ASSERT_TRUE(handle.read_range(page * kVectorPageSize, kVectorPageSize, + data.data())); + } + + ASSERT_TRUE( + memory_pool.try_charge_external(kExternalPages * kVectorPageSize)); + EXPECT_LE(memory_pool.used(), memory_pool.capacity()); + memory_pool.release_external(kExternalPages * kVectorPageSize); +} + +TEST_F(BufferPoolTest, PriorityChangeMigratesQueuedPageBeforeEviction) { + InitVecPool(/*capacity_pages=*/4, /*file_pages=*/2); + std::string file = NewFile(/*num_pages=*/2); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + ASSERT_TRUE(handle.read_range(kVectorPageSize, kVectorPageSize, data.data())); + + ASSERT_TRUE(pool.set_page_priority(0, VecBufferPool::kHighPriority)); + ASSERT_TRUE(pool.set_page_priority(1, VecBufferPool::kLowPriority)); + ASSERT_EQ(1u, BlockEvictionQueue::get_instance().batch_recycle(1)); + + EXPECT_TRUE(pool.is_page_resident(0)); + EXPECT_FALSE(pool.is_page_resident(1)); +} + +TEST_F(BufferPoolTest, DominantProtectedQueueReceivesAgingSamples) { + auto &queue = BlockEvictionQueue::get_instance(); + BlockEvictionQueue::BlockType item; + while (queue.evict_single_block(item)) { + } + + BlockEvictionQueue::BlockType block; + for (size_t i = 0; i < 64; ++i) { + ASSERT_TRUE( + queue.add_single_block(block, BlockEvictionQueue::kProbationPriority)); + } + for (size_t i = 0; i < 256; ++i) { + ASSERT_TRUE( + queue.add_single_block(block, BlockEvictionQueue::kProtectedPriority)); + } + + const uint64_t aging_before = queue.stats().protected_aging_dequeues; + // A reclaim batch samples the protected queue at most once; scalar queue + // inspection deliberately stays strict-priority and pays no aging cost. + (void)queue.batch_recycle(/*count=*/8); + EXPECT_GT(queue.stats().protected_aging_dequeues, aging_before); + + while (queue.evict_single_block(item)) { + } +} + +TEST_F(BufferPoolTest, ReadOnlyPoolDefersAdaptivePriorityUntilPressure) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/1); + std::string file = NewFile(/*num_pages=*/1); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + + EXPECT_EQ(VecBufferPool::kLowPriority, pool.page_table_.eviction_priority(0)); + EXPECT_EQ(0u, + pool.stats().priority_promotions[VecBufferPool::kNormalPriority]); +} + +TEST_F(BufferPoolTest, ReusedReadOnlyPagePromotesAfterPressure) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + + for (size_t page = 0; page < 3; ++page) { + ASSERT_TRUE(handle.read_range(page * kVectorPageSize, kVectorPageSize, + data.data())); + } + ASSERT_GT(pool.stats().evict, 0u); + + // Reuse promotion is sampled under pressure. Any run of 16 hits contains a + // policy sample regardless of the thread-local cursor's starting phase. + for (size_t i = 0; i < 16; ++i) { + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + } + EXPECT_EQ(VecBufferPool::kNormalPriority, + pool.page_table_.eviction_priority(0)); + const auto stats = pool.stats(); + EXPECT_EQ(1u, stats.priority_promotions[VecBufferPool::kNormalPriority]); + // Residency is not stable after the final read releases its pin: the + // background reclaimer may run between assertions. The priority and + // promotion counter are the durable policy outcomes under test. +} + +TEST_F(BufferPoolTest, ProtectedPageAgesThroughProbationBeforeEviction) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(buffer, + table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0)); + table.release_block(/*block_id=*/0); + + ASSERT_TRUE(table.promote_evict_priority( + /*block_id=*/0, VectorPageTable::kNormalPriority)); + ASSERT_EQ(VectorPageTable::kNormalPriority, + table.eviction_priority(/*owner_key=*/0)); + + // One CLOCK turn consumes recent activity and the next demotes the page. + // Demotion itself is the probation turn; no unconditional extra second + // chance is added under pressure. + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_EQ(VectorPageTable::kNormalPriority, + table.eviction_priority(/*owner_key=*/0)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_EQ(VectorPageTable::kLowPriority, + table.eviction_priority(/*owner_key=*/0)); + EXPECT_TRUE(table.evict_block(/*block_id=*/0)); + + const auto stats = table.stats(); + EXPECT_EQ(1u, stats.priority_promotions[VectorPageTable::kNormalPriority]); + EXPECT_EQ(1u, stats.priority_demotions[VectorPageTable::kLowPriority]); + EXPECT_EQ(1u, stats.evictions_by_priority[VectorPageTable::kLowPriority]); +} + +TEST_F(BufferPoolTest, EvictedHotPageGetsProtectedGhostAdmission) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + char *buffer = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + ASSERT_EQ(buffer, + table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0)); + table.release_block(/*block_id=*/0); + + ASSERT_TRUE(table.promote_evict_priority( + /*block_id=*/0, VectorPageTable::kNormalPriority)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); // consume reference + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); // remember and demote + ASSERT_TRUE(table.evict_block(/*block_id=*/0)); + ASSERT_EQ(1u, table.stats().ghost_hot_marks); + + char *reloaded = nullptr; + ASSERT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, reloaded)); + ASSERT_EQ(reloaded, table.set_block_acquired(/*block_id=*/0, reloaded, + /*offset=*/0)); + EXPECT_EQ(VectorPageTable::kNormalPriority, + table.eviction_priority(/*owner_key=*/0)); + EXPECT_EQ(1u, table.stats().ghost_hot_hits); + table.release_block(/*block_id=*/0); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, UnusedGhostAdmissionDoesNotRenewItself) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + auto load_page = [&table] { + char *buffer = nullptr; + EXPECT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + if (buffer == nullptr) return false; + char *installed = + table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0); + EXPECT_EQ(buffer, installed); + if (installed == nullptr) return false; + table.release_block(/*block_id=*/0); + return true; + }; + + ASSERT_TRUE(load_page()); + ASSERT_TRUE(table.promote_evict_priority( + /*block_id=*/0, VectorPageTable::kNormalPriority)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + ASSERT_TRUE(table.evict_block(/*block_id=*/0)); + + ASSERT_TRUE(load_page()); + ASSERT_EQ(VectorPageTable::kNormalPriority, + table.eviction_priority(/*owner_key=*/0)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + ASSERT_TRUE(table.evict_block(/*block_id=*/0)); + + ASSERT_TRUE(load_page()); + EXPECT_EQ(VectorPageTable::kLowPriority, + table.eviction_priority(/*owner_key=*/0)); + const auto stats = table.stats(); + EXPECT_EQ(1u, stats.ghost_hot_marks); + EXPECT_EQ(1u, stats.ghost_hot_hits); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, ReusedGhostAdmissionRenewsHotHistory) { + InitTablePool(/*capacity_pages=*/1, /*entry_num=*/1); + VectorPageTable table; + ASSERT_TRUE(table.init(/*entry_num=*/1)); + + auto load_page = [&table] { + char *buffer = nullptr; + EXPECT_TRUE(MemoryLimitPool::get_instance().try_acquire_buffer( + kVectorPageSize, buffer)); + if (buffer == nullptr) return false; + char *installed = + table.set_block_acquired(/*block_id=*/0, buffer, /*offset=*/0); + EXPECT_EQ(buffer, installed); + if (installed == nullptr) return false; + table.release_block(/*block_id=*/0); + return true; + }; + auto age_and_evict = [&table] { + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + EXPECT_FALSE(table.evict_block(/*block_id=*/0)); + return table.evict_block(/*block_id=*/0); + }; + + ASSERT_TRUE(load_page()); + ASSERT_TRUE(table.promote_evict_priority( + /*block_id=*/0, VectorPageTable::kNormalPriority)); + ASSERT_TRUE(age_and_evict()); + + ASSERT_TRUE(load_page()); + // Validate the ghost admission through the sampled reuse path. + for (size_t i = 0; i < 16; ++i) { + char *reused = table.acquire_block(/*block_id=*/0); + ASSERT_NE(nullptr, reused); + table.release_block(/*block_id=*/0); + } + ASSERT_TRUE(age_and_evict()); + + ASSERT_TRUE(load_page()); + EXPECT_EQ(VectorPageTable::kNormalPriority, + table.eviction_priority(/*owner_key=*/0)); + const auto stats = table.stats(); + EXPECT_EQ(2u, stats.ghost_hot_marks); + EXPECT_EQ(2u, stats.ghost_hot_hits); + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +TEST_F(BufferPoolTest, ReusedPageSurvivesContinuousColdStream) { + constexpr size_t kFilePages = 32; + InitVecPool(/*capacity_pages=*/3, /*file_pages=*/kFilePages); + std::string file = NewFile(kFilePages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + + for (size_t page = 1; page <= 4; ++page) { + ASSERT_TRUE(handle.read_range(page * kVectorPageSize, kVectorPageSize, + data.data())); + } + ASSERT_GT(pool.stats().evict, 0u); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + for (size_t page = 5; page < kFilePages; ++page) { + ASSERT_TRUE(handle.read_range(page * kVectorPageSize, kVectorPageSize, + data.data())); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + } + + EXPECT_TRUE(pool.is_page_resident(0)); + const auto stats = pool.stats(); + EXPECT_LT(stats.miss, kFilePages + kFilePages / 2); + EXPECT_GT(stats.priority_promotions[VecBufferPool::kNormalPriority], 0u); + EXPECT_GT(stats.evictions_by_priority[VecBufferPool::kLowPriority], 0u); + EXPECT_EQ(0u, stats.evictions_by_priority[VecBufferPool::kNormalPriority]); +} + +TEST_F(BufferPoolTest, WritablePoolDoesNotAdaptReadPriority) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/1, /*writable=*/true); + std::string file = NewFile(/*num_pages=*/1); + + VecBufferPool pool(file, /*writable=*/true); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + + EXPECT_EQ(VecBufferPool::kLowPriority, pool.page_table_.eviction_priority(0)); + EXPECT_EQ(0u, + pool.stats().priority_promotions[VecBufferPool::kNormalPriority]); +} + +TEST_F(BufferPoolTest, BypassReadDoesNotAdmitPage) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + ASSERT_TRUE( + handle.read_range_bypass(kVectorPageSize, kVectorPageSize, data.data())); + + ExpectPageContent(data.data(), 1); + EXPECT_FALSE(pool.is_page_resident(1)); + auto stats = pool.stats(); + EXPECT_EQ(1u, stats.bypass_reads); + EXPECT_EQ(kVectorPageSize, stats.bypass_bytes); + EXPECT_EQ(1u, stats.bypass_io_requests); + EXPECT_EQ(0u, stats.miss); +} + +TEST_F(BufferPoolTest, ReadAndPrefetchRangesRejectOverflow) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/2); + std::string file = NewFile(/*num_pages=*/2); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + + EXPECT_FALSE( + handle.read_range(std::numeric_limits::max(), 2, data.data())); + EXPECT_FALSE(handle.read_range(pool.file_size() - 1, 2, data.data())); + handle.prefetch_range(std::numeric_limits::max(), + std::numeric_limits::max()); + EXPECT_EQ(0u, pool.stats().miss); +} + +#if defined(__linux__) +TEST_F(BufferPoolTest, AioAdmissionUsesFreeCapacityBeforeEviction) { + InitVecPool(/*capacity_pages=*/8, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + if (!pool.aio_enabled()) { + GTEST_SKIP() << "no asynchronous backend is available"; + } + EXPECT_EQ(current_io_backend_type(), pool.io_backend_type()); + EXPECT_NE(IOBackendType::kPread, pool.io_backend_type()); + + char *resident = pool.acquire_buffer(/*page_id=*/0); + ASSERT_NE(nullptr, resident); + pool.page_table_.release_block(/*block_id=*/0); + const uint64_t evictions_before = pool.stats().evict; + + pool.prefetch_pages_aio(/*first_page=*/1, /*page_count=*/2); + + EXPECT_TRUE(pool.is_page_resident(0)); + EXPECT_TRUE(pool.is_page_resident(1)); + EXPECT_TRUE(pool.is_page_resident(2)); + EXPECT_EQ(evictions_before, pool.stats().evict); +} +#endif + +// Scattered acquisition is storage-level functionality: it preserves caller +// order, deduplicates cold I/O internally, and still returns one independent +// pin for every occurrence of a duplicate page id. +TEST_F(BufferPoolTest, BatchAcquireScatteredPagesWithDuplicates) { + InitVecPool(/*capacity_pages=*/16, /*file_pages=*/32); + std::string file = NewFile(/*num_pages=*/32); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + const block_id_t page_ids[] = {7, 1, 7, 31, 0, 16}; + char *pages[sizeof(page_ids) / sizeof(page_ids[0])] = {}; + constexpr size_t count = sizeof(page_ids) / sizeof(page_ids[0]); + + ASSERT_TRUE(handle.acquire_pages(page_ids, count, pages)); + for (size_t i = 0; i < count; ++i) { + ASSERT_NE(pages[i], nullptr); + ExpectPageContent(pages[i], page_ids[i]); + } + EXPECT_EQ(pages[0], pages[2]); + + handle.release_pages(page_ids, count); + for (block_id_t page_id : page_ids) { + EXPECT_TRUE(pool.page_table_.is_released(page_id)); + } +} + +TEST_F(BufferPoolTest, ConcurrentBatchLoadsPopulateEachPageOnce) { + static constexpr size_t kPageCount = 128; + constexpr size_t kThreadCount = 8; + InitVecPool(/*capacity_pages=*/512, /*file_pages=*/kPageCount); + std::string file = NewFile(kPageCount); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::array page_ids{}; + for (size_t i = 0; i < kPageCount; ++i) { + page_ids[i] = static_cast(i); + } + + std::atomic ready{0}; + std::atomic start{false}; + std::atomic succeeded{0}; + std::array workers; + for (auto &worker : workers) { + worker = std::thread([&] { + std::array pages{}; + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + if (!handle.acquire_pages(page_ids.data(), page_ids.size(), + pages.data())) { + return; + } + bool valid = true; + for (size_t i = 0; i < kPageCount; ++i) { + valid = valid && pages[i] != nullptr && + pages[i][0] == static_cast(i & 0xff); + } + if (valid) { + succeeded.fetch_add(1, std::memory_order_release); + } + handle.release_pages(page_ids.data(), page_ids.size()); + }); + } + while (ready.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto &worker : workers) { + worker.join(); + } + + EXPECT_EQ(kThreadCount, succeeded.load(std::memory_order_acquire)); + EXPECT_EQ(kPageCount, pool.stats().miss); +} + +TEST_F(BufferPoolTest, BatchMissesRemainProbationUntilLaterReuse) { + InitVecPool(/*capacity_pages=*/2, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + std::vector data(kVectorPageSize); + for (size_t page = 2; page < 4; ++page) { + ASSERT_TRUE(handle.read_range(page * kVectorPageSize, kVectorPageSize, + data.data())); + } + ASSERT_TRUE(handle.read_range(0, kVectorPageSize, data.data())); + ASSERT_GT(pool.stats().evict, 0u); + + const block_id_t page_ids[] = {1}; + char *pages[1] = {}; + + ASSERT_TRUE(handle.acquire_pages(page_ids, 1, pages)); + handle.release_pages(page_ids, 1); + EXPECT_EQ(VecBufferPool::kLowPriority, pool.page_table_.eviction_priority(1)); + + for (size_t i = 0; i < 16; ++i) { + ASSERT_TRUE(handle.acquire_pages(page_ids, 1, pages)); + handle.release_pages(page_ids, 1); + } + EXPECT_EQ(VecBufferPool::kNormalPriority, + pool.page_table_.eviction_priority(1)); +} + +TEST_F(BufferPoolTest, BatchAcquireRollsBackPinsOnInvalidPage) { + InitVecPool(/*capacity_pages=*/4, /*file_pages=*/4); + std::string file = NewFile(/*num_pages=*/4); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + const block_id_t page_ids[] = {1, 4}; + char *pages[2] = {}; + EXPECT_FALSE(handle.acquire_pages(page_ids, 2, pages)); + EXPECT_EQ(pages[0], nullptr); + EXPECT_EQ(pages[1], nullptr); + EXPECT_TRUE(pool.page_table_.is_released(1)); +} + +// --------------------------------------------------------------------------- +// 2. Re-touching a small hot set under memory pressure should trigger the CLOCK +// second-chance path (pages spared instead of evicted) and keep them hot. +// --------------------------------------------------------------------------- +TEST_F(BufferPoolTest, SecondChanceKeepsHotSet) { + const size_t num_pages = 128; + InitVecPool(/*capacity_pages=*/32, /*file_pages=*/num_pages); + std::string file = NewFile(num_pages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + std::vector buf(kVectorPageSize); + auto read_page = [&](size_t p) { + ASSERT_TRUE( + handle.read_range(p * kVectorPageSize, kVectorPageSize, buf.data())); + ExpectPageContent(buf.data(), p); + }; + + // Keep a small hot set (0..7) genuinely hot by re-touching it frequently + // during the cold scan so its reuse distance stays below the pool capacity + // (32). A plain round-by-round scan touches every page once per round + // (reuse distance 128 >> capacity): the hot set is evicted before it can be + // re-hit, which is correct scan-resistant behavior but never exercises the + // second-chance path. Interleaving creates real reuse -- the hot pages + // stay resident (hits) and carry a set reference bit when the evictor + // reaches them, so it spares them (second chance). + for (int round = 0; round < 20; ++round) { + for (size_t c = 8; c < num_pages; ++c) { + read_page(c); // cold churn + if ((c & 7u) == 0u) { // every 8 cold pages... + for (size_t h = 0; h < 8; ++h) read_page(h); // ...re-touch hot set + } + } + } + + VecBufferPool::Stats s = pool.stats(); + EXPECT_GT(s.hit, 0u); + EXPECT_GT(s.evict, 0u); + + // The workload above intentionally leaves eviction scheduling + // nondeterministic. Pin one resident page while setting its CLOCK bit so + // either this thread or the background reclaimer must consume a second + // chance after the pin is released. + const block_id_t hot_page = 0; + char *hot_buffer = nullptr; + ASSERT_TRUE(handle.acquire_pages(&hot_page, 1, &hot_buffer)); + ASSERT_NE(nullptr, hot_buffer); + pool.page_table_.set_evict_priority(hot_page, VecBufferPool::kLowPriority); + ASSERT_TRUE(pool.page_table_.promote_evict_priority( + hot_page, VecBufferPool::kNormalPriority)); + const uint64_t second_chance_before = pool.stats().second_chance; + handle.release_pages(&hot_page, 1); + (void)pool.page_table_.evict_block(hot_page); + ASSERT_TRUE(pool.page_table_.wait_for_block_transition(hot_page)); + EXPECT_GT(pool.stats().second_chance, second_chance_before); +} + +// --------------------------------------------------------------------------- +// 3. The background evictor should proactively reclaim resident-but-released +// pages down to the low watermark (75%) without any foreground eviction. +// --------------------------------------------------------------------------- +TEST_F(BufferPoolTest, BackgroundReclaimsToLowWatermark) { + const size_t cap_pages = 64; + const size_t num_pages = 64; + InitVecPool(cap_pages, /*file_pages=*/num_pages); + std::string file = NewFile(num_pages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + auto handle = pool.get_handle(); + + // Read every page individually so each becomes resident then released, + // filling the pool close to capacity. + std::vector buf(kVectorPageSize); + for (size_t p = 0; p < num_pages; ++p) { + ASSERT_TRUE( + handle.read_range(p * kVectorPageSize, kVectorPageSize, buf.data())); + } + + auto &mp = MemoryLimitPool::get_instance(); + const size_t low = cap_pages * kVectorPageSize / 4 * 3; // 75% page budget + // Poll up to ~2s for the background thread to reclaim down to the low mark. + for (int i = 0; i < 200 && mp.stats().page_used > low + kVectorPageSize; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_LE(mp.stats().page_used, low + kVectorPageSize); + EXPECT_GT(mp.stats().bg_evicted_buffers, 0u); +} + +TEST_F(BufferPoolTest, BackgroundBacksOffWhenAllPagesArePinned) { + constexpr size_t kPageCount = 4; + InitVecPool(/*capacity_pages=*/kPageCount, /*file_pages=*/kPageCount); + std::string file = NewFile(kPageCount); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + std::vector pinned(kPageCount, nullptr); + for (size_t page = 0; page < kPageCount; ++page) { + pinned[page] = pool.acquire_buffer(page); + ASSERT_NE(nullptr, pinned[page]); + } + + auto &memory_pool = MemoryLimitPool::get_instance(); + const uint64_t sleeps_before = memory_pool.stats().bg_no_progress_sleeps; + for (int i = 0; + i < 200 && memory_pool.stats().bg_no_progress_sleeps == sleeps_before; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + EXPECT_GT(memory_pool.stats().bg_no_progress_sleeps, sleeps_before); + + for (size_t page = 0; page < kPageCount; ++page) { + pool.page_table_.release_block(page); + } +} + +// --------------------------------------------------------------------------- +// 4. Concurrent random reads across many threads exercise the sharded +// free-list, +// concurrent acquire/release/evict and the background thread simultaneously. +// All reads must return correct data with no crash or corruption. +// --------------------------------------------------------------------------- +TEST_F(BufferPoolTest, ConcurrentRandomReads) { + const size_t num_pages = 256; + InitVecPool(/*capacity_pages=*/48, /*file_pages=*/num_pages); + std::string file = NewFile(num_pages); + + VecBufferPool pool(file, /*writable=*/false); + ASSERT_EQ(pool.init(), 0); + + const int kThreads = 8; + const int kIters = 3000; + std::atomic failed{false}; + std::vector threads; + for (int t = 0; t < kThreads; ++t) { + threads.emplace_back([&, t]() { + std::mt19937 rng(static_cast(t + 1)); + std::uniform_int_distribution dist(0, num_pages - 1); + auto handle = pool.get_handle(); + std::vector buf(kVectorPageSize); + for (int i = 0; i < kIters && !failed.load(); ++i) { + size_t p = dist(rng); + if (!handle.read_range(p * kVectorPageSize, kVectorPageSize, + buf.data())) { + failed.store(true); + break; + } + char expected = static_cast(p & 0xff); + if (buf[0] != expected || buf[kVectorPageSize - 1] != expected) { + failed.store(true); + break; + } + } + }); + } + for (auto &th : threads) th.join(); + EXPECT_FALSE(failed.load()); +} + +// --------------------------------------------------------------------------- +// 5. Sharded MemoryLimitPool: allocate/free correctness and stats accounting. +// --------------------------------------------------------------------------- +TEST_F(BufferPoolTest, ShardedPoolAllocFreeAccounting) { + const size_t cap_pages = 32; + InitPool(cap_pages); + auto &mp = MemoryLimitPool::get_instance(); + + std::vector bufs; + // Acquire up to capacity. + for (size_t i = 0; i < cap_pages; ++i) { + char *b = nullptr; + ASSERT_TRUE(mp.try_acquire_buffer(kVectorPageSize, b)); + ASSERT_NE(b, nullptr); + bufs.push_back(b); + } + // Pool is full now: further acquire must fail. + char *overflow = nullptr; + EXPECT_FALSE(mp.try_acquire_buffer(kVectorPageSize, overflow)); + EXPECT_EQ(mp.used(), cap_pages * kVectorPageSize); + + // Release everything back to the shards. + for (char *b : bufs) mp.release_buffer(b, kVectorPageSize); + EXPECT_EQ(mp.used(), 0u); + EXPECT_EQ(mp.committed(), cap_pages * kVectorPageSize); + + // Re-acquire should now be served from shard free-lists (no new slab carve). + MemoryLimitPool::PoolStats before = mp.stats(); + char *b = nullptr; + ASSERT_TRUE(mp.try_acquire_buffer(kVectorPageSize, b)); + MemoryLimitPool::PoolStats after = mp.stats(); + EXPECT_GT(after.alloc_from_freelist, before.alloc_from_freelist); + mp.release_buffer(b, kVectorPageSize); +} + +TEST_F(BufferPoolTest, SlabBaseIsFourMiBAlignedAndPagesAreDirectIoAligned) { + constexpr size_t kPages = 8; + InitPool(kPages); + auto &mp = MemoryLimitPool::get_instance(); + + std::vector pages; + uintptr_t slab_base = 0; + for (size_t i = 0; i < kPages; ++i) { + char *page = nullptr; + ASSERT_TRUE(mp.try_acquire_buffer(kVectorPageSize, page)); + ASSERT_NE(nullptr, page); + const uintptr_t address = reinterpret_cast(page); + EXPECT_EQ(0u, address % MemoryLimitPool::page_buffer_size()); + const uintptr_t current_slab = + address & ~(MemoryLimitPool::slab_alignment() - 1); + EXPECT_EQ(0u, current_slab % MemoryLimitPool::slab_alignment()); + EXPECT_GE(address - current_slab, MemoryLimitPool::page_buffer_size()); + if (slab_base == 0) { + slab_base = current_slab; + } else { + EXPECT_EQ(slab_base, current_slab); + } + pages.push_back(page); + } + + const auto allocated = mp.stats(); + EXPECT_EQ(1u, allocated.slab_count); + EXPECT_GE(allocated.slab_mapped_bytes, MemoryLimitPool::slab_size()); + EXPECT_EQ(MemoryLimitPool::page_buffer_size(), allocated.slab_header_bytes); + for (char *page : pages) { + mp.release_buffer(page, kVectorPageSize); + } +} + +TEST_F(BufferPoolTest, ReinitializationReleasesSlabMappings) { + InitPool(/*capacity_pages=*/4); + auto &mp = MemoryLimitPool::get_instance(); + + char *page = nullptr; + ASSERT_TRUE(mp.try_acquire_buffer(kVectorPageSize, page)); + mp.release_buffer(page, kVectorPageSize); + ASSERT_EQ(1u, mp.stats().slab_count); + + ASSERT_EQ(0, mp.init(8 * kVectorPageSize)); + const auto reinitialized = mp.stats(); + EXPECT_EQ(0u, reinitialized.used); + EXPECT_EQ(0u, reinitialized.committed); + EXPECT_EQ(0u, reinitialized.free_buffers); + EXPECT_EQ(0u, reinitialized.slab_count); + EXPECT_EQ(0u, reinitialized.slab_mapped_bytes); + EXPECT_EQ(0u, reinitialized.slab_header_bytes); +} diff --git a/tests/ailego/io/iouring_loader_test.cc b/tests/ailego/io/iouring_loader_test.cc new file mode 100644 index 000000000..162db4561 --- /dev/null +++ b/tests/ailego/io/iouring_loader_test.cc @@ -0,0 +1,126 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace ailego { +namespace { + +constexpr size_t kBlockSize = 512; +constexpr size_t kBlockCount = 4; + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + if (fd_ >= 0) { + ::close(fd_); + } + ::unlink(path_); + } + + int fd() const { + return fd_; + } + + private: + char path_[64] = "IoUringLoaderTest.XXXXXX"; + int fd_; +}; + +void *allocate_aligned(size_t size) { + void *buffer = nullptr; + if (::posix_memalign(&buffer, kBlockSize, size) != 0) { + return nullptr; + } + std::memset(buffer, 0, size); + return buffer; +} + +TEST(IoUringLoaderTest, ReadsBatchAndZeroPadsExpectedShortRead) { + IoUringRing ring; + if (!ring.setup(/*entries=*/8)) { + GTEST_SKIP() << "io_uring is unavailable"; + } + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + constexpr size_t kTailSize = 123; + std::vector source(kBlockSize + kTailSize); + std::fill(source.begin(), source.begin() + kBlockSize, 0x3c); + std::fill(source.begin() + kBlockSize, source.end(), 0x7d); + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + + void *output = allocate_aligned(2 * kBlockSize); + ASSERT_NE(output, nullptr); + std::vector requests; + requests.emplace_back(/*offset=*/0, kBlockSize, output); + requests.emplace_back(kBlockSize, kBlockSize, + static_cast(output) + kBlockSize, kTailSize); + + ASSERT_EQ(ring.execute(file.fd(), requests), 0); + EXPECT_EQ(std::memcmp(output, source.data(), kBlockSize), 0); + const auto *tail = static_cast(output) + kBlockSize; + EXPECT_EQ(std::memcmp(tail, source.data() + kBlockSize, kTailSize), 0); + EXPECT_TRUE(std::all_of(tail + kTailSize, tail + kBlockSize, + [](uint8_t value) { return value == 0; })); + std::free(output); +} + +TEST(IoUringLoaderTest, WritesScatteredBatch) { + IoUringRing ring; + if (!ring.setup(/*entries=*/8)) { + GTEST_SKIP() << "io_uring is unavailable"; + } + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kBlockCount * kBlockSize); + for (size_t block = 0; block < kBlockCount; ++block) { + std::memset(source.data() + block * kBlockSize, static_cast(block + 1), + kBlockSize); + } + + const std::array order = {3, 0, 2, 1}; + std::vector requests; + for (size_t block : order) { + requests.emplace_back(block * kBlockSize, kBlockSize, + source.data() + block * kBlockSize); + } + ASSERT_EQ(ring.execute_writes(file.fd(), requests), 0); + + std::vector output(source.size()); + ASSERT_EQ(::pread(file.fd(), output.data(), output.size(), 0), + static_cast(output.size())); + EXPECT_EQ(output, source); +} + +} // namespace +} // namespace ailego +} // namespace zvec + +#endif // __linux__