diff --git a/src/ailego/buffer/block_eviction_queue.cc b/src/ailego/buffer/block_eviction_queue.cc index eff930127..40d97d0c3 100644 --- a/src/ailego/buffer/block_eviction_queue.cc +++ b/src/ailego/buffer/block_eviction_queue.cc @@ -12,127 +12,1003 @@ // 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 3318db1ca..8812d440e 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 @@ -20,7 +24,13 @@ #include #include +#if defined(__linux__) +#include +#include +#endif + #if defined(_MSC_VER) +#include #ifndef NOMINMAX #define NOMINMAX #endif @@ -67,100 +77,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; } - // 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; + 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; + } + + 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; } - size_t old_count = segment_count_.load(std::memory_order_relaxed); + 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; + } + + 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 (e.ref_count.compare_exchange_weak(current_count, current_count + 1, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - return e.buffer; + 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)); + } + + 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, @@ -169,105 +625,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_); + } + } + } +} + +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; } -void VectorPageTable::evict_block(block_id_t block_id) { +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; } } } @@ -275,65 +914,277 @@ 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); fd_ = _open(filename.c_str(), flags, 0644); + meta_fd_ = _open(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 +} + +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; + } + // 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; + } +#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; } int VecBufferPool::init() { - size_t block_num = (file_size_ + kVectorPageSize - 1) / kVectorPageSize; + 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; + } + } + + 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; + } + // 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()); + } + } +#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 (exceeds " - "VectorPageTable::kMaxEntries=%zu)", + "file_size=%zu, block_num=%zu, required_metadata=%zu", file_name_.c_str(), file_size_, block_num, - VectorPageTable::kMaxEntries); + metadata_bytes_for_page_count(block_num, writable_)); return -1; } - block_mutexes_ = - std::make_unique(VecBufferPool::kMutexBucketCount); + 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_); - // 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. + initialized_ = true; 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; - }); + 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; } @@ -342,22 +1193,406 @@ VecBufferPoolHandle VecBufferPool::get_handle() { return VecBufferPoolHandle(*this); } -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; +bool VecBufferPool::enqueue_writeback(block_id_t page_id) { + if (writeback_error() != 0) { + return false; } - std::lock_guard lock( - block_mutexes_[page_id % VecBufferPool::kMutexBucketCount]); - buffer = page_table_.acquire_block(page_id); - if (buffer) { - return buffer; + 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( @@ -368,31 +1603,175 @@ 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); } +} + +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; - 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); + 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; + } + } } - 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; + 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; } - return page_table_.set_block_acquired(page_id, buffer, page_offset); + 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], " @@ -400,7 +1779,65 @@ int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) { file_name_.c_str(), offset, length); return -1; } - return 0; + 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, @@ -413,20 +1850,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); @@ -439,6 +1884,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_) { @@ -446,7 +1944,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], " @@ -461,27 +1966,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; } @@ -495,11 +2020,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, " @@ -508,44 +2038,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); @@ -554,36 +2091,209 @@ 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); +} + +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) { + page = pool_.acquire_buffer(static_cast(run_start + j), 50); + 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); + 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; + if (page_end_in_buf <= actually_read && + !pool_.page_table_.is_loaded(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); } @@ -593,6 +2303,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); @@ -611,11 +2326,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_context.cc b/src/core/algorithm/diskann/diskann_context.cc index f13affb74..6ce2712a4 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -27,7 +27,8 @@ DiskAnnContext::DiskAnnContext(const IndexMeta &meta, : dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} int DiskAnnContext::init(ContextType type, uint32_t graph_degree, - uint32_t pq_chunk_num, uint32_t element_size) { + uint32_t pq_chunk_num, uint32_t element_size, + bool setup_io_context) { type_ = type; element_size_ = element_size; pq_chunk_num_ = pq_chunk_num; @@ -66,10 +67,12 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, DiskAnnUtil::kMaxSectorReadNum * DiskAnnUtil::kSectorSize, DiskAnnUtil::kSectorSize); - ret = setup_io_ctx(io_ctx_); - if (ret != 0) { - LOG_ERROR("setup io ctx error, ret=%d", ret); - return ret; + if (setup_io_context) { + ret = setup_io_ctx(io_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error, ret=%d", ret); + return ret; + } } break; @@ -139,4 +142,4 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, } } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index ce8c0ca0c..532e090bd 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -60,7 +60,7 @@ class DiskAnnContext : public IndexContext, public: //! Init int init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, - uint32_t element_size); + uint32_t element_size, bool setup_io_context = true); //! Update context, the context may be shared by different searcher/streamer int update_context(ContextType type, const IndexMeta &meta, diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 3c6301121..2c81eb248 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -16,12 +16,14 @@ #include #include #include -#include #include #include #include +#include +#include #include #include +#include #include #include @@ -34,10 +36,6 @@ namespace core { 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; - // Ensures the I/O backend selection is logged exactly once per process, // regardless of which entry point (setup_io_ctx or register_thread) // triggers it first. @@ -46,23 +44,25 @@ static std::once_flag g_io_backend_log_once; void log_diskann_io_backend() { #if (defined(__linux) || defined(__linux__)) - auto &backend = ailego::IOBackend::Instance(); - if (backend.is_pread()) { - LOG_WARN( - "DiskAnn: no async I/O backend available. Install libaio (e.g. " - "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " - "retry. DiskAnn will fall back to synchronous pread() — performance " - "will be degraded."); - } else { - LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", - backend.name()); - } + std::call_once(g_io_backend_log_once, [] { + auto &backend = ailego::IOBackend::Instance(); + if (backend.is_pread()) { + LOG_WARN( + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); + } else { + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); + } + }); #endif } int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - std::call_once(g_io_backend_log_once, log_diskann_io_backend); + log_diskann_io_backend(); if (ailego::IOBackend::Instance().is_pread()) { // No async backend available — leave ctx null so callers fall back to // synchronous pread(). @@ -324,218 +324,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; } @@ -580,7 +368,8 @@ void LinuxAlignedFileReader::register_thread() { return; } if (ctx != nullptr) { - LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); + LOG_INFO("allocating ctx: %llu", + static_cast(reinterpret_cast(ctx))); } ctx_map[thread_id] = ctx; lk.unlock(); @@ -656,11 +445,7 @@ void LinuxAlignedFileReader::close() { } int LinuxAlignedFileReader::read(std::vector &read_reqs, - IOContext &ctx, bool async) { - if (async == true) { - LOG_WARN("Async currently not supported"); - } - + IOContext &ctx) { if (this->file_desc == -1) { LOG_ERROR("Attempt to read from invalid file descriptor"); return IndexError_Runtime; @@ -671,6 +456,250 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } +BufferPoolAlignedFileReader::BufferPoolAlignedFileReader( + std::shared_ptr pool) + : pool_(std::move(pool)) {} + +BufferPoolAlignedFileReader::~BufferPoolAlignedFileReader() = default; + +IOContext &BufferPoolAlignedFileReader::get_ctx() { + return unused_ctx_; +} + +void BufferPoolAlignedFileReader::register_thread() {} + +void BufferPoolAlignedFileReader::deregister_thread() {} + +void BufferPoolAlignedFileReader::deregister_all_threads() {} + +void BufferPoolAlignedFileReader::open(const std::string &fname) { + bypass_reader_.open(fname); +} + +void BufferPoolAlignedFileReader::close() { + bypass_reader_.close(); + pool_.reset(); +} + +int BufferPoolAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx) { + if (!pool_) { + LOG_ERROR("BufferPoolAlignedFileReader: buffer pool is not available"); + return IndexError_Runtime; + } + if (read_reqs.empty()) return 0; + + try { + struct UniquePage { + ailego::block_id_t page_id; + char *first_destination; + char *cached_page{nullptr}; + bool bypass_candidate{false}; + }; + struct PageOccurrence { + size_t unique_index; + char *destination; + }; + + size_t total_pages = 0; + for (const AlignedRead &req : read_reqs) { + if (req.buf == nullptr || req.len == 0 || + req.offset > std::numeric_limits::max() || + req.len > std::numeric_limits::max()) { + return IndexError_InvalidArgument; + } + const size_t offset = static_cast(req.offset); + const size_t length = static_cast(req.len); + if (offset % ailego::kVectorPageSize != 0 || + length % ailego::kVectorPageSize != 0 || + offset > pool_->file_size() || length > pool_->file_size() - offset) { + return IndexError_InvalidArgument; + } + const size_t pages = length / ailego::kVectorPageSize; + if (pages > std::numeric_limits::max() - total_pages) { + return IndexError_InvalidLength; + } + total_pages += pages; + } + + std::vector unique_pages; + std::vector occurrences; + unique_pages.reserve(total_pages); + occurrences.reserve(total_pages); + for (const AlignedRead &req : read_reqs) { + const size_t first_page = + static_cast(req.offset) / ailego::kVectorPageSize; + const size_t pages = + static_cast(req.len) / ailego::kVectorPageSize; + char *destination = static_cast(req.buf); + for (size_t i = 0; i < pages; ++i) { + const auto page_id = static_cast(first_page + i); + size_t unique_index = 0; + while (unique_index < unique_pages.size() && + unique_pages[unique_index].page_id != page_id) { + ++unique_index; + } + char *page_destination = destination + i * ailego::kVectorPageSize; + if (unique_index == unique_pages.size()) { + unique_pages.push_back( + UniquePage{page_id, page_destination, nullptr, false}); + } + occurrences.push_back(PageOccurrence{unique_index, page_destination}); + } + } + + std::vector admitted_ids; + std::vector admitted_indices; + std::vector admitted_pages(unique_pages.size(), nullptr); + std::vector bypass_requests; + admitted_ids.reserve(unique_pages.size()); + admitted_indices.reserve(unique_pages.size()); + bypass_requests.reserve(unique_pages.size()); + + auto release_cached_pages = [&]() { + for (UniquePage &page : unique_pages) { + if (page.cached_page != nullptr) { + pool_->release_pages(&page.page_id, 1); + page.cached_page = nullptr; + } + } + }; + struct CachedPageGuard { + decltype(release_cached_pages) &release; + ~CachedPageGuard() { + release(); + } + } cached_page_guard{release_cached_pages}; + + for (size_t i = 0; i < unique_pages.size(); ++i) { + UniquePage &page = unique_pages[i]; + page.cached_page = pool_->try_acquire_buffer(page.page_id); + if (page.cached_page != nullptr) { + continue; + } + if (pool_->should_admit_page(page.page_id)) { + admitted_ids.push_back(page.page_id); + admitted_indices.push_back(i); + } else { + page.bypass_candidate = true; + } + } + + if (!admitted_ids.empty() && + !pool_->acquire_pages(admitted_ids.data(), admitted_ids.size(), + admitted_pages.data())) { + release_cached_pages(); + return IndexError_ReadData; + } + for (size_t i = 0; i < admitted_ids.size(); ++i) { + unique_pages[admitted_indices[i]].cached_page = admitted_pages[i]; + } + + // A rejected page may have become resident or started loading while the + // admitted portion of this batch was populated. Rejoin that cache flight + // instead of issuing duplicate direct I/O. This does not record another + // admission observation. + size_t bypass_rechecks = 0; + size_t bypass_cache_joins = 0; + admitted_ids.clear(); + admitted_indices.clear(); + for (size_t i = 0; i < unique_pages.size(); ++i) { + UniquePage &page = unique_pages[i]; + if (!page.bypass_candidate) { + continue; + } + ++bypass_rechecks; + page.cached_page = pool_->try_acquire_buffer(page.page_id); + if (page.cached_page != nullptr) { + page.bypass_candidate = false; + ++bypass_cache_joins; + } else if (pool_->should_join_cache_path(page.page_id)) { + admitted_ids.push_back(page.page_id); + admitted_indices.push_back(i); + } + } + if (!admitted_ids.empty() && + pool_->acquire_pages(admitted_ids.data(), admitted_ids.size(), + admitted_pages.data())) { + for (size_t i = 0; i < admitted_ids.size(); ++i) { + UniquePage &page = unique_pages[admitted_indices[i]]; + page.cached_page = admitted_pages[i]; + page.bypass_candidate = false; + } + bypass_cache_joins += admitted_ids.size(); + } + pool_->record_bypass_recheck(bypass_rechecks, bypass_cache_joins); + + // Preserve large contiguous DiskANN reads on the direct path. Duplicate + // pages use their first destination as the canonical read target and are + // fanned out after I/O completes. + uint64_t run_offset = 0; + uint64_t run_length = 0; + char *run_destination = nullptr; + size_t bypass_page_count = 0; + auto flush_bypass_run = [&]() { + if (run_length != 0) { + bypass_requests.emplace_back(run_offset, run_length, run_destination); + run_length = 0; + } + }; + for (const PageOccurrence &occurrence : occurrences) { + const UniquePage &page = unique_pages[occurrence.unique_index]; + if (!page.bypass_candidate || + occurrence.destination != page.first_destination) { + continue; + } + const uint64_t page_offset = + static_cast(page.page_id) * ailego::kVectorPageSize; + if (run_length != 0 && run_offset + run_length == page_offset && + run_destination + run_length == occurrence.destination) { + run_length += ailego::kVectorPageSize; + } else { + flush_bypass_run(); + run_offset = page_offset; + run_length = ailego::kVectorPageSize; + run_destination = occurrence.destination; + } + ++bypass_page_count; + } + flush_bypass_run(); + + if (!bypass_requests.empty()) { +#if defined(__linux__) || defined(__linux) + // Buffer-pool hits need no DiskANN I/O context. Create it only when + // admission first chooses direct AIO; the caller already owns and + // destroys this context with its normal DiskANN context lifecycle. + if (ctx == nullptr && setup_io_ctx(ctx) != 0) { + release_cached_pages(); + return IndexError_Runtime; + } +#endif + const int read_ret = bypass_reader_.read(bypass_requests, ctx); + if (read_ret != 0) { + release_cached_pages(); + return read_ret; + } + pool_->record_bypass_read(bypass_page_count * ailego::kVectorPageSize, + bypass_requests.size()); + } + + for (const PageOccurrence &occurrence : occurrences) { + const UniquePage &page = unique_pages[occurrence.unique_index]; + if (page.cached_page != nullptr) { + std::memcpy(occurrence.destination, page.cached_page, + ailego::kVectorPageSize); + } else if (occurrence.destination != page.first_destination) { + std::memcpy(occurrence.destination, page.first_destination, + ailego::kVectorPageSize); + } + } + release_cached_pages(); + return 0; + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } +} + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 7eaf41e8b..37e5fe36e 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -24,11 +24,15 @@ #include #include +#include #include #include #include "diskann_util.h" namespace zvec { +namespace ailego { +class VecBufferPool; +} namespace core { #if (defined(__linux) || defined(__linux__)) @@ -51,7 +55,7 @@ struct IoBackend { }; Backend backend{NONE}; - IoUringRing ring{}; + ailego::IoUringRing ring{}; io_context_t aio_ctx{nullptr}; }; @@ -68,6 +72,9 @@ int destroy_io_ctx(IOContext &ctx); // Probes the backend on first call. No-op on non-Linux platforms. void log_diskann_io_backend(); +#if (defined(__linux) || defined(__linux__)) +using AlignedRead = ailego::IoUringRead; +#else struct AlignedRead { uint64_t offset; uint64_t len; @@ -82,6 +89,7 @@ struct AlignedRead { ailego_assert(reinterpret_cast(buf) % 512 == 0); } }; +#endif class AlignedFileReader { protected: @@ -100,8 +108,12 @@ class AlignedFileReader { virtual void open(const std::string &fname) = 0; virtual void close() = 0; - virtual int read(std::vector &read_reqs, IOContext &ctx, - bool async = false) = 0; + //! Submit the batch and return only after every destination is ready. + virtual int read(std::vector &read_reqs, IOContext &ctx) = 0; + + virtual bool requires_io_context() const { + return true; + } }; class LinuxAlignedFileReader : public AlignedFileReader { @@ -124,8 +136,31 @@ class LinuxAlignedFileReader : public AlignedFileReader { void open(const std::string &fname); void close(); - int read(std::vector &read_reqs, IOContext &ctx, - bool async = false); + int read(std::vector &read_reqs, IOContext &ctx); +}; + +class BufferPoolAlignedFileReader : public AlignedFileReader { + public: + explicit BufferPoolAlignedFileReader( + std::shared_ptr pool); + ~BufferPoolAlignedFileReader() override; + + IOContext &get_ctx() override; + void register_thread() override; + void deregister_thread() override; + void deregister_all_threads() override; + void open(const std::string &fname) override; + void close() override; + int read(std::vector &read_reqs, IOContext &ctx) override; + + bool requires_io_context() const override { + return false; + } + + private: + std::shared_ptr pool_; + LinuxAlignedFileReader bypass_reader_; + IOContext unused_ctx_{}; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..b27766cb3 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -38,24 +38,40 @@ DiskAnnIndexer::~DiskAnnIndexer() { int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { entity_ = &entity; - auto storage = entity.get_storage(); + storage_ = entity.get_storage(); + auto storage = storage_; auto vector_segment = entity.get_vector_segment(); pq_table_ = entity.get_pq_table(); index_segment_offset_ = vector_segment->data_offset(); - reader_.reset(new LinuxAlignedFileReader()); - - auto file_path = storage->file_path(); - reader_->open(file_path); - - storage->cleanup(); + auto pool = storage->vec_buffer_pool(); + if (pool) { + if (ailego::kVectorPageSize != DiskAnnUtil::kSectorSize) { + LOG_ERROR( + "DiskAnn BufferPool page size mismatch: page_size=%zu " + "sector_size=%zu", + ailego::kVectorPageSize, + static_cast(DiskAnnUtil::kSectorSize)); + return IndexError_Unsupported; + } + reader_ = std::make_shared(std::move(pool)); + reader_->open(storage->file_path()); + } else { + reader_ = std::make_shared(); + reader_->open(storage->file_path()); + storage->cleanup(); + storage_.reset(); + } - int ret = setup_io_ctx(init_ctx_); - if (ret != 0) { - LOG_ERROR("setup io ctx error"); - return ret; + int ret = 0; + if (reader_->requires_io_context()) { + ret = setup_io_ctx(init_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error"); + return ret; + } } max_node_size_ = entity.max_node_size(); diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index c372d288f..07fc49f30 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -51,6 +51,10 @@ class DiskAnnIndexer { int get_vector(diskann_id_t id, IndexContext::Pointer &context, std::string &vector); + bool requires_io_context() const { + return reader_ && reader_->requires_io_context(); + } + diskann_key_t get_key(diskann_id_t id) const; diskann_id_t get_id(diskann_key_t key) const; @@ -87,7 +91,7 @@ class DiskAnnIndexer { diskann_id_t medoid_; std::vector entrypoints_; - std::shared_ptr reader_{nullptr}; + std::shared_ptr reader_{nullptr}; PQTable::Pointer pq_table_; diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a34c546e5..245f15f45 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -295,7 +295,8 @@ IndexSearcher::Context::Pointer DiskAnnSearcher::create_context() const { } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), meta_.element_size(), + diskann_indexer_->requires_io_context())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h index 99584fa35..903259c86 100644 --- a/src/core/algorithm/diskann/diskann_searcher.h +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -17,8 +17,6 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..c51b42be0 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -330,7 +330,8 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), meta_.element_size(), + diskann_indexer_->requires_io_context())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; @@ -338,6 +339,7 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } ctx->set_list_size(list_size_); + ctx->set_magic(magic_); return Context::Pointer(ctx); } diff --git a/src/core/algorithm/diskann/diskann_streamer.h b/src/core/algorithm/diskann/diskann_streamer.h index ddb159ae7..4ab7e012a 100644 --- a/src/core/algorithm/diskann/diskann_streamer.h +++ b/src/core/algorithm/diskann/diskann_streamer.h @@ -17,8 +17,6 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc index 1fd6f6567..d8ca45295 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.cc +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "hnsw_algorithm.h" -#include namespace zvec { namespace core { @@ -160,12 +159,12 @@ void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, HnswDistCalculator &dc = ctx->dist_calculator(); - update_neighbors(dc, id, level, topk_heap); + update_neighbors(dc, id, level, topk_heap, ctx); // reverse update neighbors for (size_t i = 0; i < topk_heap.size(); ++i) { reverse_update_neighbors(dc, topk_heap[i].first, level, id, - topk_heap[i].second, ctx->update_heap()); + topk_heap[i].second, ctx->update_heap(), ctx); } return; @@ -177,11 +176,13 @@ void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, // Two specialized inner loops, dispatched from search_neighbors(): // // fast_search_neighbors: mmap/contiguous with direct vector pointers. -// Uses BlockHeap (AVX2) or LinearPool (scalar) -// for visited tracking and top-k maintenance. +// fast_search_neighbors_buffer: BufferStorage with page-backed MemoryBlocks. +// Both use BlockHeap (AVX2) or LinearPool +// (scalar) for visited tracking and top-k +// maintenance. // dual_heap_search_neighbors: CandidateHeap + TopkHeap + VisitFilter. // Used for add_node (use_pool=false), filtered -// search, upper levels, and BufferPool fallback. +// search and upper levels for every backend. // ============================================================================ // mmap/contiguous variant: resolve vectors via get_vector_ptr and use @@ -257,6 +258,80 @@ void fast_search_neighbors(const EntityType &entity, HeapType &pool, } } +// BufferStorage variant of the level-0 fast path. It intentionally keeps the +// MemoryBlocks alive through batch_dist(): a buffer-pool page may be evicted as +// soon as its last block is released. Apart from vector resolution, this is +// the same graph traversal used by mmap, so selecting BufferStorage does not +// silently switch HNSW to the slower dual-heap algorithm. +template +void fast_search_neighbors_buffer(const EntityType &entity, HeapType &pool, + VisitFilter &visit, HnswDistCalculator &dc, + uint32_t topk, uint32_t ef, + node_id_t entry_point, dist_t entry_dist, + uint32_t prefetch_lines, + uint32_t prefetch_offset) { + using MemBlockType = typename EntityType::MemoryBlock; + + const uint32_t max_deg = entity.max_degree(0); + const uint32_t cap = std::max(topk, ef); + pool.reset(static_cast(cap), static_cast(max_deg)); + visit.clear(); + + visit.set_visited(entry_point); + pool.push_block(&entry_dist, &entry_point, 1); + + uint32_t buf_capacity = max_deg; + std::vector neighbor_ids(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + std::vector neighbor_vec_blocks; + neighbor_vec_blocks.reserve(buf_capacity); + + while (pool.has_next()) { + const auto current_node = pool.pop(); + const auto neighbors = entity.get_neighbors_typed(0, current_node); + ailego_prefetch(neighbors.data); + + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + neighbor_vec_blocks.reserve(buf_capacity); + } + + uint32_t unvisited_count = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + const node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + neighbor_ids[unvisited_count++] = node; + } + if (unvisited_count == 0) continue; + + neighbor_vec_blocks.clear(); + if (ailego_unlikely(entity.get_vector_typed(neighbor_ids.data(), + unvisited_count, + neighbor_vec_blocks) != 0)) { + break; + } + for (uint32_t i = 0; i < unvisited_count; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + const uint32_t po = std::min(prefetch_offset, unvisited_count); + for (uint32_t i = 0; i < po; ++i) { + const char *p = static_cast(neighbor_vecs[i]); + for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { + ailego_prefetch(p + cl * 64); + } + } + + dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + pool.push_block(dists.data(), neighbor_ids.data(), + static_cast(unvisited_count)); + } +} + // ============================================================================ // dual_heap_search_neighbors: shared core for the fallback dual-heap path. // @@ -380,8 +455,8 @@ void dual_heap_search_neighbors(const EntityType &entity, level_t level, // // - add_node / filtered / upper levels → dual_heap_search_neighbors // - level-0 unfiltered search: -// MmapMemoryBlock → fast_search_neighbors (BlockHeap/LinearPool) -// BufferPool → dual_heap_search_neighbors (fallback) +// MmapMemoryBlock → fast_search_neighbors (direct pointers) +// BufferPool → fast_search_neighbors_buffer (pinned pages) // ============================================================================ template void HnswAlgorithm::search_neighbors(level_t level, @@ -436,10 +511,27 @@ void HnswAlgorithm::search_neighbors(level_t level, copy_pool_to_topk(lpool, topk); } } else { - // BufferPool entities: fallback to dual-heap path. - auto filter = [](node_id_t) { return false; }; - dual_heap_search_neighbors( - entity, level, entry_point, dist, topk, ctx, dc, filter); + const uint32_t prefetch_lines = + ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; + const uint32_t topk_v = static_cast(ctx->topk()); + const uint32_t ef_v = ctx->ef(); + const bool avx2_ok = + zvec::ailego::internal::CpuFeatures::static_flags_.AVX2; + auto &visit = ctx->visit_filter(); + + if (avx2_ok) { + auto &bpool = ctx->block_pool(); + fast_search_neighbors_buffer(entity, bpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, + ctx->po()); + copy_pool_to_topk(bpool, topk); + } else { + auto &lpool = ctx->pool(); + fast_search_neighbors_buffer(entity, lpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, + ctx->po()); + copy_pool_to_topk(lpool, topk); + } } } } @@ -567,7 +659,8 @@ void HnswAlgorithm::expand_neighbors_by_group( template void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, node_id_t id, level_t level, - TopkHeap &topk_heap) { + TopkHeap &topk_heap, + HnswContext *ctx) { topk_heap.sort(); uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level); @@ -579,24 +672,72 @@ void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, } uint32_t cur_size = 0; - for (size_t i = 0; i < topk_heap.size(); ++i) { - node_id_t cur_node = topk_heap[i].first; - dist_t cur_node_dist = topk_heap[i].second; - bool good = true; - for (uint32_t j = 0; j < cur_size; ++j) { - dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first); - if (tmp_dist <= cur_node_dist) { - good = false; - break; + if constexpr (std::is_same_v) { + const auto &read_entity = + static_cast(ctx->get_entity()); + auto &prune_ids = ctx->prune_ids(); + auto &prune_blocks = ctx->prune_blocks(); + auto &selected_indices = ctx->prune_selected_indices(); + + prune_ids.clear(); + prune_ids.reserve(topk_heap.size()); + for (const auto &candidate : topk_heap) { + prune_ids.emplace_back(candidate.first); + } + prune_blocks.clear(); + if (ailego_unlikely(read_entity.get_vector_for_prune( + prune_ids.data(), prune_ids.size(), + prune_blocks) != 0)) { + dc.set_error(); + return; + } + selected_indices.clear(); + selected_indices.reserve(max_neighbor_cnt); + + for (size_t i = 0; i < topk_heap.size(); ++i) { + node_id_t cur_node = topk_heap[i].first; + dist_t cur_node_dist = topk_heap[i].second; + bool good = true; + for (const size_t selected : selected_indices) { + const dist_t pair_distance = + dc.dist(prune_blocks[i].data(), prune_blocks[selected].data()); + if (pair_distance <= cur_node_dist) { + good = false; + break; + } + } + + if (good) { + topk_heap[cur_size].first = cur_node; + topk_heap[cur_size].second = cur_node_dist; + selected_indices.emplace_back(i); + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } } } + prune_blocks.clear(); + } else { + for (size_t i = 0; i < topk_heap.size(); ++i) { + node_id_t cur_node = topk_heap[i].first; + dist_t cur_node_dist = topk_heap[i].second; + bool good = true; + for (uint32_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, topk_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } - if (good) { - topk_heap[cur_size].first = cur_node; - topk_heap[cur_size].second = cur_node_dist; - cur_size++; - if (cur_size >= max_neighbor_cnt) { - break; + if (good) { + topk_heap[cur_size].first = cur_node; + topk_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } } } } @@ -631,60 +772,132 @@ void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, template void HnswAlgorithm::reverse_update_neighbors( HnswDistCalculator &dc, node_id_t id, level_t level, node_id_t link_id, - dist_t dist, TopkHeap &update_heap) { + dist_t dist, TopkHeap &update_heap, HnswContext *ctx) { const size_t max_neighbor_cnt = entity_.neighbor_cnt(level); - uint32_t lock_idx = id & kLockMask; - lock_pool_[lock_idx].lock(); + const uint32_t lock_idx = id & kLockMask; + std::unique_lock node_lock(lock_pool_[lock_idx]); const Neighbors neighbors = entity_.get_neighbors(level, id); - size_t size = neighbors.size(); + const size_t size = neighbors.size(); ailego_assert_with(size <= max_neighbor_cnt, "invalid neighbor size"); if (size < max_neighbor_cnt) { entity_.add_neighbor(level, id, size, link_id); - lock_pool_[lock_idx].unlock(); return; } - update_heap.emplace(link_id, dist); + if constexpr (std::is_same_v) { + const auto &read_entity = + static_cast(ctx->get_entity()); + // A wide level-0 list can require O(M^2) pairwise comparisons. Resolving + // each pair through HnswDistCalculator::dist(node_id, node_id) repeatedly + // pins and releases the same BufferStorage pages and dominates high-M + // construction. Resolve the center and every candidate once, keep those + // blocks alive for the prune, and preserve the existing scalar pruning + // order exactly. + auto &candidate_ids = ctx->prune_ids(); + auto &candidate_blocks = ctx->prune_blocks(); + + candidate_ids.clear(); + candidate_ids.reserve(size + 2); + candidate_ids.emplace_back(id); + candidate_ids.emplace_back(link_id); + for (size_t i = 0; i < size; ++i) { + candidate_ids.emplace_back(neighbors[i]); + } + + candidate_blocks.clear(); + if (ailego_unlikely(read_entity.get_vector_for_prune( + candidate_ids.data(), candidate_ids.size(), + candidate_blocks) != 0)) { + dc.set_error(); + return; + } - for (size_t i = 0; i < size; ++i) { - node_id_t node = neighbors[i]; - dist_t cur_dist = dc.dist(id, node); - update_heap.emplace(node, cur_dist); - } + update_heap.clear(); + // Use block indices until pruning is complete. This keeps sorted heap + // entries directly associated with their pinned vectors without a hash + // map. + update_heap.emplace(1, dist); + for (size_t i = 0; i < size; ++i) { + const dist_t candidate_distance = + dc.dist(candidate_blocks[0].data(), candidate_blocks[i + 2].data()); + update_heap.emplace(static_cast(i + 2), candidate_distance); + } + update_heap.sort(); + + size_t selected_count = 0; + for (size_t i = 0; i < update_heap.size(); ++i) { + const node_id_t candidate_block_index = update_heap[i].first; + const dist_t candidate_distance = update_heap[i].second; + bool good = true; + if (selected_count != 0) { + for (size_t j = 0; j < selected_count; ++j) { + const dist_t pair_distance = dc.dist( + candidate_blocks[candidate_block_index].data(), + candidate_blocks[update_heap[j].first].data()); + if (pair_distance <= candidate_distance) { + good = false; + break; + } + } + } - //! TODO: optimize prune - //! prune edges - update_heap.sort(); - size_t cur_size = 0; - for (size_t i = 0; i < update_heap.size(); ++i) { - node_id_t cur_node = update_heap[i].first; - dist_t cur_node_dist = update_heap[i].second; - bool good = true; - for (size_t j = 0; j < cur_size; ++j) { - dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first); - if (tmp_dist <= cur_node_dist) { - good = false; - break; + if (good) { + update_heap[selected_count] = {candidate_block_index, + candidate_distance}; + ++selected_count; + if (selected_count >= max_neighbor_cnt) { + break; + } } } - if (good) { - update_heap[cur_size].first = cur_node; - update_heap[cur_size].second = cur_node_dist; - cur_size++; - if (cur_size >= max_neighbor_cnt) { - break; - } + update_heap.resize(selected_count); + for (auto &selected : update_heap) { + selected.first = candidate_ids[selected.first]; } - } + entity_.update_neighbors(level, id, update_heap); - update_heap.resize(cur_size); - entity_.update_neighbors(level, id, update_heap); + node_lock.unlock(); + update_heap.clear(); + // Release BufferStorage pins before this worker consumes its next node. + candidate_blocks.clear(); + } else { + update_heap.emplace(link_id, dist); + for (size_t i = 0; i < size; ++i) { + node_id_t node = neighbors[i]; + dist_t cur_dist = dc.dist(id, node); + update_heap.emplace(node, cur_dist); + } + + update_heap.sort(); + size_t cur_size = 0; + for (size_t i = 0; i < update_heap.size(); ++i) { + node_id_t cur_node = update_heap[i].first; + dist_t cur_node_dist = update_heap[i].second; + bool good = true; + for (size_t j = 0; j < cur_size; ++j) { + dist_t tmp_dist = dc.dist(cur_node, update_heap[j].first); + if (tmp_dist <= cur_node_dist) { + good = false; + break; + } + } - lock_pool_[lock_idx].unlock(); + if (good) { + update_heap[cur_size].first = cur_node; + update_heap[cur_size].second = cur_node_dist; + cur_size++; + if (cur_size >= max_neighbor_cnt) { + break; + } + } + } - update_heap.clear(); + update_heap.resize(cur_size); + entity_.update_neighbors(level, id, update_heap); + update_heap.clear(); + } return; } diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.h b/src/core/algorithm/hnsw/hnsw_algorithm.h index 7851b1c52..6d932d96f 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.h +++ b/src/core/algorithm/hnsw/hnsw_algorithm.h @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -113,23 +114,23 @@ class HnswAlgorithm : public HnswAlgorithmBase { //! Given a node id and level, search the nearest neighbors in graph. //! Dispatches to fast_search_neighbors (pool-based, direct pointer) for - //! mmap/contiguous level-0 unfiltered search, or dual_heap_search_neighbors - //! (CandidateHeap + TopkHeap) for add_node, filtered search, upper levels, - //! and BufferPool fallback. + //! mmap/contiguous level-0 unfiltered search, a page-pinned equivalent for + //! BufferStorage, or dual_heap_search_neighbors (CandidateHeap + TopkHeap) + //! for add_node, filtered search and upper levels. //! Note: entry_point and dist will be updated to current level nearest node. void search_neighbors(level_t level, node_id_t *entry_point, dist_t *dist, TopkHeap &topk, HnswContext *ctx, bool use_pool) const; //! Update the node's neighbors void update_neighbors(HnswDistCalculator &dc, node_id_t id, level_t level, - TopkHeap &topk_heap); + TopkHeap &topk_heap, HnswContext *ctx); //! Checking linkId could be id's new neighbor, and add as neighbor if true //! @dc distance calculator //! @updateHeap temporary heap in updating neighbors void reverse_update_neighbors(HnswDistCalculator &dc, node_id_t id, level_t level, node_id_t link_id, dist_t dist, - TopkHeap &update_heap); + TopkHeap &update_heap, HnswContext *ctx); //! expand neighbors until group nums are reached void expand_neighbors_by_group(TopkHeap &topk, HnswContext *ctx) const; @@ -139,7 +140,12 @@ class HnswAlgorithm : public HnswAlgorithmBase { HnswAlgorithm &operator=(const HnswAlgorithm &) = delete; private: - static constexpr uint32_t kLockCnt{1U << 8}; + // A full reverse-neighbor update performs diversity pruning while holding + // the node stripe. BufferStorage uses more stripes so unrelated high-M + // updates do not serialize; leave every other backend's scheduling intact. + static constexpr uint32_t kLockCnt{ + std::is_same_v ? 1U << 12 + : 1U << 8}; static constexpr uint32_t kLockMask{kLockCnt - 1U}; EntityType &entity_; @@ -153,4 +159,4 @@ class HnswAlgorithm : public HnswAlgorithmBase { }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index 4e2f90c5f..aea7adeed 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -299,6 +299,21 @@ class HnswContext : public IndexContext { return update_heap_; } + //! Reusable construction scratch. Keeping this on the per-worker context + //! avoids function-local thread_local ownership and lets the algorithm pin + //! all candidates once for a full-neighbor diversity prune. + inline std::vector &prune_ids() { + return prune_ids_; + } + + inline std::vector &prune_blocks() { + return prune_blocks_; + } + + inline std::vector &prune_selected_indices() { + return prune_selected_indices_; + } + inline LinearPool &pool() { return pool_; } @@ -569,6 +584,9 @@ class HnswContext : public IndexContext { std::vector group_results_{}; TopkHeap topk_heap_{}; TopkHeap update_heap_{}; + std::vector prune_ids_{}; + std::vector prune_blocks_{}; + std::vector prune_selected_indices_{}; std::vector level_topks_{}; CandidateHeap candidates_{}; VisitFilter visit_filter_{}; diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index 2e4b22d1f..182dce27d 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -116,7 +116,7 @@ class HnswDistCalculator { inline dist_t dist(node_id_t id) { compare_cnt_++; IndexStorage::MemoryBlock vec_block; - int ret = entity_->get_vector(id, vec_block); + int ret = entity_->get_vector_borrowed(id, vec_block); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Get nullptr vector, id=%u", id); error_ = true; @@ -138,7 +138,7 @@ class HnswDistCalculator { IndexStorage::MemoryBlock vec_block_feat; - int ret = entity_->get_vector(lhs, vec_block_feat); + int ret = entity_->get_vector_borrowed(lhs, vec_block_feat); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Get nullptr vector, id=%u", lhs); error_ = true; @@ -147,7 +147,7 @@ class HnswDistCalculator { const void *feat = vec_block_feat.data(); IndexStorage::MemoryBlock vec_block_query; - ret = entity_->get_vector(rhs, vec_block_query); + ret = entity_->get_vector_borrowed(rhs, vec_block_query); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Get nullptr vector, id=%u", rhs); error_ = true; @@ -185,7 +185,7 @@ class HnswDistCalculator { compare_cnt_++; IndexStorage::MemoryBlock vec_block; - int ret = entity_->get_vector(id, vec_block); + int ret = entity_->get_vector_borrowed(id, vec_block); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Get nullptr vector, id=%u", id); error_ = true; @@ -216,6 +216,10 @@ class HnswDistCalculator { return error_; } + inline void set_error() { + error_ = true; + } + //! Get distances compute times inline uint32_t compare_cnt() const { return compare_cnt_; diff --git a/src/core/algorithm/hnsw/hnsw_entity.cc b/src/core/algorithm/hnsw/hnsw_entity.cc index aa0c0a3fe..b09776deb 100644 --- a/src/core/algorithm/hnsw/hnsw_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_entity.cc @@ -174,16 +174,19 @@ int64_t HnswEntity::dump_vectors( std::vector padding(padding_size); memset(padding.data(), 0, sizeof(char) * padding_size); - const void *data = nullptr; uint32_t crc = 0U; size_t vecs_size = 0UL; //! dump vectors for (node_id_t id = 0; id < doc_cnt(); ++id) { - data = get_vector(reorder_mapping.empty() ? id : reorder_mapping[id]); - if (ailego_unlikely(!data)) { + IndexStorage::MemoryBlock data_block; + if (ailego_unlikely(get_vector_borrowed( + reorder_mapping.empty() ? id : reorder_mapping[id], + data_block) != 0 || + !data_block.data())) { return IndexError_ReadData; } + const void *data = data_block.data(); size_t len = dumper->write(data, vector_size()); if (len != vector_size()) { LOG_ERROR("Dump vectors failed, write=%zu expect=%zu", len, diff --git a/src/core/algorithm/hnsw/hnsw_entity.h b/src/core/algorithm/hnsw/hnsw_entity.h index 639fc4ad4..e8b2a3c91 100644 --- a/src/core/algorithm/hnsw/hnsw_entity.h +++ b/src/core/algorithm/hnsw/hnsw_entity.h @@ -13,7 +13,7 @@ // limitations under the License. #pragma once -#include +#include #include #include #include @@ -78,11 +78,11 @@ struct HNSWHeader { } HNSWHeader(const HNSWHeader &header) { - memcpy(static_cast(this), &header, sizeof(header)); + std::memcpy(static_cast(this), &header, sizeof(header)); } HNSWHeader &operator=(const HNSWHeader &header) { - memcpy(static_cast(this), &header, sizeof(header)); + std::memcpy(static_cast(this), &header, sizeof(header)); return *this; } @@ -95,7 +95,7 @@ struct HNSWHeader { //! Clear all fields to init value void inline clear() { - memset(static_cast(this), 0, sizeof(HNSWHeader)); + std::memset(static_cast(this), 0, sizeof(HNSWHeader)); hnsw.entry_point = kInvalidNodeId; graph.size = sizeof(GraphHeader); hnsw.size = sizeof(HnswHeader); @@ -194,107 +194,7 @@ struct MmapMemoryBlock { void *data_{nullptr}; }; -//! Lightweight MemoryBlock for buffer pool mode: release on destruction -struct BufferPoolMemoryBlock { - BufferPoolMemoryBlock() = default; - - BufferPoolMemoryBlock(ailego::VecBufferPoolHandle *handle, size_t block_id, - void *data) - : buffer_pool_handle_(handle), buffer_block_id_(block_id), data_(data) {} - - static BufferPoolMemoryBlock MakeOwned(void *owned_data) { - BufferPoolMemoryBlock b; - b.owns_buffer_ = true; - b.data_ = owned_data; - return b; - } - - BufferPoolMemoryBlock(const BufferPoolMemoryBlock &rhs) - : buffer_pool_handle_(rhs.buffer_pool_handle_), - buffer_block_id_(rhs.buffer_block_id_), - data_(rhs.data_) { - if (rhs.owns_buffer_) { - owns_buffer_ = false; - buffer_pool_handle_ = nullptr; - } else if (buffer_pool_handle_) { - buffer_pool_handle_->acquire_one(buffer_block_id_); - } - } - - BufferPoolMemoryBlock &operator=(const BufferPoolMemoryBlock &rhs) { - if (this != &rhs) { - release(); - buffer_pool_handle_ = rhs.buffer_pool_handle_; - buffer_block_id_ = rhs.buffer_block_id_; - data_ = rhs.data_; - if (rhs.owns_buffer_) { - owns_buffer_ = false; - buffer_pool_handle_ = nullptr; - } else if (buffer_pool_handle_) { - buffer_pool_handle_->acquire_one(buffer_block_id_); - } - } - return *this; - } - - BufferPoolMemoryBlock(BufferPoolMemoryBlock &&rhs) noexcept - : buffer_pool_handle_(rhs.buffer_pool_handle_), - buffer_block_id_(rhs.buffer_block_id_), - owns_buffer_(rhs.owns_buffer_), - data_(rhs.data_) { - rhs.buffer_pool_handle_ = nullptr; - rhs.owns_buffer_ = false; - rhs.data_ = nullptr; - } - - BufferPoolMemoryBlock &operator=(BufferPoolMemoryBlock &&rhs) noexcept { - if (this != &rhs) { - release(); - buffer_pool_handle_ = rhs.buffer_pool_handle_; - buffer_block_id_ = rhs.buffer_block_id_; - owns_buffer_ = rhs.owns_buffer_; - data_ = rhs.data_; - rhs.buffer_pool_handle_ = nullptr; - rhs.owns_buffer_ = false; - rhs.data_ = nullptr; - } - return *this; - } - - ~BufferPoolMemoryBlock() { - release(); - } - - const void *data() const { - return data_; - } - - void reset(ailego::VecBufferPoolHandle *handle, size_t block_id, void *data) { - release(); - buffer_pool_handle_ = handle; - buffer_block_id_ = block_id; - data_ = data; - } - - private: - void release() { - if (owns_buffer_) { - if (data_) { - ailego_free(data_); - } - owns_buffer_ = false; - } else if (buffer_pool_handle_) { - buffer_pool_handle_->release_one(buffer_block_id_); - buffer_pool_handle_ = nullptr; - } - data_ = nullptr; - } - - ailego::VecBufferPoolHandle *buffer_pool_handle_{nullptr}; - size_t buffer_block_id_{0}; - bool owns_buffer_{false}; - void *data_{nullptr}; -}; +using BufferPoolMemoryBlock = IndexStorage::MemoryBlock; //! Typed Neighbors: holds a typed MemoryBlock to avoid runtime branching template @@ -629,6 +529,17 @@ class HnswEntity { return IndexError_NotImplemented; } + //! Fetch a vector for an operation bounded by this entity's lifetime. + //! + //! Storage-backed implementations may omit shared ownership from the + //! returned block because the caller guarantees that the entity remains + //! alive until the block is destroyed. Kept at the end of the vtable so + //! existing method slots remain stable. + virtual int get_vector_borrowed(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + return get_vector(id, block); + } + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, size_t data_size, size_t *padding_size); diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc index 1ad1ebd9c..296bc2cb6 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -228,7 +228,32 @@ int HnswStreamerEntity::get_vector(const node_id_t id, ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), "invalid chunk offset"); size_t read_size = vector_size(); - size_t ret = node_chunks_[loc.first]->read(loc.second, block, read_size); + size_t ret = + node_chunks_[loc.first]->read_immutable(loc.second, block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + return 0; +} + +int HnswStreamerEntity::get_vector_borrowed( + const node_id_t id, IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + + if (node_chunk_bases_ && loc.first < node_chunk_bases_->size() && + (*node_chunk_bases_)[loc.first]) { + block.reset((void *)((*node_chunk_bases_)[loc.first] + loc.second)); + return 0; + } + + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + const size_t read_size = vector_size(); + const size_t ret = node_chunks_[loc.first]->read_borrowed_immutable( + loc.second, block, read_size); if (ailego_unlikely(ret != read_size)) { LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", loc.second, read_size, ret); @@ -257,8 +282,8 @@ int HnswStreamerEntity::get_vector( ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), "invalid chunk offset"); size_t read_size = vector_size(); - size_t ret = - node_chunks_[loc.first]->read(loc.second, vec_blocks[i], read_size); + size_t ret = node_chunks_[loc.first]->read_immutable( + loc.second, vec_blocks[i], read_size); if (ailego_unlikely(ret != read_size)) { LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", loc.second, read_size, ret); @@ -268,6 +293,40 @@ int HnswStreamerEntity::get_vector( return 0; } +int HnswBufferPoolStreamerEntity::get_vector_for_prune( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + if (count == 0) { + return 0; + } + + // Unlike the search-time adaptive batch path, construction benefits from a + // batch even when every page is resident: each candidate is acquired only + // once and remains pinned through the entire prune. + static thread_local std::vector + batch_reads; + batch_reads.clear(); + batch_reads.reserve(count); + const size_t read_size = vector_size(); + for (uint32_t i = 0; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + batch_reads.emplace_back(node_chunks_[loc.first].get(), loc.second, + read_size, &vec_blocks[i]); + } + if (ailego_unlikely( + !batch_reads.front().segment->read_borrowed_batch_immutable( + batch_reads.data(), batch_reads.size()))) { + LOG_ERROR("Batch read prune vectors failed, count=%u, read size=%zu", + count, read_size); + return IndexError_ReadData; + } + return 0; +} + key_t HnswStreamerEntity::get_key(node_id_t id) const { if (use_key_info_map_) { auto loc = get_key_chunk_loc(id); @@ -303,16 +362,14 @@ void HnswStreamerEntity::add_neighbor(level_t level, node_id_t id, loc.second + sizeof(NeighborsHeader) + size * sizeof(node_id_t); ailego_assert_with(size < neighbor_cnt(level), "invalid neighbor size"); ailego_assert_with(offset < loc.first->data_size(), "invalid chunk offset"); - size_t ret = loc.first->write(offset, &neighbor_id, sizeof(node_id_t)); - if (ailego_unlikely(ret != sizeof(node_id_t))) { - LOG_ERROR("Write neighbor id failed, ret=%zu", ret); - return; - } - uint32_t neighbors = size + 1; - ret = loc.first->write(loc.second, &neighbors, sizeof(uint32_t)); - if (ailego_unlikely(ret != sizeof(uint32_t))) { - LOG_ERROR("Write neighbor cnt failed, ret=%zu", ret); + // Preserve publication order (payload before count), but let page-backed + // storage share one page pin and exclusive latch for both four-byte writes. + IndexStorage::SegmentData writes[2]; + writes[0] = {offset, sizeof(neighbor_id), &neighbor_id}; + writes[1] = {loc.second, sizeof(neighbors), &neighbors}; + if (ailego_unlikely(!loc.first->write_batch(writes, 2))) { + LOG_ERROR("Append neighbor failed"); } return; @@ -370,6 +427,32 @@ int HnswStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { return 0; } +void HnswStreamerEntity::protect_search_hotset() { + using CachePriority = IndexStorage::Segment::CachePriority; + + // Upper-level adjacency is a small fraction of the index and participates + // in every search, so keep it ahead of the much larger level-0/vector + // working set. BufferStorage performs a blocking batched prefetch here; + // mmap and other backends ignore the hint. + for (const auto &chunk : upper_neighbor_chunks_) { + chunk->prefetch(0, chunk->data_size(), CachePriority::kHigh); + } + + // Every search starts at the entry point. Protect its vector, key and L0 + // adjacency together because a node record may straddle two cache pages. + // Admit it last so it remains resident even when the upper graph is larger + // than an unusually small pool. + if (doc_cnt() == 0 || entry_point() == kInvalidNodeId) { + return; + } + const uint32_t chunk_idx = entry_point() >> node_index_mask_bits_; + const size_t offset = + static_cast(entry_point() & node_index_mask_) * node_size(); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + node_chunks_[chunk_idx]->prefetch(offset, node_size(), CachePriority::kHigh); +} + int HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc) { std::lock_guard lock(mutex_); @@ -456,6 +539,10 @@ int HnswStreamerEntity::open(IndexStorage::Pointer stg, uint64_t max_index_size, stats_.set_loaded_count(doc_cnt()); + if (storage_mode() == HnswStorageMode::kBufferPool) { + protect_search_hotset(); + } + return 0; } @@ -811,6 +898,39 @@ const HnswEntity::Pointer HnswMmapStreamerEntity::clone() const { return HnswEntity::Pointer(entity); } +const HnswEntity::Pointer HnswBufferPoolStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswBufferPoolStreamerEntity get node chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswBufferPoolStreamerEntity get upper chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) HnswBufferPoolStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_, + keys_map_, use_key_info_map_, std::move(node_chunks), + std::move(upper_neighbor_chunks), broker_, nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswBufferPoolStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const { std::vector node_chunks; node_chunks.reserve(node_chunks_.size()); diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity.h index 96d4b5506..30013a173 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -67,6 +67,9 @@ class HnswStreamerEntity : public HnswEntity { int get_vector(const node_id_t id, IndexStorage::MemoryBlock &block) const override; + int get_vector_borrowed(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + int get_vector( const node_id_t *ids, uint32_t count, std::vector &vec_blocks) const override; @@ -223,7 +226,7 @@ class HnswStreamerEntity : public HnswEntity { } inline size_t max_degree(level_t level) const { - return level == 0 ? neighbor_size_ : upper_neighbor_size_; + return neighbor_cnt(level); } @@ -532,6 +535,9 @@ class HnswStreamerEntity : public HnswEntity { //! Init node chunk and neighbor chunks int init_chunks(const Chunk::Pointer &header_chunk); + //! Preload the small, universally hot search root for BufferStorage. + void protect_search_hotset(); + int flush_header(void) { if (!broker_->dirty()) { // do not need to flush @@ -678,22 +684,12 @@ HnswStreamerEntity::get_neighbors_typed( } ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); IndexStorage::MemoryBlock mem_block; - size_t ret = chunk->read(offset, mem_block, nbr_size); + size_t ret = chunk->read_borrowed(offset, mem_block, nbr_size); if (ailego_unlikely(ret != nbr_size)) { LOG_ERROR("Read neighbor header failed, ret=%zu", ret); return NeighborsT(); } - BufferPoolMemoryBlock block; - if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { - block = BufferPoolMemoryBlock::MakeOwned(mem_block.data_); - mem_block.data_ = nullptr; - mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; - } else { - block = BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, - mem_block.buffer_block_id_, mem_block.data_); - mem_block.buffer_pool_handle_ = nullptr; - } - return NeighborsT(std::move(block)); + return NeighborsT(std::move(mem_block)); } //! MmapMemoryBlock specialization for batch get_vector @@ -726,33 +722,59 @@ inline int HnswStreamerEntity::get_vector_typed( const node_id_t *ids, uint32_t count, std::vector &vec_blocks) const { vec_blocks.resize(count); + if (count == 0) { + return 0; + } + + const auto first_loc = get_vector_chunk_loc(ids[0]); + ailego_assert_with(first_loc.first < node_chunks_.size(), + "invalid chunk idx"); + ailego_assert_with( + first_loc.second < node_chunks_[first_loc.first]->data_size(), + "invalid chunk offset"); + if (!node_chunks_[first_loc.first]->prefer_borrowed_batch_for( + vector_size())) { + const size_t read_size = vector_size(); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + IndexStorage::MemoryBlock mem_block; + const size_t ret = node_chunks_[loc.first]->read_borrowed_immutable( + loc.second, mem_block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + vec_blocks[i] = std::move(mem_block); + } + return 0; + } + + // Reuse request storage on each query thread. HNSW node ids span many + // storage chunks, so the batch deliberately crosses Segment boundaries; + // BufferStorage can then deduplicate all page misses in one AIO submission. + static thread_local std::vector + batch_reads; + batch_reads.clear(); + batch_reads.reserve(count); + const size_t read_size = vector_size(); for (auto i = 0U; i < count; ++i) { auto loc = get_vector_chunk_loc(ids[i]); ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), "invalid chunk offset"); - size_t read_size = vector_size(); - IndexStorage::MemoryBlock mem_block; - size_t ret = - node_chunks_[loc.first]->read(loc.second, mem_block, read_size); - if (ailego_unlikely(ret != read_size)) { - LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", - loc.second, read_size, ret); - return IndexError_ReadData; - } - vec_blocks[i] = [&]() { - if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { - BufferPoolMemoryBlock b = - BufferPoolMemoryBlock::MakeOwned(mem_block.data_); - mem_block.data_ = nullptr; - mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; - return b; - } - BufferPoolMemoryBlock b(mem_block.buffer_pool_handle_, - mem_block.buffer_block_id_, mem_block.data_); - mem_block.buffer_pool_handle_ = nullptr; - return b; - }(); + batch_reads.emplace_back(node_chunks_[loc.first].get(), loc.second, + read_size, &vec_blocks[i]); + } + if (ailego_unlikely( + !batch_reads.front().segment->read_borrowed_batch_immutable( + batch_reads.data(), batch_reads.size()))) { + LOG_ERROR("Batch read vectors failed, count=%u, read size=%zu", count, + read_size); + return IndexError_ReadData; } return 0; } @@ -789,8 +811,8 @@ inline key_t HnswStreamerEntity::get_key_typed( ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), "invalid chunk offset"); IndexStorage::MemoryBlock key_block; - size_t ret = - node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + size_t ret = node_chunks_[loc.first]->read_borrowed(loc.second, key_block, + sizeof(key_t)); if (ailego_unlikely(ret != sizeof(key_t))) { LOG_ERROR("Read key failed, ret=%zu", ret); return kInvalidKey; @@ -929,6 +951,11 @@ class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { return HnswStorageMode::kBufferPool; } + //! Keep the concrete entity type in per-query clones. HnswAlgorithm is + //! specialized on this type and statically dispatches the BufferStorage + //! accessors below. + const HnswEntity::Pointer clone() const override; + inline TypedNeighbors get_neighbors_typed(level_t level, node_id_t id) const { return HnswStreamerEntity::get_neighbors_typed(level, id); @@ -941,6 +968,10 @@ class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { ids, count, vec_blocks); } + int get_vector_for_prune( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + inline key_t get_key_typed(node_id_t id) const { return HnswStreamerEntity::get_key_typed(id); } diff --git a/src/core/algorithm/vamana/vamana_streamer_entity.h b/src/core/algorithm/vamana/vamana_streamer_entity.h index ecc3f70b4..4c0c3fa36 100644 --- a/src/core/algorithm/vamana/vamana_streamer_entity.h +++ b/src/core/algorithm/vamana/vamana_streamer_entity.h @@ -357,22 +357,13 @@ VamanaStreamerEntity::get_neighbors_typed( sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); IndexStorage::MemoryBlock mem_block; - size_t ret = node_chunks_[chunk_idx]->read(offset, mem_block, neighbor_size_); + size_t ret = + node_chunks_[chunk_idx]->read_borrowed(offset, mem_block, neighbor_size_); if (ailego_unlikely(ret != neighbor_size_)) { LOG_ERROR("Read neighbor header failed, ret=%zu", ret); return NeighborsT(); } - BufferPoolMemoryBlock block; - if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { - block = BufferPoolMemoryBlock::MakeOwned(mem_block.data_); - mem_block.data_ = nullptr; - mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; - } else { - block = BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, - mem_block.buffer_block_id_, mem_block.data_); - mem_block.buffer_pool_handle_ = nullptr; - } - return NeighborsT(std::move(block)); + return NeighborsT(std::move(mem_block)); } template <> @@ -403,25 +394,13 @@ inline int VamanaStreamerEntity::get_vector_typed( auto loc = get_vector_chunk_loc(ids[i]); ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); IndexStorage::MemoryBlock mem_block; - size_t ret = - node_chunks_[loc.first]->read(loc.second, mem_block, vector_size()); + size_t ret = node_chunks_[loc.first]->read_borrowed(loc.second, mem_block, + vector_size()); if (ailego_unlikely(ret != vector_size())) { LOG_ERROR("Read vector failed, ret=%zu", ret); return IndexError_ReadData; } - vec_blocks[i] = [&]() { - if (mem_block.type_ == IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH) { - BufferPoolMemoryBlock b = - BufferPoolMemoryBlock::MakeOwned(mem_block.data_); - mem_block.data_ = nullptr; - mem_block.type_ = IndexStorage::MemoryBlock::MBT_UNKNOWN; - return b; - } - BufferPoolMemoryBlock b(mem_block.buffer_pool_handle_, - mem_block.buffer_block_id_, mem_block.data_); - mem_block.buffer_pool_handle_ = nullptr; - return b; - }(); + vec_blocks[i] = std::move(mem_block); } return 0; } @@ -448,8 +427,8 @@ inline key_t VamanaStreamerEntity::get_key_typed( auto loc = get_key_chunk_loc(id); ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); IndexStorage::MemoryBlock key_block; - size_t ret = - node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + size_t ret = node_chunks_[loc.first]->read_borrowed(loc.second, key_block, + sizeof(key_t)); if (ailego_unlikely(ret != sizeof(key_t))) { LOG_ERROR("Read key failed, ret=%zu", ret); return kInvalidKey; diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index e377f5342..bdfcf41f3 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -18,6 +18,7 @@ #include #if DISKANN_SUPPORTED #include "algorithm/diskann/diskann_params.h" +#include "utility/utility_params.h" #include "holder_builder.h" #endif @@ -141,11 +142,7 @@ int DiskAnnIndex::Open(const std::string &file_path, file_path_ = file_path; is_read_only_ = storage_options.read_only; switch (storage_options.type) { - case StorageOptions::StorageType::kMMAP: - case StorageOptions::StorageType::kBufferPool: { - // NOTE: DiskAnn index is dumped via FileDumper (plain binary file), which - // is not compatible with BufferStorage's IndexFormat layout. Fall back to - // FileReadStorage for both MMAP and BufferPool storage types. + case StorageOptions::StorageType::kMMAP: { storage_ = core::IndexFactory::CreateStorage("FileReadStorage"); if (storage_ == nullptr) { LOG_ERROR("Failed to create FileReadStorage"); @@ -159,6 +156,22 @@ int DiskAnnIndex::Open(const std::string &file_path, } break; } + case StorageOptions::StorageType::kBufferPool: { + storage_params.set(core::BUFFER_READ_STORAGE_WARMUP_MODE, + core::BUFFER_READ_STORAGE_WARMUP_NONE); + storage_ = core::IndexFactory::CreateStorage("BufferReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create BufferReadStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init BufferReadStorage, path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } default: { LOG_ERROR("Unsupported storage type"); return core::IndexError_Unsupported; diff --git a/src/core/utility/buffer_read_storage.cc b/src/core/utility/buffer_read_storage.cc new file mode 100644 index 000000000..e45ea1542 --- /dev/null +++ b/src/core/utility/buffer_read_storage.cc @@ -0,0 +1,663 @@ +// 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. +// +// Read-only FileDumper storage backed by VecBufferPool instead of mmap. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "utility_params.h" + +namespace zvec { +namespace core { + +namespace { + +bool ResolveContainerOffset(size_t file_size, int64_t configured_offset, + bool zero_from_end, size_t *resolved_offset) { + const bool absolute = + configured_offset > 0 || (configured_offset == 0 && !zero_from_end); + if (absolute) { + const uint64_t value = static_cast(configured_offset); + if (value > file_size) { + return false; + } + *resolved_offset = static_cast(value); + return true; + } + + const uint64_t distance = + configured_offset == 0 + ? 0 + : static_cast(-(configured_offset + 1)) + 1; + if (distance > file_size) { + return false; + } + *resolved_offset = file_size - static_cast(distance); + return true; +} + +} // namespace + +/*! Buffer Read Storage (backed by VecBufferPool) + */ +class BufferReadStorage : public IndexStorage { + public: + /*! Read-only segment. Resident single-page reads pin cache pages; other + * reads use thread-local or owned scratch buffers. + */ + class Segment : public IndexStorage::Segment { + public: + //! Constructor + Segment(const std::shared_ptr &handle, + bool cache_enabled, size_t index_offset, + const IndexUnpacker::SegmentMeta &segment) + : data_offset_(index_offset + segment.data_offset()), + data_size_(segment.data_size()), + padding_size_(segment.padding_size()), + region_size_(segment.data_size() + segment.padding_size()), + data_crc_(segment.data_crc()), + handle_(handle), + cache_enabled_(cache_enabled) {} + + //! Constructor (clone) + Segment(const Segment &rhs) + : data_offset_(rhs.data_offset_), + data_size_(rhs.data_size_), + padding_size_(rhs.padding_size_), + region_size_(rhs.region_size_), + data_crc_(rhs.data_crc_), + handle_(rhs.handle_), + cache_enabled_(rhs.cache_enabled_) {} + + //! Destructor + ~Segment(void) override { + release_thread_scratch(); + } + + //! Retrieve size of data + size_t data_size(void) const override { + return data_size_; + } + + //! Retrieve the absolute data offset used by sector readers. + size_t data_offset(void) const override { + return data_offset_; + } + + //! Retrieve crc of data + uint32_t data_crc(void) const override { + return data_crc_; + } + + //! Retrieve size of padding + size_t padding_size(void) const override { + return padding_size_; + } + + //! Retrieve capacity of segment + size_t capacity(void) const override { + return region_size_; + } + + //! Fetch data from segment (copies into the caller-owned buffer) + size_t fetch(size_t offset, void *buf, size_t len) const override { + len = clamp_length(&offset, len); + if (len == 0) { + return 0; + } + if (!read_bytes(data_offset_ + offset, len, static_cast(buf))) { + LOG_ERROR( + "BufferReadStorage::Segment::fetch: read_range failed, " + "abs_offset=%zu, len=%zu", + data_offset_ + offset, len); + return 0; + } + return len; + } + + //! Read data from segment (stable until this thread's next pointer read) + size_t read(size_t offset, const void **data, size_t len) override { + if (ailego_unlikely(data == nullptr)) { + return 0; + } + auto scratch = thread_scratch(); + scratch->release_pin(); + len = clamp_length(&offset, len); + if (len == 0) { + *data = scratch->buffer.data(); + return 0; + } + const size_t abs_offset = data_offset_ + offset; + const size_t offset_in_page = abs_offset % ailego::kVectorPageSize; + bool force_bypass = !cache_enabled_; + if (cache_enabled_ && len <= ailego::kVectorPageSize - offset_in_page) { + size_t page_id = 0; + char *raw = scratch->handle->get_single_page(abs_offset, len, page_id); + if (raw != nullptr) { + scratch->pin(page_id); + *data = raw; + return len; + } + force_bypass = true; + } + scratch->buffer.resize(len); + const bool read_ok = + force_bypass + ? scratch->handle->read_range_bypass( + abs_offset, len, + reinterpret_cast(scratch->buffer.data())) + : read_bytes(abs_offset, len, + reinterpret_cast(scratch->buffer.data())); + if (!read_ok) { + LOG_ERROR( + "BufferReadStorage::Segment::read: read_range failed, " + "abs_offset=%zu, len=%zu", + abs_offset, len); + *data = nullptr; + return 0; + } + *data = scratch->buffer.data(); + return len; + } + + //! Read data from segment into a MemoryBlock + size_t read(size_t offset, MemoryBlock &data, size_t len) override { + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); + return 0; + } + return read_memory_block(data_offset_ + offset, data, len, + /*borrow_handle=*/false); + } + + //! Borrowed read; the caller keeps this Segment alive until block release. + size_t read_borrowed(size_t offset, MemoryBlock &data, + size_t len) override { + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); + return 0; + } + return read_memory_block(data_offset_ + offset, data, len, + /*borrow_handle=*/true); + } + + //! Read scattered data (stable until this thread's next pointer read) + bool read(SegmentData *iovec, size_t count) override { + ailego_false_if_false(iovec != nullptr && count != 0); + size_t total = 0u; + for (size_t i = 0; i < count; ++i) { + const SegmentData &item = iovec[i]; + ailego_false_if_false(item.offset <= region_size_); + ailego_false_if_false(item.length <= region_size_ - item.offset); + ailego_false_if_false(item.length <= + std::numeric_limits::max() - total); + total += item.length; + } + ailego_false_if_false(total != 0); + + auto scratch = thread_scratch(); + scratch->release_pin(); + scratch->buffer.resize(total); + uint8_t *buf = scratch->buffer.data(); + for (size_t i = 0; i < count; ++i) { + SegmentData *it = &iovec[i]; + ailego_false_if_false(read_bytes(data_offset_ + it->offset, it->length, + reinterpret_cast(buf))); + it->data = buf; + buf += it->length; + } + return true; + } + + size_t write(size_t, const void *, size_t) override { + return IndexError_NotImplemented; + } + + size_t resize(size_t) override { + return IndexError_NotImplemented; + } + + void update_data_crc(uint32_t) override {} + + //! Clone the segment + IndexStorage::Segment::Pointer clone(void) override { + return std::make_shared(*this); + } + + void prefetch(size_t offset, size_t len, + CachePriority priority = CachePriority::kLow) override { + if (!cache_enabled_) return; + len = clamp_length(&offset, len); + if (len == 0) return; + handle_->prefetch_range(data_offset_ + offset, len, + static_cast(priority)); + } + + //! Cached pages do not expose a stable base address. + const uint8_t *base_data(void) const override { + return nullptr; + } + + private: + struct ThreadScratch { + ThreadScratch( + const std::shared_ptr &owner_arg, + const std::shared_ptr &handle_arg) + : owner(owner_arg), handle(handle_arg) {} + + ThreadScratch(const ThreadScratch &) = delete; + ThreadScratch &operator=(const ThreadScratch &) = delete; + + ~ThreadScratch() { + release_pin(); + } + + void pin(size_t page_id) { + pinned_page_id = page_id; + page_pinned = true; + } + + void release_pin() { + if (page_pinned) { + handle->release_one(pinned_page_id); + page_pinned = false; + } + } + + std::weak_ptr owner; + std::vector buffer; + std::shared_ptr handle; + size_t pinned_page_id{0}; + bool page_pinned{false}; + }; + + struct ThreadScratchRegistry { + std::unordered_map scratches; + const uint8_t *last_key{nullptr}; + ThreadScratch *last_scratch{nullptr}; + }; + + static ThreadScratchRegistry &thread_scratch_registry() { + static thread_local ThreadScratchRegistry registry; + return registry; + } + + void release_thread_scratch() const { + ThreadScratchRegistry ®istry = thread_scratch_registry(); + const uint8_t *key = scratch_token_.get(); + if (registry.last_key == key) { + registry.last_key = nullptr; + registry.last_scratch = nullptr; + } + registry.scratches.erase(key); + } + + ThreadScratch *thread_scratch() const { + // Keep one scratch/pin per (thread, Segment); weak tokens clean up dead + // segments on long-lived worker threads. + ThreadScratchRegistry ®istry = thread_scratch_registry(); + const uint8_t *key = scratch_token_.get(); + if (registry.last_key == key && registry.last_scratch != nullptr && + !registry.last_scratch->owner.expired()) { + return registry.last_scratch; + } + + // Only Segment switches scan expired TLS entries. + registry.last_key = nullptr; + registry.last_scratch = nullptr; + for (auto iter = registry.scratches.begin(); + iter != registry.scratches.end();) { + if (iter->second.owner.expired()) { + iter = registry.scratches.erase(iter); + } else { + ++iter; + } + } + auto result = + registry.scratches.try_emplace(key, scratch_token_, handle_); + registry.last_key = key; + registry.last_scratch = &result.first->second; + return registry.last_scratch; + } + + bool read_bytes(size_t abs_offset, size_t len, char *out) const { + return cache_enabled_ ? handle_->read_range(abs_offset, len, out) + : handle_->read_range_bypass(abs_offset, len, out); + } + + size_t read_memory_block(size_t abs_offset, MemoryBlock &data, size_t len, + bool borrow_handle) const { + const size_t offset_in_page = abs_offset % ailego::kVectorPageSize; + bool force_bypass = !cache_enabled_; + if (cache_enabled_ && len <= ailego::kVectorPageSize - offset_in_page) { + size_t page_id = 0; + char *raw = handle_->get_single_page(abs_offset, len, page_id); + if (raw != nullptr) { + if (borrow_handle) { + data.reset(handle_.get(), page_id, raw); + } else { + data.reset(handle_, page_id, raw); + } + return len; + } + force_bypass = true; + } + // Copy cross-page and bypass reads into owned memory. + return read_owned(abs_offset, data, len, force_bypass); + } + + size_t read_owned(size_t abs_offset, MemoryBlock &data, size_t len, + bool force_bypass) const { + static constexpr size_t kAlign = 4096UL; + if (ailego_unlikely(len > + std::numeric_limits::max() - (kAlign - 1))) { + LOG_ERROR( + "BufferReadStorage::Segment::read(MemoryBlock&): cross-page " + "length overflow, abs_offset=%zu, len=%zu", + abs_offset, len); + return 0; + } + size_t alloc_size = (len + (kAlign - 1UL)) & ~(kAlign - 1UL); + char *tmp = + static_cast(ailego_aligned_malloc(alloc_size, kAlign)); + if (!tmp) { + LOG_ERROR( + "BufferReadStorage::Segment::read(MemoryBlock&): cross-page alloc " + "failed, abs_offset=%zu, len=%zu", + abs_offset, len); + return 0; + } + const bool read_ok = + force_bypass ? handle_->read_range_bypass(abs_offset, len, tmp) + : read_bytes(abs_offset, len, tmp); + if (!read_ok) { + ailego_free(tmp); + LOG_ERROR( + "BufferReadStorage::Segment::read(MemoryBlock&): cross-page " + "read_range failed, abs_offset=%zu, len=%zu", + abs_offset, len); + return 0; + } + data = MemoryBlock::MakeOwned(tmp, len); + return len; + } + size_t clamp_length(size_t *offset, size_t len) const { + if (ailego_unlikely(*offset > region_size_)) { + *offset = region_size_; + return 0; + } + return std::min(len, region_size_ - *offset); + } + + size_t data_offset_{0u}; + size_t data_size_{0u}; + size_t padding_size_{0u}; + size_t region_size_{0u}; + uint32_t data_crc_{0u}; + std::shared_ptr scratch_token_{ + std::make_shared(0)}; + std::shared_ptr handle_{nullptr}; + bool cache_enabled_{false}; + }; + + //! Destructor + ~BufferReadStorage(void) override = default; + + //! Initialize container + int init(const ailego::Params ¶ms) override { + params.get(BUFFER_READ_STORAGE_CHECKSUM_VALIDATION, &checksum_validation_); + params.get(BUFFER_READ_STORAGE_HEADER_OFFSET, &header_offset_); + params.get(BUFFER_READ_STORAGE_FOOTER_OFFSET, &footer_offset_); + params.get(BUFFER_READ_STORAGE_WARMUP_MODE, &warmup_mode_); + if (warmup_mode_ != BUFFER_READ_STORAGE_WARMUP_NONE && + warmup_mode_ != BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL) { + LOG_ERROR("Invalid BufferReadStorage warmup mode: %s", + warmup_mode_.c_str()); + return IndexError_InvalidArgument; + } + return 0; + } + + int flush(void) override { + return 0; + } + + int append(const std::string &, size_t) override { + return IndexError_NotImplemented; + } + + void refresh(uint64_t) override {} + + uint64_t check_point(void) const override { + return 0; + } + + //! Cleanup container + int cleanup(void) override { + return this->close(); + } + + //! Load an index file into the container + int open(const std::string &path, bool) override { + try { + // Publish new state only after open succeeds. + auto candidate_pool = + std::make_shared(path, /*writable=*/false); + auto candidate_handle = + std::make_shared(candidate_pool); + + const size_t file_size = candidate_pool->file_size(); + const size_t page_count = + file_size == 0 ? 0 : (file_size - 1) / ailego::kVectorPageSize + 1; + const size_t metadata_bytes = + ailego::VecBufferPool::metadata_bytes_for_page_count(page_count); + size_t candidate_index_offset = 0; + size_t end_offset = 0; + if (!ResolveContainerOffset(file_size, header_offset_, + /*zero_from_end=*/false, + &candidate_index_offset) || + !ResolveContainerOffset(file_size, footer_offset_, + /*zero_from_end=*/true, &end_offset) || + candidate_index_offset >= end_offset) { + LOG_ERROR( + "Invalid BufferReadStorage container offsets: path=%s " + "file_size=%zu header_offset=%lld footer_offset=%lld", + path.c_str(), file_size, static_cast(header_offset_), + static_cast(footer_offset_)); + return IndexError_InvalidArgument; + } + const size_t container_size = end_offset - candidate_index_offset; + + // IndexUnpacker requires a stable pointer until its next callback. + std::vector scratch; + auto read_data = [&candidate_handle, &scratch, candidate_index_offset, + container_size](size_t offset, const void **data, + size_t len) -> size_t { + if (offset > container_size) { + offset = container_size; + len = 0; + } else { + len = std::min(len, container_size - offset); + } + scratch.resize(len); + *data = scratch.data(); + if (len == 0) { + return 0; + } + const size_t file_offset = candidate_index_offset + offset; + if (candidate_handle->get_meta( + file_offset, len, reinterpret_cast(scratch.data())) != + 0) { + return 0; + } + return len; + }; + + IndexUnpacker unpacker; + if (!unpacker.unpack(read_data, container_size, checksum_validation_)) { + LOG_ERROR("Failed to unpack file: %s", path.c_str()); + return IndexError_UnpackIndex; + } + auto candidate_segments = std::move(*unpacker.mutable_segments()); + for (const auto &item : candidate_segments) { + const auto &segment = item.second; + const size_t segment_offset = segment.data_offset(); + if (segment_offset > container_size || + segment.data_size() > container_size - segment_offset || + segment.padding_size() > + container_size - segment_offset - segment.data_size()) { + LOG_ERROR( + "Invalid BufferReadStorage segment bounds: path=%s id=%s " + "container_size=%zu offset=%zu data_size=%zu padding_size=%zu", + path.c_str(), item.first.c_str(), container_size, segment_offset, + segment.data_size(), segment.padding_size()); + return IndexError_InvalidLength; + } + } + const uint32_t candidate_magic = unpacker.magic(); + + // Fall back to bypass-only mode when metadata plus one page cannot fit. + bool candidate_cache_enabled = false; + const size_t available = + ailego::MemoryLimitPool::get_instance().available(); + if (metadata_bytes != std::numeric_limits::max() && + metadata_bytes <= available && + ailego::kVectorPageSize <= available - metadata_bytes) { + candidate_cache_enabled = candidate_pool->init() == 0; + } + if (!candidate_cache_enabled) { + LOG_INFO( + "BufferReadStorage opened in bypass-only mode: path=%s " + "available=%zu metadata=%zu page_size=%zu", + path.c_str(), available, metadata_bytes, ailego::kVectorPageSize); + } + if (candidate_cache_enabled && + warmup_mode_ == BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL) { + candidate_pool->warmup(); + } + + file_path_ = path; + index_offset_ = candidate_index_offset; + magic_ = candidate_magic; + segments_ = std::move(candidate_segments); + handle_ = std::move(candidate_handle); + buffer_pool_ = std::move(candidate_pool); + cache_enabled_ = candidate_cache_enabled; + return 0; + } catch (const std::bad_alloc &) { + LOG_ERROR("Out of memory opening BufferReadStorage: %s", path.c_str()); + return IndexError_NoMemory; + } catch (const std::runtime_error &error) { + LOG_ERROR("Failed to open BufferReadStorage file %s: %s", path.c_str(), + error.what()); + return IndexError_OpenFile; + } catch (const std::exception &error) { + LOG_ERROR("Unexpected BufferReadStorage open failure for %s: %s", + path.c_str(), error.what()); + return IndexError_Runtime; + } catch (...) { + LOG_ERROR("Unknown BufferReadStorage open failure for %s", path.c_str()); + return IndexError_Runtime; + } + } + + int close(void) override { + segments_.clear(); + handle_ = nullptr; + buffer_pool_ = nullptr; + cache_enabled_ = false; + return 0; + } + + //! Retrieve a segment by id + IndexStorage::Segment::Pointer get(const std::string &id, int) override { + if (!handle_) { + return {}; + } + auto it = segments_.find(id); + if (it == segments_.end()) { + return {}; + } + return std::make_shared( + handle_, cache_enabled_, index_offset_, it->second); + } + + std::map get_all( + void) const override { + std::map result; + if (handle_) { + for (const auto &it : segments_) { + result.emplace(it.first, + std::make_shared( + handle_, cache_enabled_, index_offset_, it.second)); + } + } + return result; + } + + //! Test if a segment exists + bool has(const std::string &id) const override { + return segments_.find(id) != segments_.end(); + } + + //! Retrieve magic number of index + uint32_t magic(void) const override { + return magic_; + } + + //! Reads go through the VecBufferPool paged cache. + MemoryBlock::MemoryBlockType memory_block_type(void) const override { + return MemoryBlock::MBT_BUFFERPOOL; + } + + std::shared_ptr vec_buffer_pool(void) const override { + return cache_enabled_ ? buffer_pool_ : nullptr; + } + + //! Path of the opened index file (diagnostics / backend consistency). + std::string file_path(void) const override { + return file_path_; + } + + private: + bool checksum_validation_{false}; + std::string warmup_mode_{BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL}; + int64_t header_offset_{0}; + int64_t footer_offset_{0}; + size_t index_offset_{0}; + uint32_t magic_{0}; + std::string file_path_{}; + std::map segments_{}; + std::shared_ptr buffer_pool_{nullptr}; + std::shared_ptr handle_{nullptr}; + bool cache_enabled_{false}; +}; + +INDEX_FACTORY_REGISTER_STORAGE(BufferReadStorage); + +} // namespace core +} // namespace zvec diff --git a/src/core/utility/buffer_storage.cc b/src/core/utility/buffer_storage.cc index 3db6c58d9..bf980bbc8 100644 --- a/src/core/utility/buffer_storage.cc +++ b/src/core/utility/buffer_storage.cc @@ -12,18 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include #include +#include #include #include +#include #include -#include +#include #include #include #include +#include #include #include -#include #include #include #include @@ -34,16 +35,17 @@ namespace zvec { namespace core { namespace { -// Cross-compiler helpers for lock-free 64-bit acquire/release access -// to SegmentMeta::data_size / padding_size. -// -// These fields are POD (uint64_t) inside a serialised struct so we cannot -// change their type to std::atomic<>; std::atomic_ref is C++20 and the -// project targets C++17. GCC/Clang have native __atomic_* builtins that -// emit single ldar/stlr on arm64 and plain mov on x86_64. MSVC lacks -// these builtins, so we fall back to volatile load/store paired with a -// std::atomic_thread_fence, which is correct on all targets MSVC ships -// (x86_64 / arm64 desktop) and equivalent in cost. +constexpr size_t kBufferAlignment = 4096UL; + +inline bool AlignBufferSize(size_t size, size_t *aligned_size) { + if (size > std::numeric_limits::max() - (kBufferAlignment - 1)) { + return false; + } + *aligned_size = (size + (kBufferAlignment - 1)) & ~(kBufferAlignment - 1); + return true; +} + +// C++17-compatible atomic access to serialized uint64_t fields. inline uint64_t bs_load_acquire(const uint64_t *p) { #if defined(__GNUC__) || defined(__clang__) return __atomic_load_n(p, __ATOMIC_ACQUIRE); @@ -81,11 +83,8 @@ inline void bs_store_relaxed(uint64_t *p, uint64_t v) { } // namespace -// The legacy read(const void**) overload guarantees the returned pointer -// stays valid until close_index(). Single-page reads pin the page -// (never released); cross-page reads allocate a temp buffer owned by -// tmp_buffers_ (freed in close_index()). Callers wanting bounded -// lifetime should use the read(MemoryBlock&) overload. +// Legacy pointer reads stay valid until close_index(); prefer MemoryBlock for +// bounded ownership. /*! Buffer Storage */ @@ -96,32 +95,30 @@ class BufferStorage : public IndexStorage { class WrappedSegment : public IndexStorage::Segment, public std::enable_shared_from_this { public: - //! Index Storage Pointer - typedef std::shared_ptr Pointer; - //! Constructor. See segment_info_ for the pointer-stability contract. WrappedSegment(BufferStorage *owner, IndexMapping::SegmentInfo *info, - size_t segment_id) + const std::string *segment_id) : segment_info_(info), owner_(owner), segment_id_(segment_id), capacity_(static_cast(info->segment.meta()->data_size + info->segment.meta()->padding_size)) {} //! Destructor - ~WrappedSegment(void) override {} - - //! Retrieve size of data - //! - //! data_size / padding_size are mutated lock-free by concurrent - //! writers (write/resize) and observed by concurrent readers on the - //! lock-free hot path. Use acquire/release ordering so weakly-ordered - //! ARM (e.g. Android arm64) cannot see stale values that would cause - //! read() to truncate len to 0. + ~WrappedSegment(void) override = default; + + //! Retrieve size of data. Paired acquire/release operations publish + //! concurrent write/resize metadata changes to lock-free readers. size_t data_size(void) const override { return static_cast( bs_load_acquire(&segment_info_->segment.meta()->data_size)); } + size_t data_offset(void) const override { + return segment_info_->segment_header_start_offset + + segment_info_->segment_header->content_offset + + segment_info_->segment.meta()->data_index; + } + //! Retrieve crc of data uint32_t data_crc(void) const override { return segment_info_->segment.meta()->data_crc; @@ -139,102 +136,71 @@ class BufferStorage : public IndexStorage { } //! Fetch data from segment (with own buffer) - //! - //! C1: pool/handle are stable for the lifetime of the index - //! (no retire/rebuild), so no lock is needed on the hot path. size_t fetch(size_t offset, void *buf, size_t len) const override { if (ailego_unlikely(!owner_->buffer_pool_handle_)) { - LOG_ERROR("WrappedSegment::fetch: handle is null, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + LOG_ERROR("WrappedSegment::fetch: handle is null, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } - const size_t data_size = - bs_load_acquire(&segment_info_->segment.meta()->data_size); - if (ailego_unlikely(offset > data_size || len > data_size - offset)) { - if (offset > data_size) { - offset = data_size; - } - len = data_size - offset; + len = clamp_length(&offset, len); + if (len == 0) { + return 0; } - size_t abs_offset = segment_info_->segment_header_start_offset + - segment_info_->segment_header->content_offset + - segment_info_->segment.meta()->data_index + offset; - if (!owner_->buffer_pool_handle_->read_range(abs_offset, len, - static_cast(buf))) { + const size_t abs_offset = absolute_offset(offset); + if (!owner_->read_range(abs_offset, len, static_cast(buf))) { LOG_ERROR( - "WrappedSegment::fetch: read_range failed, file[%s], id[%zu], " + "WrappedSegment::fetch: read_range failed, file[%s], id[%s], " "abs_offset=%zu, len=%zu", - owner_->file_name_.c_str(), segment_id_, abs_offset, len); + owner_->file_name_.c_str(), segment_id_->c_str(), abs_offset, len); return 0; } return len; } //! Read data from segment - //! C1: lock-free hot path (pool/handle never change during operation). size_t read(size_t offset, const void **data, size_t len) override { if (ailego_unlikely(!owner_->buffer_pool_handle_)) { - LOG_ERROR("WrappedSegment::read: handle is null, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + LOG_ERROR("WrappedSegment::read: handle is null, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); *data = nullptr; return 0; } - const size_t data_size = - bs_load_acquire(&segment_info_->segment.meta()->data_size); - if (ailego_unlikely(offset > data_size || len > data_size - offset)) { - if (offset > data_size) { - offset = data_size; - } - len = data_size - offset; + len = clamp_length(&offset, len); + if (len == 0) { + *data = nullptr; + return 0; } - size_t abs_offset = segment_info_->segment_header_start_offset + - segment_info_->segment_header->content_offset + - segment_info_->segment.meta()->data_index + offset; + const size_t abs_offset = absolute_offset(offset); + // Without a page cache, retain a copied result for the legacy pointer + // lifetime. Cached single-page reads below pin the page until close(), + // including for writable pools, instead of retaining one copy per read. size_t first_page = abs_offset / ailego::kVectorPageSize; - size_t last_page = (len == 0) - ? first_page - : (abs_offset + len - 1) / ailego::kVectorPageSize; - if (first_page == last_page) { + size_t last_page = (abs_offset + len - 1) / ailego::kVectorPageSize; + bool force_bypass = !owner_->cache_enabled_; + if (owner_->cache_enabled_ && first_page == last_page) { size_t page_id = 0; char *raw = owner_->buffer_pool_handle_->get_single_page(abs_offset, len, page_id); - if (!raw) { - LOG_ERROR( - "WrappedSegment::read: single-page acquire failed, file[%s], " - "id[%zu], abs_offset=%zu, len=%zu, page=%zu", - owner_->file_name_.c_str(), segment_id_, abs_offset, len, - first_page); - *data = nullptr; - return 0; - } - *data = raw; - // Pin held until close_index() per the never-released contract - // of this overload. Record the page so close_index() can release - // the pin before tearing down the pool; otherwise the lingering - // ref_count would trip ~VecBufferPool's "all blocks released" - // assertion. - { - std::lock_guard pin_latch(owner_->pinned_pages_mutex_); - owner_->pinned_pages_.push_back(page_id); + if (raw != nullptr) { + *data = raw; + // Legacy pointer reads retain their pin until close_index(). + { + std::lock_guard pin_latch(owner_->pinned_pages_mutex_); + owner_->pinned_pages_.push_back(page_id); + } + return len; } - return len; + force_bypass = true; + } + // Keep scratch buffers 4K-aligned without over-allocating on platforms + // whose native page size is larger. + size_t alloc_size = 0; + if (ailego_unlikely(!AlignBufferSize(len, &alloc_size))) { + *data = nullptr; + return 0; } - // Cross-page path: see file-level banner. C11 aligned_alloc requires - // size to be a multiple of alignment, and alignment must be a power - // of two. Use a fixed 4096-byte alignment for the dst buffer: 4K is - // the minimum page granularity across all supported platforms - // (always a divisor of the 16K/64K page sizes used on Apple Silicon - // and some Android arm64 configurations) and is sufficient for the - // downstream SIMD/DMA-friendly access contract. Pinning kAlign to - // 4096 also avoids over-allocating 16KB per cross-page read on - // large-page platforms. - static constexpr size_t kAlign = 4096UL; - size_t alloc_size = (len + (kAlign - 1UL)) & ~(kAlign - 1UL); - // Allocate a 4K-aligned slot from the per-storage arena pool. - // This batches page-aligned allocation: under heap fragmentation - // (notably Android Bionic scudo), one large posix_memalign per - // arena via the secondary (mmap-backed) allocator is far more - // reliable than many independent posix_memalign(4K, 4K) calls. + // The arena amortizes aligned allocation and avoids fragmented-heap + // failures observed with many small aligned allocations on Android. char *tmp = nullptr; { std::lock_guard tmp_latch(owner_->tmp_buffers_mutex_); @@ -243,21 +209,24 @@ class BufferStorage : public IndexStorage { if (!tmp) { LOG_ERROR( "WrappedSegment::read: cross-page alloc failed, file[%s], " - "id[%zu], abs_offset=%zu, len=%zu, alloc_size=%zu, align=%zu", - owner_->file_name_.c_str(), segment_id_, abs_offset, len, - alloc_size, kAlign); + "id[%s], abs_offset=%zu, len=%zu, alloc_size=%zu, align=%zu", + owner_->file_name_.c_str(), segment_id_->c_str(), abs_offset, len, + alloc_size, kBufferAlignment); *data = nullptr; return 0; } - if (!owner_->buffer_pool_handle_->read_range(abs_offset, len, tmp)) { + const bool read_ok = force_bypass + ? owner_->buffer_pool_handle_->read_range_bypass( + abs_offset, len, tmp) + : owner_->read_range(abs_offset, len, tmp); + if (!read_ok) { LOG_ERROR( "WrappedSegment::read: cross-page read_range failed, file[%s], " - "id[%zu], abs_offset=%zu, len=%zu, first_page=%zu, last_page=%zu", - owner_->file_name_.c_str(), segment_id_, abs_offset, len, + "id[%s], abs_offset=%zu, len=%zu, first_page=%zu, last_page=%zu", + owner_->file_name_.c_str(), segment_id_->c_str(), abs_offset, len, first_page, last_page); - // The arena slot is intentionally not rolled back: rolling back - // would require holding the arena lock across read_range, while - // the worst-case leak per failed read is one slot (alloc_size). + // Avoid holding the arena lock across I/O; a failed read loses one + // temporary slot. *data = nullptr; return 0; } @@ -265,83 +234,527 @@ class BufferStorage : public IndexStorage { return len; } - //! C1: lock-free hot path (pool/handle never change during operation). + //! Read data into a bounded-lifetime block. size_t read(size_t offset, MemoryBlock &data, size_t len) override { if (ailego_unlikely(!owner_->buffer_pool_handle_)) { LOG_ERROR( "WrappedSegment::read(MemoryBlock&): handle is null, file[%s], " - "id[%zu]", - owner_->file_name_.c_str(), segment_id_); + "id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } - const size_t data_size = - bs_load_acquire(&segment_info_->segment.meta()->data_size); - if (ailego_unlikely(offset > data_size || len > data_size - offset)) { - if (offset > data_size) { - offset = data_size; - } - len = data_size - offset; + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); + return 0; } - size_t abs_offset = segment_info_->segment_header_start_offset + - segment_info_->segment_header->content_offset + - segment_info_->segment.meta()->data_index + offset; - size_t first_page = abs_offset / ailego::kVectorPageSize; - size_t last_page = (len == 0) - ? first_page - : (abs_offset + len - 1) / ailego::kVectorPageSize; - if (first_page == last_page) { - size_t page_id = 0; - char *raw = owner_->buffer_pool_handle_->get_single_page(abs_offset, - len, page_id); - if (!raw) { - LOG_ERROR("read error (single-page acquire failed)."); - return 0; - } - data.reset(owner_->buffer_pool_handle_.get(), page_id, raw); - return len; + return read_memory_block(absolute_offset(offset), data, len, + /*borrow_handle=*/false, + /*immutable=*/false); + } + + size_t read_immutable(size_t offset, MemoryBlock &data, + size_t len) override { + if (ailego_unlikely(!owner_->buffer_pool_handle_)) { + return 0; } - // C11 aligned_alloc requires the requested size to be a multiple of - // the alignment, and alignment must be a power of two. See the - // sibling read(const void**) overload above for the rationale of - // pinning kAlign to a fixed 4096 instead of sysconf(_SC_PAGESIZE). - static constexpr size_t kAlign = 4096UL; - size_t alloc_size = (len + (kAlign - 1UL)) & ~(kAlign - 1UL); - char *tmp = - static_cast(ailego_aligned_malloc(alloc_size, kAlign)); - if (!tmp) { - LOG_ERROR("read error (alloc cross-page temp buffer failed)."); + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); return 0; } - if (!owner_->buffer_pool_handle_->read_range(abs_offset, len, tmp)) { - ailego_free(tmp); - LOG_ERROR("read error (cross-page read_range failed)."); + return read_memory_block(absolute_offset(offset), data, len, + /*borrow_handle=*/false, + /*immutable=*/true); + } + + //! Borrowed read: the caller keeps the Segment alive until block release. + size_t read_borrowed(size_t offset, MemoryBlock &data, + size_t len) override { + if (ailego_unlikely(!owner_->buffer_pool_handle_)) { + LOG_ERROR( + "WrappedSegment::read_borrowed: handle is null, file[%s], " + "id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } - data = MemoryBlock::MakeOwned(tmp, len); - return len; + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); + return 0; + } + return read_memory_block(absolute_offset(offset), data, len, + /*borrow_handle=*/true, + /*immutable=*/false); + } + + size_t read_borrowed_immutable(size_t offset, MemoryBlock &data, + size_t len) override { + if (ailego_unlikely(!owner_->buffer_pool_handle_)) { + return 0; + } + len = clamp_length(&offset, len); + if (len == 0) { + data.reset(); + return 0; + } + return read_memory_block(absolute_offset(offset), data, len, + /*borrow_handle=*/true, + /*immutable=*/true); + } + + bool prefer_borrowed_batch() const override { + return owner_->cache_enabled_ && owner_->buffer_pool_ != nullptr && + !owner_->buffer_pool_->writable() && + owner_->buffer_pool_->has_evicted(); + } + + bool prefer_borrowed_batch_for(size_t value_size) const override { + if (!owner_->cache_enabled_ || owner_->buffer_pool_ == nullptr) { + return false; + } + if (!owner_->buffer_pool_->writable()) { + return owner_->buffer_pool_->has_evicted(); + } + // With uniformly distributed records, crossing a page costs roughly + // value_size/page_size. Below one eighth, scalar pins are cheaper than + // constructing and resolving a batch (notably 512 B on macOS 16 KiB). + return value_size >= std::max(1, ailego::kVectorPageSize / 8); + } + + bool read_borrowed_batch(BorrowedRead *reads, size_t count) override { + return read_borrowed_batch_impl(reads, count, /*immutable=*/false); + } + + bool read_borrowed_batch_immutable(BorrowedRead *reads, + size_t count) override { + return read_borrowed_batch_impl(reads, count, /*immutable=*/true); + } + + bool read_borrowed_batch_impl(BorrowedRead *reads, size_t count, + bool immutable) { + if (count == 0) { + return true; + } + if (reads == nullptr || owner_->buffer_pool_handle_ == nullptr || + owner_->buffer_pool_ == nullptr) { + return false; + } + + // Mutable reads from writable pools retain snapshot ownership rules. + if (!owner_->cache_enabled_ || + (owner_->buffer_pool_->writable() && !immutable)) { + return immutable + ? IndexStorage::Segment::read_borrowed_batch_immutable(reads, + count) + : IndexStorage::Segment::read_borrowed_batch(reads, count); + } + + // Cross-page vectors need a contiguous copy. A thread-local scratch arena + // is reused across calls instead of a per-vector malloc/free: the arena + // is safe to recycle because the previous hop's blocks are destroyed + // before this call runs (see fast_search_neighbors_buffer). Cross-page + // results become non-owning views into the arena. On by default; + // ZVEC_CROSS_ARENA=0 restores the per-vector malloc path. + static const bool cross_arena = [] { + const char *v = std::getenv("ZVEC_CROSS_ARENA"); + return !(v != nullptr && *v == '0'); + }(); + // Cap the per-thread arena: a batch needing more falls back to malloc so + // an abnormally wide batch cannot pin unbounded out-of-pool RSS. HNSW + // hops need only about max_degree * vector_size (tens of KiB). + static constexpr size_t kMaxArenaBytes = 2UL << 20; // 2 MiB / thread + bool batch_arena = false; + + struct BatchState { + BorrowedRead *request{nullptr}; + size_t abs_offset{0}; + size_t first_page_index{0}; + size_t page_count{0}; + bool copy_result{false}; + char *owned{nullptr}; + }; + struct PageUse { + size_t unique_index{0}; + size_t state_index{0}; + size_t source_offset{0}; + size_t destination_offset{0}; + size_t length{0}; + }; + struct BatchScratch { + std::vector states; + std::vector page_ids; + std::vector unique_page_ids; + std::vector pages; + std::vector page_uses; + std::vector admitted_ids; + std::vector admitted_indices; + std::vector admitted_pages; + std::vector arena; // reused cross-page scratch (arena mode) + std::vector bypass_page; + }; + static thread_local BatchScratch scratch; + + scratch.states.clear(); + scratch.page_ids.clear(); + scratch.unique_page_ids.clear(); + scratch.pages.clear(); + scratch.page_uses.clear(); + scratch.admitted_ids.clear(); + scratch.admitted_indices.clear(); + scratch.admitted_pages.clear(); + scratch.states.reserve(count); + for (size_t i = 0; i < count; ++i) { + if (reads[i].block != nullptr) { + reads[i].block->reset(); + } + } + + auto cleanup_owned = [&]() { + // Arena slices are non-owning; only malloc-mode buffers are freed. + for (BatchState &state : scratch.states) { + if (!batch_arena && state.owned != nullptr) { + ailego_free(state.owned); + } + state.owned = nullptr; + } + }; + auto release_pages = [&]() { + for (size_t i = 0; i < scratch.pages.size(); ++i) { + if (scratch.pages[i] != nullptr) { + owner_->buffer_pool_handle_->release_one( + scratch.unique_page_ids[i]); + scratch.pages[i] = nullptr; + } + } + }; + auto fail = [&]() { + cleanup_owned(); + release_pages(); + for (size_t i = 0; i < count; ++i) { + if (reads[i].block != nullptr) { + reads[i].block->reset(); + } + } + return false; + }; + + size_t total_pages = 0; + for (size_t i = 0; i < count; ++i) { + BorrowedRead &request = reads[i]; + auto *segment = dynamic_cast(request.segment); + if (segment == nullptr || segment->owner_ != owner_ || + request.block == nullptr) { + cleanup_owned(); + return immutable + ? IndexStorage::Segment::read_borrowed_batch_immutable( + reads, count) + : IndexStorage::Segment::read_borrowed_batch(reads, count); + } + + const size_t data_size = segment->data_size(); + if (request.offset > data_size || + request.length > data_size - request.offset) { + return fail(); + } + + BatchState state; + state.request = &request; + if (request.length == 0) { + scratch.states.emplace_back(state); + continue; + } + + const size_t data_offset = segment->data_offset(); + if (data_offset > std::numeric_limits::max() - request.offset) { + return fail(); + } + state.abs_offset = data_offset + request.offset; + if (state.abs_offset > + std::numeric_limits::max() - (request.length - 1)) { + return fail(); + } + const size_t first_page = state.abs_offset / ailego::kVectorPageSize; + const size_t last_page = + (state.abs_offset + request.length - 1) / ailego::kVectorPageSize; + state.first_page_index = total_pages; + state.page_count = last_page - first_page + 1; + if (state.page_count > + std::numeric_limits::max() - total_pages) { + return fail(); + } + total_pages += state.page_count; + scratch.states.emplace_back(state); + for (size_t page = first_page; page <= last_page; ++page) { + scratch.page_ids.emplace_back(page); + } + } + + // Pin unique pages rather than one pin per vector occurrence. Besides + // removing duplicate work, this lets a pressured batch preserve all + // resident hits and fall back only for pages that really cannot be + // admitted. MemoryLimitPool itself enforces the shared non-page reserve; + // checking available bytes here used to turn the entire batch into + // per-vector bypass reads as soon as the cache reached that reserve. + scratch.unique_page_ids = scratch.page_ids; + std::sort(scratch.unique_page_ids.begin(), + scratch.unique_page_ids.end()); + scratch.unique_page_ids.erase( + std::unique(scratch.unique_page_ids.begin(), + scratch.unique_page_ids.end()), + scratch.unique_page_ids.end()); + scratch.pages.assign(scratch.unique_page_ids.size(), nullptr); + + // Take resident hits without I/O first, then admit only unique misses + // selected by the compact frequency policy. One-off cold pages bypass + // the cache, while repeated HNSW graph pages become resident without + // continuously replacing useful members of the working set. + for (size_t i = 0; i < scratch.unique_page_ids.size(); ++i) { + scratch.pages[i] = owner_->buffer_pool_->try_acquire_buffer( + scratch.unique_page_ids[i]); + if (scratch.pages[i] == nullptr) { + // Under pressure, a first-touch HNSW page is usually a one-off + // random candidate. Bypass it once; a repeated touch is admitted by + // the pool's compact frequency policy. This avoids replacing one + // useful resident page for every miss when the working set is much + // larger than the budget. + if (owner_->buffer_pool_->should_admit_page( + scratch.unique_page_ids[i])) { + scratch.admitted_ids.push_back(scratch.unique_page_ids[i]); + scratch.admitted_indices.push_back(i); + } + } + } + + // Bound each acquisition so an oversized diagnostic batch cannot hold + // the entire cache pinned. Linux retains batched AIO within each chunk; + // on macOS, failed capacity resolution stops after one chunk instead of + // rescanning the same pinned cache for every remaining page. + static constexpr size_t kAdmissionBatchPages = 64; + scratch.admitted_pages.resize(scratch.admitted_ids.size(), nullptr); + size_t admitted_begin = 0; + while (admitted_begin < scratch.admitted_ids.size()) { + const size_t admitted_count = + std::min(kAdmissionBatchPages, + scratch.admitted_ids.size() - admitted_begin); + const bool acquired = owner_->buffer_pool_handle_->acquire_pages( + scratch.admitted_ids.data() + admitted_begin, admitted_count, + scratch.admitted_pages.data() + admitted_begin); + if (acquired) { + for (size_t j = 0; j < admitted_count; ++j) { + scratch.pages[scratch.admitted_indices[admitted_begin + j]] = + scratch.admitted_pages[admitted_begin + j]; + } + admitted_begin += admitted_count; + continue; + } + + // acquire_pages() rolls its pins back on failure, but it may have + // populated part of the chunk before capacity ran out. Re-pin those + // pages once so that useful work and concurrent single-flight loads + // are not discarded, then bypass the remaining cold pages. + for (size_t j = 0; j < admitted_count; ++j) { + const size_t admitted_index = admitted_begin + j; + const size_t unique_index = + scratch.admitted_indices[admitted_index]; + scratch.pages[unique_index] = + owner_->buffer_pool_->try_acquire_buffer( + scratch.admitted_ids[admitted_index]); + } + break; + } + + // Close races with other query threads after admission and AIO. This is + // a hit-only probe; it never turns a bypass candidate into new I/O. + for (size_t i = 0; i < scratch.unique_page_ids.size(); ++i) { + if (scratch.pages[i] == nullptr) { + scratch.pages[i] = owner_->buffer_pool_->try_acquire_buffer( + scratch.unique_page_ids[i]); + } + } + + auto unique_index_for = [&](ailego::block_id_t page_id) -> size_t { + auto it = std::lower_bound(scratch.unique_page_ids.begin(), + scratch.unique_page_ids.end(), page_id); + return static_cast(it - scratch.unique_page_ids.begin()); + }; + + // Cross-page values always need a contiguous copy. A single-page value + // only needs one when that page is on the bypass side of the hybrid + // batch; resident single-page values keep their normal zero-copy pin. + for (BatchState &state : scratch.states) { + state.copy_result = state.page_count > 1; + for (size_t j = 0; j < state.page_count && !state.copy_result; ++j) { + const size_t unique_index = + unique_index_for(scratch.page_ids[state.first_page_index + j]); + state.copy_result = scratch.pages[unique_index] == nullptr; + } + } + + if (cross_arena) { + // Reserve copied values once, then hand out bump slices. This includes + // single-page bypass values, so the direct-read page scratch can be + // immediately reused for the next unique miss. + static constexpr size_t kArenaAlign = 64; + size_t total = 0; + for (const BatchState &state : scratch.states) { + if (!state.copy_result) { + continue; + } + const size_t length = state.request->length; + if (length > std::numeric_limits::max() - + (kArenaAlign - 1)) { + return fail(); + } + const size_t aligned = + (length + kArenaAlign - 1) & ~(kArenaAlign - 1); + if (aligned > kMaxArenaBytes - total) { + total = kMaxArenaBytes + 1; + break; + } + total += aligned; + } + batch_arena = total <= kMaxArenaBytes; + if (batch_arena) { + if (scratch.arena.size() < total) { + scratch.arena.resize(total); + } + size_t off = 0; + for (BatchState &state : scratch.states) { + if (!state.copy_result) { + continue; + } + state.owned = scratch.arena.data() + off; + off += (state.request->length + kArenaAlign - 1) & + ~(kArenaAlign - 1); + } + } + } + if (!batch_arena) { + for (BatchState &state : scratch.states) { + if (!state.copy_result) { + continue; + } + const size_t length = state.request->length; + size_t alloc_size = 0; + if (!AlignBufferSize(length, &alloc_size)) { + return fail(); + } + state.owned = static_cast( + ailego_aligned_malloc(alloc_size, kBufferAlignment)); + if (state.owned == nullptr) { + return fail(); + } + } + } + + // Describe every copied fragment, sort by unique source page, and read + // each bypass page exactly once before scattering it to all vectors that + // overlap that page. + scratch.page_uses.reserve(total_pages); + for (size_t state_index = 0; state_index < scratch.states.size(); + ++state_index) { + const BatchState &state = scratch.states[state_index]; + if (!state.copy_result) { + continue; + } + size_t remaining = state.request->length; + size_t destination_offset = 0; + size_t source_offset = state.abs_offset % ailego::kVectorPageSize; + for (size_t j = 0; j < state.page_count; ++j) { + const size_t length = + std::min(remaining, ailego::kVectorPageSize - source_offset); + scratch.page_uses.push_back(PageUse{ + unique_index_for( + scratch.page_ids[state.first_page_index + j]), + state_index, source_offset, destination_offset, length}); + destination_offset += length; + remaining -= length; + source_offset = 0; + } + } + std::sort(scratch.page_uses.begin(), scratch.page_uses.end(), + [](const PageUse &lhs, const PageUse &rhs) { + return lhs.unique_index < rhs.unique_index; + }); + + if (scratch.bypass_page.size() < ailego::kVectorPageSize) { + scratch.bypass_page.resize(ailego::kVectorPageSize); + } + size_t use_begin = 0; + while (use_begin < scratch.page_uses.size()) { + const size_t unique_index = + scratch.page_uses[use_begin].unique_index; + const char *source = scratch.pages[unique_index]; + if (source == nullptr) { + const size_t page_offset = + scratch.unique_page_ids[unique_index] * + ailego::kVectorPageSize; + const size_t read_length = std::min( + ailego::kVectorPageSize, + owner_->buffer_pool_->file_size() - page_offset); + if (!owner_->buffer_pool_handle_->read_range_bypass( + page_offset, read_length, scratch.bypass_page.data())) { + return fail(); + } + source = scratch.bypass_page.data(); + } + size_t use_end = use_begin; + while (use_end < scratch.page_uses.size() && + scratch.page_uses[use_end].unique_index == unique_index) { + const PageUse &use = scratch.page_uses[use_end]; + BatchState &state = scratch.states[use.state_index]; + std::memcpy(state.owned + use.destination_offset, + source + use.source_offset, use.length); + ++use_end; + } + use_begin = use_end; + } + + for (BatchState &state : scratch.states) { + if (state.page_count == 0) { + state.request->block->reset(); + continue; + } + if (!state.copy_result) { + const size_t unique_index = unique_index_for( + scratch.page_ids[state.first_page_index]); + const size_t offset_in_page = + state.abs_offset % ailego::kVectorPageSize; + owner_->buffer_pool_handle_->acquire_one( + scratch.unique_page_ids[unique_index]); + state.request->block->reset(owner_->buffer_pool_handle_.get(), + scratch.unique_page_ids[unique_index], + scratch.pages[unique_index] + + offset_in_page); + continue; + } + *state.request->block = + batch_arena + ? MemoryBlock::MakeBorrowedView(state.owned) + : MemoryBlock::MakeOwned(state.owned, state.request->length); + state.owned = nullptr; + } + release_pages(); + return true; } - //! Write data into the storage with offset. - //! - //! Locking: shared shard latch pairs with flush_index()'s exclusive - //! all-shards latch -- excludes CRC compute over meta_buf while we - //! mutate (data_size, padding_size). meta_mtx_ additionally - //! serialises concurrent writers on the SAME segment so the pair - //! stays consistent (sum == capacity_). + //! Write data into the storage with offset. The shard latch excludes + //! flush; meta_mtx_ serializes metadata changes within this segment. size_t write(size_t offset, const void *data, size_t len) override { std::shared_lock latch( owner_->mapping_shards_[owner_->mapping_shard_id()].mtx); if (ailego_unlikely(!owner_->buffer_pool_handle_ || !owner_->buffer_pool_)) { - LOG_ERROR("WrappedSegment::write: pool is null, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + LOG_ERROR("WrappedSegment::write: pool is null, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } if (ailego_unlikely(owner_->corrupted_.load(std::memory_order_acquire))) { LOG_ERROR( "WrappedSegment::write: storage is marked corrupted, refusing " - "write, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + "write, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } // In read-only mode the write is a silent no-op so that callers that @@ -359,39 +772,97 @@ class BufferStorage : public IndexStorage { size_t abs_offset = segment_info_->segment_header_start_offset + segment_info_->segment_header->content_offset + meta->data_index + offset; - // Write the bytes BEFORE publishing the new data_size to readers. - // Lock-free readers observe data_size with acquire ordering; the - // release-store below establishes happens-before with the page - // contents written above. Publishing data_size first (the previous - // ordering) allowed a reader on weakly-ordered ARM to see the new - // length but still read stale page contents -- or, in the inverse - // direction, see a stale length and truncate len to 0 - // (root cause of "Read sparse vector failed ... ret=0"). + // Publish data_size only after the page bytes are visible to readers. if (owner_->buffer_pool_handle_->write_range( abs_offset, len, static_cast(data)) != 0) { LOG_ERROR("write() page-cache write_range failed at abs_offset=%zu", abs_offset); return 0; } - { + const uint64_t write_end = offset + len; + if (write_end > bs_load_acquire(&meta->data_size)) { std::lock_guard meta_latch(meta_mtx_); uint64_t cur = bs_load_relaxed(&meta->data_size); - if (offset + len > cur) { - uint64_t new_size = offset + len; - // padding_size is paired with data_size; publish it first - // (relaxed) so readers that acquire data_size see a - // consistent (data_size + padding_size == capacity_) pair. - bs_store_relaxed(&meta->padding_size, capacity_ - new_size); - bs_store_release(&meta->data_size, new_size); + if (write_end > cur) { + // Publish padding before data_size to keep the pair consistent. + bs_store_relaxed(&meta->padding_size, capacity_ - write_end); + bs_store_release(&meta->data_size, write_end); } } - // Mark dirty unconditionally even when data_size did not grow: - // fixed-size in-place rewrites (e.g. chunk_meta_segment) must still - // trigger flush_all() before the next append_segment(). + // Fixed-size rewrites are dirty even when data_size is unchanged. owner_->set_as_dirty(); return len; } + bool write_batch(const SegmentData *writes, size_t count) override { + if (count == 0) { + return true; + } + if (writes == nullptr || count > kMaxCoalescedWrites) { + return IndexStorage::Segment::write_batch(writes, count); + } + + std::shared_lock latch( + owner_->mapping_shards_[owner_->mapping_shard_id()].mtx); + if (ailego_unlikely(!owner_->buffer_pool_handle_ || + !owner_->buffer_pool_ || + owner_->corrupted_.load(std::memory_order_acquire))) { + return false; + } + if (!owner_->buffer_pool_->writable()) { + return true; + } + + auto meta = segment_info_->segment.meta(); + const size_t data_base = segment_info_->segment_header_start_offset + + segment_info_->segment_header->content_offset + + meta->data_index; + std::array + fragments{}; + size_t page_id = std::numeric_limits::max(); + uint64_t write_end = 0; + bool has_data = false; + for (size_t i = 0; i < count; ++i) { + const auto &write = writes[i]; + if (ailego_unlikely(write.offset > capacity_ || + write.length > capacity_ - write.offset || + (write.length != 0 && write.data == nullptr))) { + return false; + } + if (write.length == 0) { + continue; + } + const size_t abs_offset = data_base + write.offset; + fragments[i] = {abs_offset, write.length, + static_cast(write.data)}; + write_end = std::max(write_end, write.offset + write.length); + const size_t write_page = abs_offset / ailego::kVectorPageSize; + const size_t offset_in_page = abs_offset % ailego::kVectorPageSize; + if (write.length > ailego::kVectorPageSize - offset_in_page || + (has_data && write_page != page_id)) { + latch.unlock(); + return IndexStorage::Segment::write_batch(writes, count); + } + page_id = write_page; + has_data = true; + } + + if (has_data && owner_->buffer_pool_handle_->write_fragments( + fragments.data(), count) != 0) { + return false; + } + if (write_end > bs_load_acquire(&meta->data_size)) { + std::lock_guard meta_latch(meta_mtx_); + const uint64_t current = bs_load_relaxed(&meta->data_size); + if (write_end > current) { + bs_store_relaxed(&meta->padding_size, capacity_ - write_end); + bs_store_release(&meta->data_size, write_end); + } + } + owner_->set_as_dirty(); + return true; + } + //! Resize size of data. See write() for the locking contract. size_t resize(size_t size) override { std::shared_lock latch( @@ -399,8 +870,8 @@ class BufferStorage : public IndexStorage { if (ailego_unlikely(owner_->corrupted_.load(std::memory_order_acquire))) { LOG_ERROR( "WrappedSegment::resize: storage is marked corrupted, refusing " - "resize, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + "resize, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return 0; } auto meta = segment_info_->segment.meta(); @@ -412,9 +883,7 @@ class BufferStorage : public IndexStorage { if (size > capacity_) { size = capacity_; } - // See write() for the publish ordering rationale: padding first - // (relaxed), then release-store data_size so concurrent lock-free - // readers observe a consistent pair. + // Match write(): padding first, then release-store data_size. bs_store_relaxed(&meta->padding_size, capacity_ - size); bs_store_release(&meta->data_size, size); changed = true; @@ -433,8 +902,8 @@ class BufferStorage : public IndexStorage { if (ailego_unlikely(owner_->corrupted_.load(std::memory_order_acquire))) { LOG_ERROR( "WrappedSegment::update_data_crc: storage is marked corrupted, " - "refusing CRC update, file[%s], id[%zu]", - owner_->file_name_.c_str(), segment_id_); + "refusing CRC update, file[%s], id[%s]", + owner_->file_name_.c_str(), segment_id_->c_str()); return; } { @@ -449,19 +918,120 @@ class BufferStorage : public IndexStorage { return shared_from_this(); } - protected: - friend BufferStorage; - // Pointer into BufferStorage::segments_ (unordered_map mapped value). - // The address is stable across map insertions, so re-parses after - // append_segment() are picked up without recreating WrappedSegment. + //! Preload a read-only range and attach its eviction priority. Writable + //! storage deliberately skips this hint: construction has a different + //! access pattern and dirty-page lifetime must drive admission there. + void prefetch(size_t offset, size_t len, + CachePriority priority = CachePriority::kLow) override { + if (!owner_->cache_enabled_ || !owner_->buffer_pool_ || + !owner_->buffer_pool_handle_ || owner_->buffer_pool_->writable()) { + return; + } + const size_t data_size = + bs_load_acquire(&segment_info_->segment.meta()->data_size); + if (offset >= data_size || len == 0) { + return; + } + len = std::min(len, data_size - offset); + const size_t abs_offset = data_offset() + offset; + owner_->buffer_pool_handle_->prefetch_range( + abs_offset, len, static_cast(priority)); + } + + private: + static constexpr size_t kMaxCoalescedWrites = 8; + // Stable unordered_map value address; reparses update this object in place. IndexMapping::SegmentInfo *segment_info_{nullptr}; - // Serialises hot-path writers on the SAME segment so - // (data_size, padding_size, data_crc) updates do not interleave. + // Serializes metadata writers within this segment. mutable std::mutex meta_mtx_{}; - private: + size_t clamp_length(size_t *offset, size_t len) const { + const size_t current_size = data_size(); + if (ailego_unlikely(*offset > current_size)) { + *offset = current_size; + return 0; + } + return std::min(len, current_size - *offset); + } + + size_t absolute_offset(size_t offset) const { + return data_offset() + offset; + } + + size_t read_memory_block(size_t abs_offset, MemoryBlock &data, size_t len, + bool borrow_handle, bool immutable) const { + const bool can_pin = owner_->cache_enabled_ && + (immutable || !owner_->buffer_pool_->writable()); + bool force_bypass = !owner_->cache_enabled_; + const size_t offset_in_page = abs_offset % ailego::kVectorPageSize; + if (can_pin && len <= ailego::kVectorPageSize - offset_in_page) { + size_t page_id = 0; + char *raw = owner_->buffer_pool_handle_->get_single_page(abs_offset, + len, page_id); + if (raw != nullptr) { + if (borrow_handle) { + data.reset(owner_->buffer_pool_handle_.get(), page_id, raw); + } else { + data.reset(owner_->buffer_pool_handle_, page_id, raw); + } + return len; + } + force_bypass = true; + } + + // Mutable graph snapshots are typically only a few hundred bytes. They + // are copied from cached pages and are never submitted to O_DIRECT, so a + // page-aligned allocation only adds allocator and RSS overhead. Keep + // alignment for larger/vector snapshots whose distance kernels benefit + // from it, but let the allocator's thread cache handle small objects. + static constexpr size_t kSmallSnapshotBytes = 512; + size_t alloc_size = len; + char *tmp = nullptr; + if (!immutable && len <= kSmallSnapshotBytes) { + tmp = static_cast(ailego_malloc(len)); + } else { + if (ailego_unlikely(!AlignBufferSize(len, &alloc_size))) { + return 0; + } + tmp = static_cast( + ailego_aligned_malloc(alloc_size, kBufferAlignment)); + } + if (tmp == nullptr) { + LOG_ERROR( + "WrappedSegment::read: owned buffer allocation failed, file[%s], " + "id[%s], abs_offset=%zu, len=%zu", + owner_->file_name_.c_str(), segment_id_->c_str(), abs_offset, len); + return 0; + } + bool read_ok = + force_bypass ? owner_->buffer_pool_handle_->read_range_bypass( + abs_offset, len, tmp) + : immutable ? owner_->buffer_pool_handle_->read_range_immutable( + abs_offset, len, tmp) + : owner_->read_range(abs_offset, len, tmp); + if (!read_ok && immutable && !force_bypass) { + // A cross-page immutable read can lose a capacity race after choosing + // the cache path. The bytes are immutable, so retry directly instead + // of failing the HNSW expansion. + read_ok = owner_->buffer_pool_handle_->read_range_bypass(abs_offset, + len, tmp); + } + if (!read_ok) { + ailego_free(tmp); + LOG_ERROR( + "WrappedSegment::read: owned read failed, file[%s], id[%s], " + "abs_offset=%zu, len=%zu", + owner_->file_name_.c_str(), segment_id_->c_str(), abs_offset, len); + return 0; + } + data = MemoryBlock::MakeOwned(tmp, len); + return len; + } + BufferStorage *owner_{nullptr}; - size_t segment_id_{}; + // Stable alongside segment_info_; unordered_map rehash preserves element + // references and pointers. + const std::string *segment_id_{nullptr}; size_t capacity_{}; }; @@ -475,6 +1045,10 @@ class BufferStorage : public IndexStorage { return MemoryBlock::MBT_BUFFERPOOL; } + std::shared_ptr vec_buffer_pool(void) const override { + return cache_enabled_ ? buffer_pool_ : nullptr; + } + //! Initialize storage int init(const ailego::Params ¶ms) override { uint32_t val = params.get_as_uint32(MMAPFILE_STORAGE_SEGMENT_META_CAPACITY); @@ -506,41 +1080,54 @@ class BufferStorage : public IndexStorage { } } - // Open in writable mode when the caller expects to modify the index - // (create_if_missing=true implies write intent, same as MMapFileStorage). + // create_if_missing also indicates write intent, matching MMapFileStorage. buffer_pool_ = std::make_shared( path, /*writable=*/create_if_missing); - buffer_pool_handle_ = std::make_shared( - buffer_pool_->get_handle()); + buffer_pool_handle_ = + std::make_shared(buffer_pool_); int ret = ParseToMapping(); if (ret != 0) { this->close_index(); return ret; } - ret = buffer_pool_->init(); - if (ret != 0) { + const size_t file_size = buffer_pool_->file_size(); + const size_t page_count = + file_size == 0 ? 0 : (file_size - 1) / ailego::kVectorPageSize + 1; + const size_t metadata_bytes = + ailego::VecBufferPool::metadata_bytes_for_page_count( + page_count, /*writable=*/create_if_missing); + const size_t available = + ailego::MemoryLimitPool::get_instance().available(); + const bool cache_can_fit = + metadata_bytes != std::numeric_limits::max() && + metadata_bytes <= available && + ailego::kVectorPageSize <= available - metadata_bytes; + ret = cache_can_fit ? buffer_pool_->init() : -1; + if (ret != 0 && create_if_missing) { this->close_index(); - return ret; + return IndexError_NoMemory; } - LOG_INFO( - "BufferStorage opened: file=%s, writable=%d, max_segment_size=%" PRIu64 - ", segment_count=%zu", - file_name_.c_str(), static_cast(create_if_missing), - max_segment_size_, segments_.size()); + cache_enabled_ = ret == 0; + if (!cache_enabled_) { + LOG_INFO( + "Read-only BufferStorage opened in bypass-only mode: file=%s " + "available=%zu metadata=%zu page_size=%zu", + path.c_str(), available, metadata_bytes, ailego::kVectorPageSize); + } + LOG_INFO("BufferStorage opened: file=%s, writable=%d, segment_count=%zu", + file_name_.c_str(), static_cast(create_if_missing), + segments_.size()); return 0; } - // PRECONDITION (also for ParseFooter/ParseSegment/ParseToMapping): - // caller holds either single-threaded open() or AllShardsExclusiveLatch. - // Do NOT add an internal lock here -- std::shared_mutex is not reentrant. + // Called from single-threaded open or under AllShardsExclusiveLatch. int ParseHeader(size_t offset, IndexFormat::MetaHeader *out) { constexpr size_t kHeaderSize = sizeof(IndexFormat::MetaHeader); - std::unique_ptr buffer(new char[kHeaderSize]); - if (buffer_pool_handle_->get_meta(offset, kHeaderSize, buffer.get()) != 0) { + if (buffer_pool_handle_->get_meta(offset, kHeaderSize, + reinterpret_cast(out)) != 0) { LOG_ERROR("Get segment header failed."); return IndexError_Runtime; } - memcpy(out, buffer.get(), kHeaderSize); if (out->meta_header_size != kHeaderSize) { LOG_ERROR("Header meta size is invalid."); return IndexError_InvalidLength; @@ -553,72 +1140,69 @@ class BufferStorage : public IndexStorage { return 0; } - int ParseFooter(size_t offset) { - std::unique_ptr buffer(new char[sizeof(footer_)]); - if (buffer_pool_handle_->get_meta(offset, sizeof(footer_), buffer.get()) != - 0) { + int ParseFooter(size_t offset, IndexFormat::MetaFooter *footer) { + if (buffer_pool_handle_->get_meta(offset, sizeof(*footer), + reinterpret_cast(footer)) != 0) { LOG_ERROR("Get segment footer failed."); return IndexError_Runtime; } - uint8_t *footer_ptr = reinterpret_cast(buffer.get()); - memcpy(&footer_, footer_ptr, sizeof(footer_)); - if (offset < (size_t)footer_.segments_meta_size) { + if (offset < static_cast(footer->segments_meta_size)) { LOG_ERROR("Footer meta size is invalid."); return IndexError_InvalidLength; } - if (ailego::Crc32c::Hash(&footer_, sizeof(footer_), footer_.footer_crc) != - footer_.footer_crc) { + if (ailego::Crc32c::Hash(footer, sizeof(*footer), footer->footer_crc) != + footer->footer_crc) { LOG_ERROR("Footer meta checksum is invalid."); return IndexError_InvalidChecksum; } return 0; } - int ParseSegment(size_t offset, IndexFormat::MetaHeader *chain_header, - uint32_t *out_segment_ids_offset) { - std::unique_ptr segment_buffer = - std::make_unique(footer_.segments_meta_size); - if (buffer_pool_handle_->get_meta(offset, footer_.segments_meta_size, + int ParseSegment(size_t offset, uint64_t header_start_offset, + IndexFormat::MetaHeader *chain_header, + const IndexFormat::MetaFooter &footer, + uint32_t &out_segment_ids_offset, + std::unique_ptr &segment_buffer) { + segment_buffer = std::make_unique(footer.segments_meta_size); + if (buffer_pool_handle_->get_meta(offset, footer.segments_meta_size, segment_buffer.get()) != 0) { LOG_ERROR("Get segment meta failed."); return IndexError_Runtime; } - if (ailego::Crc32c::Hash(segment_buffer.get(), footer_.segments_meta_size, - 0u) != footer_.segments_meta_crc) { + if (ailego::Crc32c::Hash(segment_buffer.get(), footer.segments_meta_size, + 0u) != footer.segments_meta_crc) { LOG_ERROR("Index segments meta checksum is invalid."); return IndexError_InvalidChecksum; } + if (sizeof(IndexFormat::SegmentMeta) * footer.segment_count > + footer.segments_meta_size) { + return IndexError_InvalidLength; + } IndexFormat::SegmentMeta *segment_start = reinterpret_cast(segment_buffer.get()); - uint32_t segment_ids_offset = footer_.segments_meta_size; + uint32_t segment_ids_offset = footer.segments_meta_size; for (IndexFormat::SegmentMeta *iter = segment_start, - *end = segment_start + footer_.segment_count; + *end = segment_start + footer.segment_count; iter != end; ++iter) { - if (iter->segment_id_offset >= footer_.segments_meta_size) { + if (iter->segment_id_offset >= footer.segments_meta_size) { return IndexError_InvalidValue; } - if (iter->data_index > footer_.content_size) { + if (iter->data_index > footer.content_size) { return IndexError_InvalidValue; } - if (iter->data_index + iter->data_size > footer_.content_size) { + if (iter->data_index + iter->data_size > footer.content_size) { return IndexError_InvalidLength; } if (iter->segment_id_offset < segment_ids_offset) { segment_ids_offset = iter->segment_id_offset; } - // Use id_hash_.size() (not segments_.size()) for the block_id: - // segments_ is intentionally NOT cleared between appends to keep - // existing WrappedSegment pointers valid, so it carries stale entries. - // - // Bound the C-string scan to the segments_meta buffer so a missing - // NUL terminator cannot walk past the buffer end (defence against - // crafted-CRC inputs; CRC already covers benign bit flips). + // Bound ID parsing to the metadata buffer. const char *seg_name_start = reinterpret_cast(segment_start) + iter->segment_id_offset; const size_t seg_name_max = - footer_.segments_meta_size - iter->segment_id_offset; + footer.segments_meta_size - iter->segment_id_offset; const size_t seg_name_len = ::strnlen(seg_name_start, seg_name_max); if (seg_name_len == seg_name_max) { LOG_ERROR("ParseSegment: segment_id missing NUL terminator, file[%s]", @@ -626,36 +1210,21 @@ class BufferStorage : public IndexStorage { return IndexError_InvalidValue; } const std::string seg_name(seg_name_start, seg_name_len); - const size_t seg_id = id_hash_.size(); - id_hash_[seg_name] = seg_id; - // In-place update so existing WrappedSegment pointers see the - // refreshed meta_ptr_ after re-parse. chain_header MUST be the - // per-chain owning copy (not a shared &header_) -- see - // chain_headers_ field comment. - segments_[seg_name] = - IndexMapping::SegmentInfo{IndexMapping::Segment{iter}, - current_header_start_offset_, chain_header}; - max_segment_size_ = - std::max(max_segment_size_, iter->data_size + iter->padding_size); - if (sizeof(IndexFormat::SegmentMeta) * footer_.segment_count > - footer_.segments_meta_size) { - return IndexError_InvalidLength; - } - } - buffer_pool_buffers_.push_back(std::move(segment_buffer)); - if (out_segment_ids_offset) { - *out_segment_ids_offset = segment_ids_offset; + // Update in place so existing WrappedSegment pointers remain valid. + segments_[seg_name] = IndexMapping::SegmentInfo{ + IndexMapping::Segment{iter}, header_start_offset, chain_header}; } + out_segment_ids_offset = segment_ids_offset; return 0; } int ParseToMapping() { + uint64_t header_start_offset = 0; while (true) { int ret; - // Per-chain owning MetaHeader; see chain_headers_ field comment. - chain_headers_.emplace_back(std::make_unique()); - IndexFormat::MetaHeader *chain_header = chain_headers_.back().get(); - ret = ParseHeader(current_header_start_offset_, chain_header); + auto header = std::make_unique(); + IndexFormat::MetaHeader *chain_header = header.get(); + ret = ParseHeader(header_start_offset, chain_header); if (ret != 0) { LOG_ERROR("Failed to parse header, errno %d, %s", ret, IndexError::What(ret)); @@ -678,18 +1247,19 @@ class BufferStorage : public IndexStorage { return IndexError_Unsupported; } uint64_t footer_offset = - chain_header->meta_footer_offset + current_header_start_offset_; + chain_header->meta_footer_offset + header_start_offset; // Reject uint64 wrap-around and offsets past file_size. - if (footer_offset < current_header_start_offset_ || + if (footer_offset < header_start_offset || footer_offset + sizeof(IndexFormat::MetaFooter) > buffer_pool_->file_size()) { LOG_ERROR("ParseToMapping: invalid footer_offset=%" PRIu64 " (header=%" PRIu64 ", file_size=%zu), file[%s]", - footer_offset, current_header_start_offset_, - buffer_pool_->file_size(), file_name_.c_str()); + footer_offset, header_start_offset, buffer_pool_->file_size(), + file_name_.c_str()); return IndexError_InvalidValue; } - ret = ParseFooter(footer_offset); + IndexFormat::MetaFooter footer{}; + ret = ParseFooter(footer_offset, &footer); if (ret != 0) { LOG_ERROR("Failed to parse footer, errno %d, %s", ret, IndexError::What(ret)); @@ -697,54 +1267,47 @@ class BufferStorage : public IndexStorage { } // Unpack segment table - if (sizeof(IndexFormat::SegmentMeta) * footer_.segment_count > - footer_.segments_meta_size) { - return IndexError_InvalidLength; - } const uint64_t segment_start_offset = - footer_offset - footer_.segments_meta_size; - uint32_t segment_ids_offset = footer_.segments_meta_size; + footer_offset - footer.segments_meta_size; + uint32_t segment_ids_offset = footer.segments_meta_size; + std::unique_ptr segment_buffer; ret = - ParseSegment(segment_start_offset, chain_header, &segment_ids_offset); + ParseSegment(segment_start_offset, header_start_offset, chain_header, + footer, segment_ids_offset, segment_buffer); if (ret != 0) { LOG_ERROR("Failed to parse segment, errno %d, %s", ret, IndexError::What(ret)); return ret; } - // Record per-chain metadata offsets so flush_index() can write - // updated segment metas and footers back to the backing file. - meta_chains_.push_back({current_header_start_offset_, footer_offset, - segment_start_offset, footer_.segments_meta_size, - segment_ids_offset, footer_}); + // Own all state referenced by this metadata chain in one object. + meta_chains_.push_back({std::move(header), std::move(segment_buffer), + header_start_offset, segment_ids_offset, footer}); - if (footer_.next_meta_header_offset == 0) { + if (footer.next_meta_header_offset == 0) { break; } - // Reject self-reference / backward jumps and offsets past file_size: - // such a corrupted next_meta_header_offset would otherwise drive the - // loop into infinite chain growth -> OOM. - const uint64_t next_off = footer_.next_meta_header_offset; - if (next_off <= current_header_start_offset_ || + // Reject invalid links before following the metadata chain. + const uint64_t next_off = footer.next_meta_header_offset; + if (next_off <= header_start_offset || next_off + sizeof(IndexFormat::MetaHeader) > buffer_pool_->file_size()) { LOG_ERROR("ParseToMapping: invalid next_meta_header_offset=%" PRIu64 " (current=%" PRIu64 ", file_size=%zu), file[%s]", - next_off, current_header_start_offset_, - buffer_pool_->file_size(), file_name_.c_str()); + next_off, header_start_offset, buffer_pool_->file_size(), + file_name_.c_str()); return IndexError_InvalidValue; } - // Bound chain count: 1024 chains @ default 1MB segment_meta_capacity - // covers >1GB of metadata, far above realistic load. + // Bound corrupted metadata chains. constexpr size_t kMaxChains = 1024; - if (chain_headers_.size() >= kMaxChains) { + if (meta_chains_.size() >= kMaxChains) { LOG_ERROR( "ParseToMapping: chain count exceeds limit %zu, file[%s] may " "be corrupted", kMaxChains, file_name_.c_str()); return IndexError_InvalidLength; } - current_header_start_offset_ = next_off; + header_start_offset = next_off; } return 0; } @@ -772,7 +1335,7 @@ class BufferStorage : public IndexStorage { //! Retrieve check point of storage uint64_t check_point(void) const override { - return footer_.check_point; + return meta_chains_.empty() ? 0 : meta_chains_.back().footer.check_point; } //! Retrieve a segment by id @@ -781,33 +1344,29 @@ class BufferStorage : public IndexStorage { mapping_shards_[mapping_shard_id()].mtx); auto seg_iter = segments_.find(id); if (seg_iter == segments_.end()) { - return WrappedSegment::Pointer{}; - } - auto id_iter = id_hash_.find(id); - if (id_iter == id_hash_.end()) { - return WrappedSegment::Pointer{}; + return {}; } return std::make_shared(this, &seg_iter->second, - id_iter->second); + &seg_iter->first); } //! Test if it a segment exists bool has(const std::string &id) const override { - return this->has_segment(id); + std::shared_lock latch( + mapping_shards_[mapping_shard_id()].mtx); + return segments_.find(id) != segments_.end(); } //! Retrieve magic number of index uint32_t magic(void) const override { - if (chain_headers_.empty()) { + if (meta_chains_.empty()) { return 0u; } - return chain_headers_.front()->magic; + return meta_chains_.front().header->magic; } protected: - //! Initialize index version segment (writes content into an IndexMapping). - //! Only intended to be called from init_index() while `mapping` is still - //! open in create-mode. + //! Write the version segment while the new mapping is open. int init_version_segment(IndexMapping &mapping) { size_t data_size = std::strlen(IndexVersion::Details()); int error_code = mapping.append(INDEX_VERSION_SEGMENT_NAME, data_size); @@ -823,16 +1382,13 @@ class BufferStorage : public IndexStorage { size_t capacity = static_cast(meta->padding_size + meta->data_size); memcpy(segment->data(), IndexVersion::Details(), data_size); segment->set_dirty(); - set_as_dirty(); meta->data_crc = ailego::Crc32c::Hash(segment->data(), data_size, 0); meta->data_size = data_size; meta->padding_size = capacity - data_size; return 0; } - //! Create the initial on-disk index structure and write the mandatory - //! version segment. Uses IndexMapping (the same engine as MMapFileStorage) - //! so the produced file is fully compatible with both storage backends. + //! Create an index compatible with the mmap storage format. int init_index(const std::string &path) { IndexMapping mapping; int ret = mapping.create(path, segment_meta_capacity_); @@ -865,20 +1421,14 @@ class BufferStorage : public IndexStorage { return index_dirty_.load(std::memory_order_relaxed); } - //! Mark the index as dirty. HOT PATH: store(true) unconditionally -- - //! a load-then-store guard could let a stale cached `true` skip the - //! store after flush_index() CAS'd dirty=false on another core, losing - //! the writer's modification. + //! Publish dirty unconditionally to avoid racing a concurrent flush. void set_as_dirty(void) { index_dirty_.store(true, std::memory_order_relaxed); } //! Refresh meta information (checksum, update time, etc.) void refresh_index(uint64_t chkp) { - // CAS-loop max: callers may invoke refresh() out of order, and the - // persisted check_point must be non-decreasing. Relaxed ordering is - // sufficient because flush_index() takes AllShardsExclusiveLatch which - // establishes the necessary happens-before for the disk write. + // Checkpoints are monotonic; the flush latch provides synchronization. if (chkp != 0) { uint64_t cur = pending_check_point_.load(std::memory_order_relaxed); while (chkp > cur) { @@ -898,19 +1448,14 @@ class BufferStorage : public IndexStorage { if (!index_dirty_.load(std::memory_order_relaxed)) { return 0; } - // Exclusive all-shards latch excludes the lock-free hot path while we - // hash meta_buf and pwrite footer; without it segments_meta_crc would - // not match the bytes on disk. + // Exclude metadata mutation while hashing and persisting it. AllShardsExclusiveLatch latch(mapping_shards_); return flush_index_locked(); } - //! PRECONDITION: caller holds AllShardsExclusiveLatch. Used by - //! flush_index() (acquires the latch) and close_index() (must flush - //! and tear down under one continuous latch hold). + //! Requires AllShardsExclusiveLatch. int flush_index_locked(void) { - // No-op on never-opened / already-closed storage: close_index() - // unconditionally calls us during teardown. + // close_index() may call this before open or after close. if (!buffer_pool_ || !buffer_pool_handle_) { index_dirty_.store(false, std::memory_order_relaxed); return 0; @@ -927,23 +1472,17 @@ class BufferStorage : public IndexStorage { index_dirty_.store(false, std::memory_order_relaxed); return 0; } - // Claim dirty atomically AT THE START so any concurrent write() that - // lands during this flush re-sets dirty=true and is picked up by the - // next flush; an unconditional store(false) at the end would silently - // swallow it. + // Claim the current dirty generation; concurrent writes start the next one. bool expected_dirty = true; if (!index_dirty_.compare_exchange_strong(expected_dirty, false, std::memory_order_relaxed)) { - // Another thread already claimed; bail out. + // Another thread already claimed it. return 0; } - // Snapshot pending_check_point_ AFTER claiming dirty: any newer chkp - // stored by a concurrent refresh_index() will be preserved by the - // CAS-reset at the end (and refresh_index() will have re-set dirty). + // Snapshot after claiming dirty so newer checkpoints survive the final CAS. const uint64_t consumed_chkp = pending_check_point_.load(std::memory_order_relaxed); - // Restore consumed_chkp on failure paths (CAS-loop max, same as - // refresh_index()) so a concurrent larger chkp wins. + // Restore on failure without overwriting a newer checkpoint. auto restore_chkp_on_failure = [this, consumed_chkp]() { if (consumed_chkp == 0) return; uint64_t cur = pending_check_point_.load(std::memory_order_relaxed); @@ -961,16 +1500,15 @@ class BufferStorage : public IndexStorage { LOG_ERROR("flush_all data blocks failed: file[%s]", file_name_.c_str()); return IndexError_WriteData; } - // Per-chain: recompute segments_meta CRC, refresh footer, pwrite both. - for (size_t ci = 0; - ci < meta_chains_.size() && ci < buffer_pool_buffers_.size(); ++ci) { + // Refresh and persist metadata for each chain. + for (size_t ci = 0; ci < meta_chains_.size(); ++ci) { MetaChain &mchain = meta_chains_[ci]; - const char *seg_buf = buffer_pool_buffers_[ci].get(); + const char *seg_buf = mchain.segment_meta.get(); mchain.footer.segments_meta_crc = - ailego::Crc32c::Hash(seg_buf, mchain.segment_meta_size, 0u); + ailego::Crc32c::Hash(seg_buf, mchain.footer.segments_meta_size, 0u); IndexFormat::UpdateMetaFooter(&mchain.footer, consumed_chkp); - if (buffer_pool_handle_->write_meta(mchain.segment_meta_file_offset, - mchain.segment_meta_size, + if (buffer_pool_handle_->write_meta(mchain.segment_meta_file_offset(), + mchain.footer.segments_meta_size, seg_buf) != 0) { LOG_ERROR("Failed to write segment meta: file[%s], chain[%zu]", file_name_.c_str(), ci); @@ -979,7 +1517,7 @@ class BufferStorage : public IndexStorage { return IndexError_WriteData; } if (buffer_pool_handle_->write_meta( - mchain.footer_file_offset, sizeof(mchain.footer), + mchain.footer_file_offset(), sizeof(mchain.footer), reinterpret_cast(&mchain.footer)) != 0) { LOG_ERROR("Failed to write footer: file[%s], chain[%zu]", file_name_.c_str(), ci); @@ -988,12 +1526,7 @@ class BufferStorage : public IndexStorage { return IndexError_WriteData; } } - if (!meta_chains_.empty()) { - footer_ = meta_chains_.back().footer; - } - // CAS-reset pending: only consume the chkp we observed. A concurrent - // larger chkp survives and will be flushed next round (refresh_index() - // also re-set dirty). + // Consume only the checkpoint observed by this flush. uint64_t expected_chkp = consumed_chkp; pending_check_point_.compare_exchange_strong(expected_chkp, 0, std::memory_order_relaxed); @@ -1002,16 +1535,12 @@ class BufferStorage : public IndexStorage { //! Close index storage void close_index(void) { - // Hold ONE continuous all-shards latch across flush + teardown so no - // writer can slip in between (which would dirty meta_buf only to have - // the page table reset under it, dropping the modification). + // Keep writers excluded across both flush and teardown. AllShardsExclusiveLatch latch(mapping_shards_); flush_index_locked(); file_name_.clear(); - id_hash_.clear(); segments_.clear(); - chain_headers_.clear(); - memset(&footer_, 0, sizeof(footer_)); + meta_chains_.clear(); { std::lock_guard tmp_latch(tmp_buffers_mutex_); for (const ArenaBlock &b : tmp_buffers_) { @@ -1021,11 +1550,7 @@ class BufferStorage : public IndexStorage { } tmp_buffers_.clear(); } - // Release every page pinned by the single-page read(const void**) - // overload (the never-released contract holds the pin until here). - // Each pin incremented the page table ref_count, so we must drop the - // matching reference before resetting the pool, otherwise - // ~VecBufferPool asserts that all blocks were released. + // Drop persistent pointer-read pins before destroying the pool. { std::lock_guard pin_latch(pinned_pages_mutex_); if (buffer_pool_handle_) { @@ -1037,20 +1562,15 @@ class BufferStorage : public IndexStorage { } buffer_pool_handle_.reset(); buffer_pool_.reset(); - max_segment_size_ = 0; - buffer_pool_buffers_.clear(); - meta_chains_.clear(); - current_header_start_offset_ = 0; + cache_enabled_ = false; pending_check_point_.store(0, std::memory_order_relaxed); index_dirty_.store(false, std::memory_order_relaxed); corrupted_.store(false, std::memory_order_relaxed); } - //! Append a segment into storage. C1: page table extends in-place; - //! latch held only briefly to protect segments_/id_hash_ insertion. + //! Append a segment into storage. int append_segment(const std::string &id, size_t size) { - // Persist any pending data_size/padding/CRC mutations from prior - // write()/resize() before we re-hash and rewrite the segment_meta. + // Persist pending metadata before re-hashing it. this->flush_index(); AllShardsExclusiveLatch latch(mapping_shards_); @@ -1077,8 +1597,7 @@ class BufferStorage : public IndexStorage { if (segments_.find(id) != segments_.end()) { return IndexError_Duplicate; } - if (meta_chains_.empty() || chain_headers_.empty() || - buffer_pool_buffers_.empty()) { + if (meta_chains_.empty()) { LOG_ERROR("append_segment: invalid state, file[%s]", file_name_.c_str()); return IndexError_Runtime; } @@ -1087,21 +1606,19 @@ class BufferStorage : public IndexStorage { const size_t page_size = ailego::kVectorPageSize; const size_t padded_size = (size + page_size - 1) / page_size * page_size; - // The current last chain owns footer_ (overwritten by ParseFooter). size_t id_size = id.length() + 1; size_t need_size = sizeof(IndexFormat::SegmentMeta) + id_size; MetaChain *chain = &meta_chains_.back(); - IndexFormat::MetaHeader *header = chain_headers_.back().get(); - char *meta_buf = buffer_pool_buffers_.back().get(); + IndexFormat::MetaHeader *header = chain->header.get(); + char *meta_buf = chain->segment_meta.get(); + IndexFormat::MetaFooter *footer = &chain->footer; - // Rollback handle for an in-memory-committed chain split. Default - // no-op; populated only after Step 1 commits, so a Step 2 failure - // can fully undo the split (otherwise an orphan empty chain would - // remain linked in the file). - std::function rollback_step1 = []() {}; + const auto footer_before_split = *footer; + const uint64_t old_footer_file_offset = chain->footer_file_offset(); + bool chain_split = false; - // ---- Step 1: chain split if current chain has no meta capacity left. - if (sizeof(IndexFormat::SegmentMeta) * footer_.segment_count + need_size > + // Step 1: split a full metadata chain. + if (sizeof(IndexFormat::SegmentMeta) * footer->segment_count + need_size > chain->segment_ids_offset) { size_t new_chain_start = buffer_pool_->file_size(); new_chain_start = @@ -1114,29 +1631,25 @@ class BufferStorage : public IndexStorage { new_meta_total - sizeof(IndexFormat::MetaHeader) - sizeof(IndexFormat::MetaFooter)); - // Stage the linked old footer without mutating footer_ yet. - const auto saved_footer_before_split = footer_; - IndexFormat::MetaFooter linked_footer = footer_; + // Stage the linked old footer without mutating the chain yet. + IndexFormat::MetaFooter linked_footer = footer_before_split; linked_footer.next_meta_header_offset = new_chain_start; IndexFormat::UpdateMetaFooter(&linked_footer, 0); if (buffer_pool_handle_->write_meta( - chain->footer_file_offset, sizeof(linked_footer), + chain->footer_file_offset(), sizeof(linked_footer), reinterpret_cast(&linked_footer)) != 0) { LOG_ERROR("append_segment: write old footer failed, file[%s]", file_name_.c_str()); return IndexError_WriteData; } - // Best-effort restore of the old footer if any subsequent write in - // this split block fails. If the restore itself fails, mark the - // storage corrupted -- on-disk old footer now points at a partial - // new chain region. - auto undo_old_footer = [this, chain, &saved_footer_before_split]() { + // Restore the old link after a failed split; reject writes if rollback + // also fails. + auto undo_old_footer = [this, chain, &footer_before_split]() { if (buffer_pool_handle_->write_meta( - chain->footer_file_offset, sizeof(saved_footer_before_split), - reinterpret_cast(&saved_footer_before_split)) != - 0) { + chain->footer_file_offset(), sizeof(footer_before_split), + reinterpret_cast(&footer_before_split)) != 0) { LOG_ERROR( "append_segment: rollback write of old footer FAILED, file[%s] " "is now in an inconsistent state -- marking storage as " @@ -1177,8 +1690,6 @@ class BufferStorage : public IndexStorage { undo_old_footer(); return IndexError_WriteData; } - uint64_t new_segment_meta_file_offset = - new_chain_start + sizeof(IndexFormat::MetaHeader); uint64_t new_footer_file_offset = new_chain_start + new_header->meta_footer_offset; if (buffer_pool_handle_->write_meta( @@ -1188,20 +1699,8 @@ class BufferStorage : public IndexStorage { return IndexError_WriteData; } - // Snapshot the OLD chain's pre-commit state for rollback_step1 - // (captured by value: `chain` is reassigned below). - const auto saved_old_chain_footer = chain->footer; - const uint64_t saved_old_footer_file_offset = chain->footer_file_offset; - const uint64_t saved_current_header_start = current_header_start_offset_; - - // Strong exception guarantee: reserve() FIRST so the three - // push_back's cannot throw mid-way and leave - // chain_headers_/buffer_pool_buffers_/meta_chains_ at mismatched - // sizes (which flush_index_locked() would silently skip while - // ParseToMapping() on next open follows the on-disk forward link). + // Reserve before committing the split so pointer updates cannot fail. try { - chain_headers_.reserve(chain_headers_.size() + 1); - buffer_pool_buffers_.reserve(buffer_pool_buffers_.size() + 1); meta_chains_.reserve(meta_chains_.size() + 1); } catch (const std::bad_alloc &) { LOG_ERROR( @@ -1212,76 +1711,54 @@ class BufferStorage : public IndexStorage { } chain = &meta_chains_.back(); chain->footer = linked_footer; // old chain keeps linked footer - chain_headers_.push_back(std::move(new_header)); - buffer_pool_buffers_.push_back(std::move(new_meta_buf)); - meta_chains_.push_back(MetaChain{ - new_chain_start, new_footer_file_offset, new_segment_meta_file_offset, - new_segments_meta_size, new_segments_meta_size, new_footer}); - footer_ = new_footer; - current_header_start_offset_ = new_chain_start; + meta_chains_.push_back(MetaChain{std::move(new_header), + std::move(new_meta_buf), new_chain_start, + new_segments_meta_size, new_footer}); chain = &meta_chains_.back(); - header = chain_headers_.back().get(); - meta_buf = buffer_pool_buffers_.back().get(); - - // Install rollback for the committed split. Captures by value so - // later reassignment of chain/header/meta_buf does not corrupt the - // closure. - rollback_step1 = [this, saved_footer_before_split, saved_old_chain_footer, - saved_old_footer_file_offset, - saved_current_header_start]() { - // 1. Drop the forward link on the old footer. If this fails the - // on-disk old footer still points at the popped new chain - // region -- mark corrupted. + header = chain->header.get(); + meta_buf = chain->segment_meta.get(); + footer = &chain->footer; + chain_split = true; + } + + auto rollback_chain_split = [&]() { + if (chain_split) { + // Unlink the new chain before dropping its in-memory state. if (buffer_pool_handle_->write_meta( - saved_old_footer_file_offset, sizeof(saved_footer_before_split), - reinterpret_cast(&saved_footer_before_split)) != - 0) { + old_footer_file_offset, sizeof(footer_before_split), + reinterpret_cast(&footer_before_split)) != 0) { LOG_ERROR( - "append_segment: rollback_step1 write of old footer FAILED, " + "append_segment: chain-split rollback write of old footer " + "FAILED, " "file[%s] is now in an inconsistent state -- marking storage " "as corrupted; further writes will be rejected.", file_name_.c_str()); corrupted_.store(true, std::memory_order_release); } - // 2. Pop the freshly-pushed new chain (releases its unique_ptrs). - if (!meta_chains_.empty()) meta_chains_.pop_back(); - if (!chain_headers_.empty()) chain_headers_.pop_back(); - if (!buffer_pool_buffers_.empty()) buffer_pool_buffers_.pop_back(); - // 3. Restore the old chain's in-memory footer (forward link cleared). - if (!meta_chains_.empty()) { - meta_chains_.back().footer = saved_old_chain_footer; - } - // 4. Restore footer_ + current_header_start_offset_. The on-disk - // file size is intentionally NOT shrunk: the orphan region is - // unreachable (step 1 cleared the link) and reusable by the - // next split via file_size() realignment. - footer_ = saved_footer_before_split; - current_header_start_offset_ = saved_current_header_start; - }; - } + meta_chains_.pop_back(); + meta_chains_.back().footer = footer_before_split; + // The unreachable file tail is reusable; shrinking is unnecessary. + } + }; - // ---- Step 2: append SegmentMeta + ID into the (possibly new) last - // chain, then persist meta_buf and footer. - uint64_t new_data_index = footer_.content_size; + // Step 2: append the segment metadata and persist the last chain. + uint64_t new_data_index = footer->content_size; uint64_t new_seg_abs_offset = chain->header_start_offset + header->content_offset + new_data_index; uint64_t new_file_size = new_seg_abs_offset + padded_size; if (new_file_size > buffer_pool_->file_size()) { if (!buffer_pool_->extend_file(new_file_size)) { + rollback_chain_split(); return IndexError_Runtime; } } - // Save mutable state for rollback if a Step 2 disk write fails. The - // meta_buf regions that get overwritten (SegmentMeta entry + ID - // string) are also snapshotted so they can be restored exactly, - // keeping CRC consistent for a later flush_index(). - const auto saved_footer = footer_; - const auto saved_chain_footer = chain->footer; + // Snapshot every overwritten metadata region for rollback. + const auto saved_footer = *footer; const auto saved_segment_ids_offset = chain->segment_ids_offset; const size_t meta_entry_off = - sizeof(IndexFormat::SegmentMeta) * footer_.segment_count; + sizeof(IndexFormat::SegmentMeta) * footer->segment_count; const uint32_t new_ids_off = chain->segment_ids_offset - static_cast(id_size); char saved_meta_entry[sizeof(IndexFormat::SegmentMeta)]; @@ -1293,7 +1770,7 @@ class BufferStorage : public IndexStorage { chain->segment_ids_offset -= static_cast(id_size); IndexFormat::SegmentMeta *new_seg = reinterpret_cast(meta_buf) + - footer_.segment_count; + footer->segment_count; new_seg->segment_id_offset = chain->segment_ids_offset; new_seg->data_index = new_data_index; new_seg->data_size = 0; @@ -1301,36 +1778,27 @@ class BufferStorage : public IndexStorage { new_seg->padding_size = padded_size; std::memcpy(meta_buf + chain->segment_ids_offset, id.c_str(), id_size); - footer_.segment_count += 1; - footer_.content_size += padded_size; - footer_.total_size += padded_size; - footer_.segments_meta_crc = - ailego::Crc32c::Hash(meta_buf, chain->segment_meta_size, 0u); - IndexFormat::UpdateMetaFooter(&footer_, 0); - chain->footer = footer_; // sync in-memory copy for flush_index - - // Rollback for Step 2: restore in-memory state AND best-effort - // rewrite the OLD segments_meta + footer back to disk. Without the - // disk rewrite, a write_meta(footer) failure (or post-write OOM) - // would tell the caller the append failed yet leave on-disk bytes - // describing the failed append -- ParseToMapping() on next open - // would surface a ghost segment with no entry in segments_/id_hash_. - // - // If the rewrite itself fails the file is unrepairable from here: - // raise corrupted_ so subsequent writers refuse to proceed. + footer->segment_count += 1; + footer->content_size += padded_size; + footer->total_size += padded_size; + footer->segments_meta_crc = + ailego::Crc32c::Hash(meta_buf, footer->segments_meta_size, 0u); + IndexFormat::UpdateMetaFooter(footer, 0); + + // Restore memory and disk together; a failed restore marks the file bad. auto rollback_step2 = [&]() { std::memcpy(meta_buf + meta_entry_off, saved_meta_entry, sizeof(IndexFormat::SegmentMeta)); std::memcpy(meta_buf + new_ids_off, saved_id_bytes.get(), id_size); - footer_ = saved_footer; - chain->footer = saved_chain_footer; + *footer = saved_footer; chain->segment_ids_offset = saved_segment_ids_offset; - const int rc_meta = buffer_pool_handle_->write_meta( - chain->segment_meta_file_offset, chain->segment_meta_size, meta_buf); + const int rc_meta = + buffer_pool_handle_->write_meta(chain->segment_meta_file_offset(), + footer->segments_meta_size, meta_buf); const int rc_footer = buffer_pool_handle_->write_meta( - chain->footer_file_offset, sizeof(footer_), - reinterpret_cast(&footer_)); + chain->footer_file_offset(), sizeof(*footer), + reinterpret_cast(footer)); if (rc_meta != 0 || rc_footer != 0) { LOG_ERROR( "append_segment: rollback_step2 disk rewrite FAILED " @@ -1342,108 +1810,83 @@ class BufferStorage : public IndexStorage { } }; - if (buffer_pool_handle_->write_meta(chain->segment_meta_file_offset, - chain->segment_meta_size, + if (buffer_pool_handle_->write_meta(chain->segment_meta_file_offset(), + footer->segments_meta_size, meta_buf) != 0) { LOG_ERROR("append_segment: write segment_meta failed, file[%s]", file_name_.c_str()); rollback_step2(); - rollback_step1(); + rollback_chain_split(); return IndexError_WriteData; } if (buffer_pool_handle_->write_meta( - chain->footer_file_offset, sizeof(footer_), - reinterpret_cast(&footer_)) != 0) { + chain->footer_file_offset(), sizeof(*footer), + reinterpret_cast(footer)) != 0) { LOG_ERROR("append_segment: write footer failed, file[%s]", file_name_.c_str()); rollback_step2(); - rollback_step1(); + rollback_chain_split(); return IndexError_WriteData; } - // Strong exception guarantee for the in-memory commit: emplace into - // segments_ and id_hash_ as one transactional unit -- if id_hash_ - // throws after segments_ succeeded, undo segments_ before - // propagating. unordered_map::emplace() leaves existing element - // addresses stable, so WrappedSegment instances pointing into - // segments_ remain valid. - auto seg_ins = segments_.end(); - bool seg_inserted = false; + // Publish the segment only after its metadata is durable. try { auto ins = segments_.emplace( id, IndexMapping::SegmentInfo{IndexMapping::Segment{new_seg}, chain->header_start_offset, header}); if (!ins.second) { - // Cannot happen under the exclusive latch we hold (find() above - // checked), but be defensive. + // Defensive: the exclusive latch should make this unreachable. LOG_ERROR( "append_segment: duplicate id appeared after commit, file[%s], " "id[%s]", file_name_.c_str(), id.c_str()); rollback_step2(); - rollback_step1(); + rollback_chain_split(); return IndexError_Duplicate; } - seg_ins = ins.first; - seg_inserted = true; - const size_t new_id = id_hash_.size(); - id_hash_.emplace(id, new_id); } catch (const std::bad_alloc &) { LOG_ERROR( "append_segment: in-memory commit OOM, rolling back, file[%s], " "id[%s]", file_name_.c_str(), id.c_str()); - if (seg_inserted) { - segments_.erase(seg_ins); - } rollback_step2(); - rollback_step1(); + rollback_chain_split(); return IndexError_Runtime; } - max_segment_size_ = std::max(max_segment_size_, padded_size); - // C1: extend_file() already extended the page table in-place; no pool - // rotation or flush_all needed. + // extend_file() already grew the page table in place. return 0; } - //! Test if a segment exists - bool has_segment(const std::string &id) const { - std::shared_lock latch( - mapping_shards_[mapping_shard_id()].mtx); - return (segments_.find(id) != segments_.end()); + private: + bool read_range(size_t offset, size_t len, char *out) const { + return cache_enabled_ + ? buffer_pool_handle_->read_range(offset, len, out) + : buffer_pool_handle_->read_range_bypass(offset, len, out); } - private: std::atomic index_dirty_{false}; std::atomic pending_check_point_{0}; - // Set when an append_segment() rollback fails to restore on-disk state. - // Once set, all writers (write/append_segment/flush_index_locked) refuse - // to proceed. Only ever raised; cleared only by close_index(). + // Raised after an unrecoverable rollback; blocks writes until close. std::atomic corrupted_{false}; - // Sharded reader-writer lock: each reader hashes to its own shard to - // avoid cache-line ping-pong on the reader counter; writers lock all - // shards. + // Readers hash across shards; metadata writers lock every shard. static constexpr size_t kMappingMutexShards = 32; struct alignas(64) MutexShard { std::shared_mutex mtx; }; mutable MutexShard mapping_shards_[kMappingMutexShards]{}; - // Per-(thread, instance) shard selection. Combining thread::id with - // `this` ensures two BufferStorage instances on the same thread map to - // different shards (a thread_local-only id collapses them onto one - // shard). boost-style hash_combine disperses skewed thread::id - // distributions across the 32 shards. + // Mix the thread and instance identities to distribute reader latches. size_t mapping_shard_id() const { - size_t seed = std::hash()(std::this_thread::get_id()); + static thread_local const size_t thread_seed = + std::hash()(std::this_thread::get_id()); + size_t seed = thread_seed; size_t inst = std::hash()(static_cast(this)); - // boost::hash_combine(seed, inst) seed ^= inst + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); return seed % kMappingMutexShards; } - // RAII guard that locks ALL shards exclusively (for writers). + // Exclusive guard for all metadata shards. struct AllShardsExclusiveLatch { MutexShard *shards_; AllShardsExclusiveLatch(MutexShard *shards) : shards_(shards) { @@ -1457,29 +1900,14 @@ class BufferStorage : public IndexStorage { delete; }; - // Arena slab for cross-page temp buffers handed out by - // WrappedSegment::read(const void**). The legacy contract requires - // every returned pointer to stay valid until close_index(), so slots - // are never freed individually -- they are carved out of large - // 4K-aligned arenas which are released in bulk. - // - // Why an arena instead of one posix_memalign(4K, 4K) per read: - // Android Bionic scudo's small-class chunk pool is prone to large- - // alignment starvation under fragmentation (we observed sporadic - // posix_memalign(4096, 4096) returning ENOMEM even with plenty of - // free memory). A single large request (>= kArenaSize) is served - // from scudo's secondary allocator (mmap-backed), which is reliable - // up to the true OOM boundary. + // Aligned arenas retain legacy pointer-read results until close_index(). struct ArenaBlock { char *base{nullptr}; size_t size{0}; // Total bytes in this arena (4K-aligned). size_t used{0}; // Bytes already handed out (4K-aligned). }; - // Caller MUST hold tmp_buffers_mutex_. alloc_size MUST be a - // multiple of 4096. Returns nullptr only if scudo cannot satisfy a - // fresh arena allocation, i.e. effectively true OOM. + // Requires tmp_buffers_mutex_; alloc_size must be 4K-aligned. char *tmp_arena_alloc_locked(size_t alloc_size) { - static constexpr size_t kAlign = 4096UL; static constexpr size_t kArenaSize = 1UL << 20; // 1 MiB if (!tmp_buffers_.empty()) { ArenaBlock &back = tmp_buffers_.back(); @@ -1490,7 +1918,8 @@ class BufferStorage : public IndexStorage { } } size_t new_size = alloc_size > kArenaSize ? alloc_size : kArenaSize; - char *p = static_cast(ailego_aligned_malloc(new_size, kAlign)); + char *p = + static_cast(ailego_aligned_malloc(new_size, kBufferAlignment)); if (!p) { return nullptr; } @@ -1500,49 +1929,39 @@ class BufferStorage : public IndexStorage { std::vector tmp_buffers_{}; mutable std::mutex tmp_buffers_mutex_{}; - // Page ids pinned by the single-page read(const void**) overload, which - // keeps the pin alive until close_index() (the never-released contract). - // Each pin increments the page table ref_count, so close_index() must - // release every recorded pin before tearing down the pool; otherwise - // ~VecBufferPool fires its "all blocks released" assertion. Duplicates are - // allowed: repeated single-page reads of the same page each take a pin and - // therefore each need a matching release. + // One entry per legacy pointer-read pin; duplicates are intentional. std::vector pinned_pages_{}; mutable std::mutex pinned_pages_mutex_{}; // buffer manager std::string file_name_; - // Per-chain owning copies of MetaHeader. segments_[name].segment_header - // points into one of these; using a single shared header_ would let the - // next chain's ParseHeader overwrite earlier-chain content_offset. - std::vector> chain_headers_{}; - IndexFormat::MetaFooter footer_{}; std::unordered_map segments_{}; - std::unordered_map id_hash_{}; - uint64_t max_segment_size_{0}; - std::vector> buffer_pool_buffers_{}; ailego::VecBufferPool::Pointer buffer_pool_{nullptr}; + bool cache_enabled_{false}; ailego::VecBufferPoolHandle::Pointer buffer_pool_handle_{nullptr}; - uint64_t current_header_start_offset_{0u}; - // Capacity (in bytes) of the segment metadata section written by - // init_index(). + // Segment metadata capacity written by init_index(). uint32_t segment_meta_capacity_{4096u}; - // Per-header-chain file offsets used by flush_index() and append_segment(). + // Per-chain ownership and state used by parsing, flushing, and appending. struct MetaChain { + // Pointees stay stable across vector moves and are referenced by segments_. + std::unique_ptr header; + std::unique_ptr segment_meta; uint64_t header_start_offset; - uint64_t footer_file_offset; - uint64_t segment_meta_file_offset; - uint32_t segment_meta_size; - // Lowest segment-ID-string offset within segment_meta; equals - // segment_meta_size when empty, decreases by strlen(id)+1 per append. - // Used to detect when a chain split is needed. + // Lowest ID-string offset; detects metadata exhaustion. uint32_t segment_ids_offset; - // In-memory copy of this chain's MetaFooter, kept in sync with disk by - // flush_index() and append_segment() to avoid a pread per chain. + // Authoritative in-memory footer for this chain. IndexFormat::MetaFooter footer; + + uint64_t footer_file_offset() const { + return header_start_offset + header->meta_footer_offset; + } + + uint64_t segment_meta_file_offset() const { + return footer_file_offset() - footer.segments_meta_size; + } }; std::vector meta_chains_{}; }; diff --git a/src/core/utility/utility_params.h b/src/core/utility/utility_params.h index c57e6e980..4fc8b7c0a 100644 --- a/src/core/utility/utility_params.h +++ b/src/core/utility/utility_params.h @@ -60,6 +60,18 @@ static const std::string MMAPFILE_READ_STORAGE_HEADER_OFFSET = static const std::string MMAPFILE_READ_STORAGE_FOOTER_OFFSET = "proxima.mmap_file.container.footer_offset"; +//! BufferReadStorage (read-only storage backed by VecBufferPool) +static const std::string BUFFER_READ_STORAGE_CHECKSUM_VALIDATION = + "proxima.buffer.read_storage.checksum_validation"; +static const std::string BUFFER_READ_STORAGE_HEADER_OFFSET = + "proxima.buffer.read_storage.header_offset"; +static const std::string BUFFER_READ_STORAGE_FOOTER_OFFSET = + "proxima.buffer.read_storage.footer_offset"; +static const std::string BUFFER_READ_STORAGE_WARMUP_MODE = + "proxima.buffer.read_storage.warmup_mode"; +static const std::string BUFFER_READ_STORAGE_WARMUP_NONE = "none"; +static const std::string BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL = "sequential"; + //! MMapFileStorage static const std::string MMAPFILE_STORAGE_MEMORY_LOCKED = "proxima.mmap_file.storage.memory_locked"; @@ -231,4 +243,4 @@ static const std::string PARAM_FORWARD_MULTI_COUNT = "proxima.param.forward.multi_count"; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/utility/visit_filter.h b/src/core/utility/visit_filter.h index 8874a949a..b55f4baba 100644 --- a/src/core/utility/visit_filter.h +++ b/src/core/utility/visit_filter.h @@ -144,7 +144,8 @@ class VisitBloomFilter { c->filter = new (std::nothrow) ailego::BloomFilter(maxScanNum, p); if (c->filter == nullptr) { - LOG_ERROR("New BloomFilter failed, reuse old one"); + LOG_ERROR("New BloomFilter failed"); + delete c; return IndexError_NoMemory; } genRandomHashBits(c); @@ -368,6 +369,9 @@ class VisitFilter { }; VisitFilter() : mode_(0), ctx_(nullptr) {}; + ~VisitFilter() { + destroy(); + } inline bool visited(id_t idx) { PROXIMA_HNSW_VISITFILTER_CALL_IMPL(visited, idx); @@ -397,18 +401,60 @@ class VisitFilter { return true; } - inline void destroy() { - if (ctx_ != nullptr) { - PROXIMA_HNSW_VISITFILTER_CALL_IMPL(destroy); + inline void destroy() noexcept { + void *ctx = ctx_; + const int mode = mode_; + ctx_ = nullptr; + mode_ = Default; + if (ctx == nullptr) { + return; + } + + switch (mode) { + case BloomFilter: + VisitBloomFilter::destroy( + static_cast(ctx)); + return; + case BitMap: + VisitBitMap::destroy(static_cast(ctx)); + return; + case ByteMap: + VisitByteMap::destroy(static_cast(ctx)); + return; + default: + return; } } int init(int mode, uint64_t maxDocCnt, uint64_t maxScanNum, float negativeProbability) { + destroy(); mode_ = mode; - PROXIMA_HNSW_VISITFILTER_CALL_IMPL(init, &ctx_, maxDocCnt, maxScanNum, - std::make_tuple(negativeProbability)); - return 0; // place holder + int ret = IndexError_InvalidArgument; + switch (mode_) { + case BloomFilter: + ret = VisitBloomFilter::init( + static_cast(nullptr), &ctx_, maxDocCnt, + maxScanNum, std::make_tuple(negativeProbability)); + break; + case BitMap: + ret = VisitBitMap::init(static_cast(nullptr), + &ctx_, maxDocCnt, maxScanNum, + std::make_tuple(negativeProbability)); + break; + case ByteMap: + ret = VisitByteMap::init(static_cast(nullptr), + &ctx_, maxDocCnt, maxScanNum, + std::make_tuple(negativeProbability)); + break; + default: + break; + } + if (ret != 0) { + ctx_ = nullptr; + mode_ = Default; + } + return ret; } int get_mode(void) const { diff --git a/src/db/collection.cc b/src/db/collection.cc index d6900dd31..cd2c8853c 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -989,10 +989,10 @@ std::vector CollectionImpl::build_compact_task( if (current_actual_doc_count + actual_doc_count > max_doc_count_per_segment) { // only create SegmentCompactTask when rebuild=true - task = SegmentTask::CreateCompactTask( - CompactTask{path_, schema, current_group, - allocate_segment_id_for_tmp_segment(), filter, - !options_.enable_mmap_, concurrency}); + task = SegmentTask::CreateCompactTask(CompactTask{ + path_, schema, current_group, + allocate_segment_id_for_tmp_segment(), filter, + !options_.enable_mmap_, options_.enable_mmap_, concurrency}); } } else { if (current_doc_count + doc_count > max_doc_count_per_segment) { @@ -1003,10 +1003,10 @@ std::vector CollectionImpl::build_compact_task( current_group[0], "", nullptr, concurrency}); skip_task = current_group[0]->all_vector_index_ready(); } else { - task = SegmentTask::CreateCompactTask( - CompactTask{path_, schema, current_group, - allocate_segment_id_for_tmp_segment(), nullptr, - !options_.enable_mmap_, concurrency}); + task = SegmentTask::CreateCompactTask(CompactTask{ + path_, schema, current_group, + allocate_segment_id_for_tmp_segment(), nullptr, + !options_.enable_mmap_, options_.enable_mmap_, concurrency}); } } } @@ -1034,7 +1034,8 @@ std::vector CollectionImpl::build_compact_task( } else { task = SegmentTask::CreateCompactTask(CompactTask{ path_, schema, current_group, allocate_segment_id_for_tmp_segment(), - rebuild ? filter : nullptr, !options_.enable_mmap_, concurrency}); + rebuild ? filter : nullptr, !options_.enable_mmap_, + options_.enable_mmap_, concurrency}); } tasks.push_back(task); } @@ -1744,7 +1745,12 @@ Result CollectionImpl::Query(const MultiQuery &query) const { // Single-segment queries have no segment-level fanout; multi-segment queries // already use the query pool per sub-query. if (segments.size() == 1) { - auto group = GlobalResource::Instance().query_thread_pool()->make_group(); + auto *pool = GlobalResource::Instance().query_thread_pool(); + if (pool == nullptr) { + return tl::make_unexpected( + Status::InternalError("Query thread pool initialization failed")); + } + auto group = pool->make_group(); for (size_t i = 0; i < pending_queries.size(); ++i) { group->execute( [&, i]() { results[i] = execute_query(pending_queries[i]); }); diff --git a/src/db/common/config.cc b/src/db/common/config.cc index 14b53a6ae..4a63b4e3d 100644 --- a/src/db/common/config.cc +++ b/src/db/common/config.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include @@ -87,9 +88,17 @@ Status GlobalConfig::Validate(const ConfigData &config) const { } // Validate log configuration - if (config.log_config->GetLoggerType() == FILE_LOG_TYPE_NAME) { + if (!config.log_config) { + return Status::InvalidArgument("log_config cannot be null"); + } + const std::string logger_type = config.log_config->GetLoggerType(); + if (logger_type == FILE_LOG_TYPE_NAME) { auto log_config = std::dynamic_pointer_cast(config.log_config); + if (!log_config) { + return Status::InvalidArgument( + "AppendLogger configuration must use FileLogConfig"); + } // Validate file log specific configurations if (log_config->dir.empty()) { @@ -112,57 +121,167 @@ Status GlobalConfig::Validate(const ConfigData &config) const { return Status::InvalidArgument( "log_overdue_days must be greater than 0 when set to FileLogger"); } + } else if (logger_type == CONSOLE_LOG_TYPE_NAME) { + if (!std::dynamic_pointer_cast(config.log_config)) { + return Status::InvalidArgument( + "ConsoleLogger configuration must use ConsoleLogConfig"); + } + } else { + return Status::InvalidArgument("unsupported logger type: ", logger_type); } return Status::OK(); } Status GlobalConfig::Initialize(const ConfigData &config) { - // Use atomic compare-exchange to ensure only one initialization - bool expected = false; - if (!initialized_.compare_exchange_strong(expected, true)) { - return Status::OK(); + { + std::unique_lock lock(initialization_mutex_); + initialization_cv_.wait(lock, [this] { + return initialization_state_ != InitializationState::kInitializing; + }); + if (initialization_state_ == InitializationState::kInitialized || + initialization_state_ == InitializationState::kFailed) { + return initialization_status_; + } + initialization_state_ = InitializationState::kInitializing; } - auto s = Validate(config); - CHECK_RETURN_STATUS(s); + Status result; + bool validation_failed = false; + try { + result = Validate(config); + validation_failed = !result.ok(); - // Preserve the SDK-set jieba_dict_dir when caller didn't specify one. - // Lock spans the bulk assign so readers never see a half-written string. - { - std::lock_guard lk(mutex_); - std::string final_jieba = config.jieba_dict_dir.empty() - ? config_.jieba_dict_dir - : config.jieba_dict_dir; - config_ = config; - config_.jieba_dict_dir = std::move(final_jieba); - } + std::unique_lock config_write_lock; + ConfigData effective_config; + std::shared_ptr published_config; + if (result.ok()) { + // Serialize snapshot writers through the whole one-time initialization. + // A concurrent SDK jieba setter then runs either before this snapshot is + // prepared or after it is published, so an empty Initialize() value can + // never overwrite a newer setter update. + config_write_lock = std::unique_lock(mutex_); + // Preserve the SDK-set jieba_dict_dir when caller didn't specify one. + // Prepare a complete snapshot locally, but do not publish it until all + // side-effecting initialization stages have succeeded. + effective_config = config; + if (effective_config.jieba_dict_dir.empty()) { + effective_config.jieba_dict_dir = config_snapshot()->jieba_dict_dir; + } + // The public ConfigData remains mutable to preserve API compatibility. + // Clone the built-in logger configurations so caller-side mutations + // after Initialize() cannot silently modify the published snapshot. + if (auto file = std::dynamic_pointer_cast( + effective_config.log_config)) { + effective_config.log_config = std::make_shared(*file); + } else if (auto console = std::dynamic_pointer_cast( + effective_config.log_config)) { + effective_config.log_config = + std::make_shared(*console); + } + // Allocate the immutable publication object before logs, thread pools or + // the memory pool are changed. The final atomic store cannot fail. + published_config = std::make_shared(effective_config); + } - s = LogUtil::Init(log_dir(), log_file_basename(), int(log_level()), - log_type(), log_file_size(), log_overdue_days()); - CHECK_RETURN_STATUS(s); + if (result.ok()) { + static const bool exit_handler_registered = + std::atexit(ExitLogHandler) == 0; + if (!exit_handler_registered) { + std::cerr << "Failed to register exit handler" << std::endl; + result = Status::InternalError("Failed to register exit handler"); + } + } - if (std::atexit(ExitLogHandler) != 0) { - std::cerr << "Failed to register exit handler" << std::endl; - return Status::InternalError("Failed to register exit handler"); + bool log_setup_attempted = false; + bool log_initialized = false; + if (result.ok()) { + Status log_status; + const int resource_result = + GlobalResource::Instance().initialize_with_setup( + effective_config.memory_limit_bytes, + effective_config.query_thread_count, + effective_config.query_thread_binding, + effective_config.optimize_thread_count, + effective_config.optimize_thread_binding, [&] { + log_setup_attempted = true; + const auto *file_config = dynamic_cast( + effective_config.log_config.get()); + static const std::string empty; + log_status = + LogUtil::Init(file_config ? file_config->dir : empty, + file_config ? file_config->basename : empty, + int(effective_config.log_config->level), + effective_config.log_config->GetLoggerType(), + file_config ? file_config->file_size : 0, + file_config ? file_config->overdue_days : 0); + log_initialized = log_status.ok(); + return log_initialized ? 0 : -1; + }); + if (resource_result != 0 && log_setup_attempted) { + // LogUtil::Init may have made partial progress even when it returns an + // error, so normalize every failed setup/resource transaction back to + // an uninitialized logger. + LogUtil::Shutdown(); + } + if (!log_status.ok()) { + result = log_status; + } else if (resource_result != 0) { + // A predictable configuration mismatch is rejected before the setup + // callback. If a later memory-pool stage fails, undo the newly-created + // logger so Initialize() does not leave a half-published subsystem. + result = Status::InternalError( + "Failed to initialize the process-wide global resources"); + } + } + + if (result.ok()) { + auto old_config = config_snapshot(); + if (old_config->log_config != published_config->log_config) { + retained_log_config_ = old_config->log_config; + } + std::atomic_store_explicit(&config_, std::move(published_config), + std::memory_order_release); + } + } catch (const std::exception &e) { + result = Status::InternalError( + "Global configuration initialization threw: ", e.what()); + } catch (...) { + result = Status::InternalError( + "Global configuration initialization threw an unknown exception"); } - GlobalResource::Instance().initialize(); - return Status::OK(); + { + std::lock_guard lock(initialization_mutex_); + initialization_status_ = result; + // Invalid input has no side effects and may be corrected by a later call. + // A later failure may already have registered the harmless process-exit + // callback, so retain and replay that terminal error rather than retrying + // a one-time process initialization with ambiguous ownership. + initialization_state_ = + validation_failed ? InitializationState::kUninitialized + : (result.ok() ? InitializationState::kInitialized + : InitializationState::kFailed); + } + initialization_cv_.notify_all(); + return result; } void GlobalConfig::set_default_jieba_dict_dir(const std::string &dir) { std::lock_guard lk(mutex_); - config_.jieba_dict_dir = dir; + auto next = std::make_shared(*config_snapshot()); + next->jieba_dict_dir = dir; + std::atomic_store_explicit(&config_, + std::shared_ptr(std::move(next)), + std::memory_order_release); } std::string GlobalConfig::jieba_dict_dir() const { - std::lock_guard lk(mutex_); - return config_.jieba_dict_dir; + return config_snapshot()->jieba_dict_dir; } uint64_t GlobalConfig::memory_limit_bytes() const noexcept { - return config_.memory_limit_bytes; + return config_snapshot()->memory_limit_bytes; } FACTORY_REGISTER_LOGGER(AppendLogger); diff --git a/src/db/common/global_resource.cc b/src/db/common/global_resource.cc index 62d1966cd..c4026de86 100644 --- a/src/db/common/global_resource.cc +++ b/src/db/common/global_resource.cc @@ -12,24 +12,108 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "db/common/global_resource.h" +#include #include #include +#include #include namespace zvec { -void GlobalResource::initialize() { - static std::once_flag flag; - std::call_once(flag, [this]() { - this->query_thread_pool_.reset(new ailego::ThreadPool( - GlobalConfig::Instance().query_thread_count(), - GlobalConfig::Instance().query_thread_binding())); - this->optimize_thread_pool_.reset(new ailego::ThreadPool( - GlobalConfig::Instance().optimize_thread_count(), - GlobalConfig::Instance().optimize_thread_binding())); - zvec::ailego::MemoryLimitPool::get_instance().init( - GlobalConfig::Instance().memory_limit_bytes()); - }); +int GlobalResource::initialize() { + const auto &config = GlobalConfig::Instance(); + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + // Standalone/core users may configure the process-wide pool before the DB + // layer lazily creates its thread pools. In that case the first published + // pool budget remains authoritative; a lazy accessor must not try to resize + // a live cache merely because GlobalConfig still contains its default. + const uint64_t effective_memory_limit = memory_pool.initialized() + ? memory_pool.capacity() + : config.memory_limit_bytes(); + return initialize(effective_memory_limit, config.query_thread_count(), + config.query_thread_binding(), + config.optimize_thread_count(), + config.optimize_thread_binding()); +} + +int GlobalResource::initialize(uint64_t memory_limit_bytes, + uint32_t query_thread_count, + bool query_thread_binding, + uint32_t optimize_thread_count, + bool optimize_thread_binding) { + return initialize_with_setup(memory_limit_bytes, query_thread_count, + query_thread_binding, optimize_thread_count, + optimize_thread_binding, {}); +} + +int GlobalResource::initialize_with_setup(uint64_t memory_limit_bytes, + uint32_t query_thread_count, + bool query_thread_binding, + uint32_t optimize_thread_count, + bool optimize_thread_binding, + const std::function &setup) { + std::lock_guard lock(initialization_mutex_); + try { + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + if (query_thread_pool_ && optimize_thread_pool_) { + if (memory_limit_bytes_ == memory_limit_bytes && + query_thread_count_ == query_thread_count && + optimize_thread_count_ == optimize_thread_count && + query_thread_binding_ == query_thread_binding && + optimize_thread_binding_ == optimize_thread_binding) { + return setup ? setup() : 0; + } + LOG_ERROR( + "GlobalResource::initialize rejected configuration change after " + "thread pools were created"); + return -1; + } + + // An explicit GlobalConfig initialization must never publish a memory + // limit different from an already configured lower-level pool. Reject it + // before logger setup or thread-pool publication; lazy initialize() above + // deliberately passes the existing capacity and therefore remains + // compatible with standalone/core callers. + if (setup && memory_pool.initialized() && + memory_pool.capacity() != memory_limit_bytes) { + LOG_ERROR( + "GlobalResource::initialize rejected memory limit change after " + "MemoryLimitPool was configured: requested_capacity=%llu " + "current_capacity=%llu", + static_cast(memory_limit_bytes), + static_cast(memory_pool.capacity())); + return -1; + } + + auto query_thread_pool = std::make_unique( + query_thread_count, query_thread_binding); + auto optimize_thread_pool = std::make_unique( + optimize_thread_count, optimize_thread_binding); + // Run side-effecting setup only after every fallible resource allocation + // and compatibility check that can be performed without mutating global + // state. Holding initialization_mutex_ closes the race with lazy callers + // attempting to initialize a different resource configuration. + if (setup && setup() != 0) { + return -1; + } + if (memory_pool.init(memory_limit_bytes) != 0) { + return -1; + } + + memory_limit_bytes_ = memory_limit_bytes; + query_thread_count_ = query_thread_count; + optimize_thread_count_ = optimize_thread_count; + query_thread_binding_ = query_thread_binding; + optimize_thread_binding_ = optimize_thread_binding; + this->query_thread_pool_ = std::move(query_thread_pool); + this->optimize_thread_pool_ = std::move(optimize_thread_pool); + return 0; + } catch (const std::exception &e) { + LOG_ERROR("GlobalResource::initialize failed: %s", e.what()); + } catch (...) { + LOG_ERROR("GlobalResource::initialize failed with an unknown exception"); + } + return -1; } } // namespace zvec diff --git a/src/db/common/global_resource.h b/src/db/common/global_resource.h index 816e45080..a04631cd6 100644 --- a/src/db/common/global_resource.h +++ b/src/db/common/global_resource.h @@ -13,27 +13,54 @@ // limitations under the License. #pragma once +#include +#include #include +#include #include #include namespace zvec { +class GlobalConfig; + class GlobalResource : public ailego::Singleton { public: - void initialize(); + int initialize(); ailego::ThreadPool *query_thread_pool() { - initialize(); + if (initialize() != 0) { + return nullptr; + } return query_thread_pool_.get(); } ailego::ThreadPool *optimize_thread_pool() { - initialize(); + if (initialize() != 0) { + return nullptr; + } return optimize_thread_pool_.get(); } private: + friend class GlobalConfig; + + int initialize(uint64_t memory_limit_bytes, uint32_t query_thread_count, + bool query_thread_binding, uint32_t optimize_thread_count, + bool optimize_thread_binding); + int initialize_with_setup(uint64_t memory_limit_bytes, + uint32_t query_thread_count, + bool query_thread_binding, + uint32_t optimize_thread_count, + bool optimize_thread_binding, + const std::function &setup); + + std::mutex initialization_mutex_; + uint64_t memory_limit_bytes_{0}; + uint32_t query_thread_count_{0}; + uint32_t optimize_thread_count_{0}; + bool query_thread_binding_{false}; + bool optimize_thread_binding_{false}; std::unique_ptr query_thread_pool_; std::unique_ptr optimize_thread_pool_; }; diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 36a3b5e9b..1c3ba49fd 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -287,6 +287,9 @@ class SegmentImpl : public Segment, const std::string &index_file_path, const std::string &column, const FieldSchema &field, int concurrency); + Status reopen_vector_indexer_for_serving( + const VectorColumnIndexer::Ptr &vector_indexer); + // Helper functions for Insert/Update/Upsert/Delete template Status InsertScalar(InvertedColumnIndexer::Ptr &indexer, const Doc &doc, @@ -1660,6 +1663,10 @@ Result SegmentImpl::merge_vector_indexer( vector_column_params::MergeOptions merge_options; if (concurrency == 0) { merge_options.pool = GlobalResource::Instance().optimize_thread_pool(); + if (merge_options.pool == nullptr) { + return tl::make_unexpected( + Status::InternalError("Optimize thread pool initialization failed")); + } merge_options.write_concurrency = static_cast(merge_options.pool->count()); } else { @@ -1670,9 +1677,27 @@ Result SegmentImpl::merge_vector_indexer( s = vector_indexer->Flush(); CHECK_RETURN_STATUS_EXPECTED(s); + s = reopen_vector_indexer_for_serving(vector_indexer); + CHECK_RETURN_STATUS_EXPECTED(s); + return vector_indexer; } +Status SegmentImpl::reopen_vector_indexer_for_serving( + const VectorColumnIndexer::Ptr &vector_indexer) { + if (options_.enable_mmap_) { + return Status::OK(); + } + if (vector_indexer == nullptr) { + return Status::InvalidArgument("Vector indexer is null"); + } + + auto s = vector_indexer->Close(); + CHECK_RETURN_STATUS(s); + return vector_indexer->Open( + vector_column_params::ReadOptions{false, false, true}); +} + Status SegmentImpl::create_vector_index( const std::string &column, const IndexParams::Ptr &index_params, int concurrency, SegmentMeta::Ptr *segment_meta, @@ -1867,6 +1892,9 @@ Status SegmentImpl::drop_vector_index( s = new_vector_indexer->Flush(); CHECK_RETURN_STATUS(s); + s = reopen_vector_indexer_for_serving(new_vector_indexer); + CHECK_RETURN_STATUS(s); + (*vector_indexers)[column] = new_vector_indexer; new_segment_meta->remove_vector_persisted_block( column, vector_index_params->quantize_type() != QuantizeType::UNDEFINED); @@ -3452,11 +3480,6 @@ Status SegmentImpl::alter_column(const std::string &column_name, persist_stores_.erase(persist_stores_.begin() + local_idx); } - if (!options_.enable_mmap_) { - zvec::ailego::MemoryLimitPool::get_instance().init( - GlobalConfig::Instance().memory_limit_bytes()); - } - // delete single column store file for (auto block_id : will_del_block_ids) { // delete forward store file @@ -3547,11 +3570,6 @@ Status SegmentImpl::drop_column(const std::string &column_name) { persist_stores_.erase(persist_stores_.begin() + local_idx); } - if (!options_.enable_mmap_) { - zvec::ailego::MemoryLimitPool::get_instance().init( - GlobalConfig::Instance().memory_limit_bytes()); - } - // delete single column store file for (auto block_id : will_del_block_ids) { // delete forward store file @@ -4112,7 +4130,7 @@ VectorColumnIndexer::Ptr SegmentImpl::create_vector_indexer( auto vector_indexer = std::make_shared(index_file_path, field); - vector_column_params::ReadOptions options{true, true}; + vector_column_params::ReadOptions options{options_.enable_mmap_, true}; auto status = vector_indexer->Open(options); if (!status.ok()) { LOG_ERROR("Failed to open vector indexer for field: %s, err: %s", @@ -4395,6 +4413,8 @@ Status SegmentImpl::finish_memory_components() { // remove indexer from memory to persist for (auto &[column_name, indexer] : memory_vector_indexers_) { + s = reopen_vector_indexer_for_serving(indexer); + CHECK_RETURN_STATUS(s); auto block_id = memory_vector_block_ids_[column_name]; BlockMeta vb = BlockMeta{block_id, BlockType::VECTOR_INDEX, block.min_doc_id_, @@ -4411,6 +4431,8 @@ Status SegmentImpl::finish_memory_components() { // remove quant indexer from memory to persist for (auto &[column_name, indexer] : quant_memory_vector_indexers_) { + s = reopen_vector_indexer_for_serving(indexer); + CHECK_RETURN_STATUS(s); auto block_id = quant_memory_vector_block_ids_[column_name]; BlockMeta block_meta(block_id, BlockType::VECTOR_INDEX_QUANTIZE, block.min_doc_id_, block.max_doc_id_, block.doc_count_, diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index f1d588111..4d6cb92dd 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -150,7 +150,8 @@ Status SegmentHelper::ExecuteCompactTask(CompactTask &task) { s = ReduceVectorIndex(schema, input_segments, output_segment_path, row_id_filter, block_id_generator, min_doc_id, - max_doc_id, doc_count, task.concurrency_, &block_metas); + max_doc_id, doc_count, task.enable_mmap_, + task.concurrency_, &block_metas); CHECK_RETURN_STATUS(s); LOG_INFO("Compacted vector index"); @@ -640,7 +641,7 @@ Status SegmentHelper::ReduceVectorIndex( const std::vector &input_segments, const std::string &output_segment_path, const IndexFilter::Ptr &filter, std::function &block_id_generator, uint64_t min_doc_id, - uint64_t max_doc_id, uint32_t doc_count, int concurrency, + uint64_t max_doc_id, uint32_t doc_count, bool enable_mmap, int concurrency, std::vector *output_block_metas) { Status s; @@ -671,7 +672,7 @@ Status SegmentHelper::ReduceVectorIndex( s = MergeWithOptionalReuse( vector_index_path, *field, collect_merge_indexers(&Segment::get_vector_indexer), filter, - concurrency, nullptr); + enable_mmap, concurrency, nullptr); CHECK_RETURN_STATUS(s); BlockMeta new_block_meta; @@ -695,7 +696,7 @@ Status SegmentHelper::ReduceVectorIndex( s = MergeWithOptionalReuse( vector_index_path, *field_without_quantize, collect_merge_indexers(&Segment::get_vector_indexer), filter, - concurrency, &vector_indexer); + enable_mmap, concurrency, &vector_indexer); CHECK_RETURN_STATUS(s); // The training step (for RABITQ) and the subsequent quantize merge both @@ -735,8 +736,8 @@ Status SegmentHelper::ReduceVectorIndex( : collect_merge_indexers(&Segment::get_quant_vector_indexer); s = MergeWithOptionalReuse(vector_quan_index_path, *field_for_quantize, - quant_merge_sources, filter, concurrency, - nullptr); + quant_merge_sources, filter, enable_mmap, + concurrency, nullptr); CHECK_RETURN_STATUS(s); s = vector_indexer->Close(); @@ -791,11 +792,15 @@ bool CanReuseFirstIndexer(const std::vector &indexers, Status SegmentHelper::MergeWithOptionalReuse( const std::string &output_index_path, const FieldSchema &index_field, std::vector source_indexers, - const IndexFilter::Ptr &filter, int concurrency, + const IndexFilter::Ptr &filter, bool enable_mmap, int concurrency, VectorColumnIndexer::Ptr *merged_indexer) { vector_column_params::MergeOptions merge_options; if (concurrency == 0) { merge_options.pool = GlobalResource::Instance().optimize_thread_pool(); + if (merge_options.pool == nullptr) { + return Status::InternalError( + "Optimize thread pool initialization failed"); + } merge_options.write_concurrency = static_cast(merge_options.pool->count()); } else { @@ -806,7 +811,12 @@ Status SegmentHelper::MergeWithOptionalReuse( bool reused_base_index = false; Status s; - if (CanReuseFirstIndexer(source_indexers, index_field, filter)) { + // BufferStorage currently treats opening an existing file as read-only. + // Until writable reopen is explicit in IndexStorage::open(), rebuilding a + // BufferStorage target is the only way to keep every build page charged to + // MemoryLimitPool without silently falling back to mmap. + if (enable_mmap && + CanReuseFirstIndexer(source_indexers, index_field, filter)) { const auto &first_indexer = source_indexers.front(); LOG_INFO( "Reusing first indexer as merge base. " @@ -816,7 +826,8 @@ Status SegmentHelper::MergeWithOptionalReuse( if (FileHelper::CopyFile(first_indexer->index_file_path(), output_index_path)) { // Open the copied file in-place (create_new=false). - s = vector_indexer->Open(vector_column_params::ReadOptions{true, false}); + s = vector_indexer->Open( + vector_column_params::ReadOptions{enable_mmap, false}); CHECK_RETURN_STATUS(s); source_indexers.erase(source_indexers.begin()); @@ -831,7 +842,8 @@ Status SegmentHelper::MergeWithOptionalReuse( } if (!reused_base_index) { - s = vector_indexer->Open(vector_column_params::ReadOptions{true, true}); + s = vector_indexer->Open( + vector_column_params::ReadOptions{enable_mmap, true}); CHECK_RETURN_STATUS(s); s = vector_indexer->Merge(source_indexers, filter, merge_options); diff --git a/src/db/index/segment/segment_helper.h b/src/db/index/segment/segment_helper.h index ab8c3a9df..588075914 100644 --- a/src/db/index/segment/segment_helper.h +++ b/src/db/index/segment/segment_helper.h @@ -35,13 +35,14 @@ struct CompactTask { const CollectionSchema::Ptr &schema, const std::vector &input_segments, SegmentID output_segment_id, const IndexFilter::Ptr filter, - bool forward_use_parquet, int concurrency) + bool forward_use_parquet, bool enable_mmap, int concurrency) : collection_path_(collection_path), schema_(schema), input_segments_(input_segments), output_segment_id_(output_segment_id), filter_(std::move(filter)), forward_use_parquet_(forward_use_parquet), + enable_mmap_(enable_mmap), concurrency_(concurrency) {} const std::string collection_path_; @@ -52,6 +53,7 @@ struct CompactTask { SegmentID output_segment_id_; const IndexFilter::Ptr filter_; bool forward_use_parquet_; + bool enable_mmap_; int concurrency_; // output @@ -255,8 +257,8 @@ class SegmentHelper { const std::vector &input_segments, const std::string &output_segment_path, const IndexFilter::Ptr &filter, std::function &block_id_generator, uint64_t min_doc_id, - uint64_t max_doc_id, uint32_t doc_count, int concurrency, - std::vector *output_block_metas); + uint64_t max_doc_id, uint32_t doc_count, bool enable_mmap, + int concurrency, std::vector *output_block_metas); // Merges `source_indexers` into a new VectorColumnIndexer at // `output_index_path`. When the first indexer is eligible for reuse (see @@ -267,7 +269,7 @@ class SegmentHelper { static Status MergeWithOptionalReuse( const std::string &output_index_path, const FieldSchema &index_field, std::vector source_indexers, - const IndexFilter::Ptr &filter, int concurrency, + const IndexFilter::Ptr &filter, bool enable_mmap, int concurrency, VectorColumnIndexer::Ptr *merged_indexer); // Returns a FieldSchema clone whose index_params is ready for building the diff --git a/src/db/index/storage/parquet_buffer_pool.cc b/src/db/index/storage/parquet_buffer_pool.cc index 48879bcc4..398d50631 100644 --- a/src/db/index/storage/parquet_buffer_pool.cc +++ b/src/db/index/storage/parquet_buffer_pool.cc @@ -13,7 +13,11 @@ // limitations under the License. #include "parquet_buffer_pool.h" +#include +#include #include +#include +#include #include #include #include @@ -25,6 +29,63 @@ namespace zvec { +namespace { + +void RetainAndDetachArrowBuffers( + const std::shared_ptr &data, + detail::ParquetBufferPayload *payload, + std::unordered_set *seen, size_t *size) { + if (!data) { + return; + } + for (auto &buffer : data->buffers) { + if (!buffer) { + continue; + } + std::shared_ptr retained = buffer; + if (seen->insert(retained.get()).second) { + payload->arrow_refs.emplace_back(retained); + const int64_t arrow_capacity = retained->capacity(); + if (arrow_capacity < 0 || static_cast(arrow_capacity) > + std::numeric_limits::max()) { + *size = std::numeric_limits::max(); + } else { + const size_t capacity = static_cast(arrow_capacity); + *size = capacity > std::numeric_limits::max() - *size + ? std::numeric_limits::max() + : *size + capacity; + } + } + buffer = + std::shared_ptr(retained.get(), [](arrow::Buffer *) {}); + } + for (const auto &child : data->child_data) { + RetainAndDetachArrowBuffers(child, payload, seen, size); + } + RetainAndDetachArrowBuffers(data->dictionary, payload, seen, size); +} + +std::shared_ptr CloneWithPinnedBuffers( + const std::shared_ptr &data, + const std::shared_ptr &pin) { + if (!data) { + return nullptr; + } + auto clone = data->Copy(); + for (auto &buffer : clone->buffers) { + if (buffer) { + buffer = std::shared_ptr(pin, buffer.get()); + } + } + for (auto &child : clone->child_data) { + child = CloneWithPinnedBuffers(child, pin); + } + clone->dictionary = CloneWithPinnedBuffers(clone->dictionary, pin); + return clone; +} + +} // namespace + ParquetBufferID::ParquetBufferID(const std::string &filename, int column, int row_group) : filename(filename), column(column), row_group(row_group) { @@ -32,8 +93,10 @@ ParquetBufferID::ParquetBufferID(const std::string &filename, int column, if (stat(filename.c_str(), &file_stat) == 0) { file_id = file_stat.st_ino; std::filesystem::path p(filename); - auto ftime = std::filesystem::last_write_time(p); - mtime = static_cast(ftime.time_since_epoch().count()); + std::error_code error; + auto ftime = std::filesystem::last_write_time(p, error); + mtime = error ? static_cast(file_stat.st_mtime) + : static_cast(ftime.time_since_epoch().count()); } } @@ -61,6 +124,25 @@ ParquetBufferContextHandle::~ParquetBufferContextHandle() { } } +std::shared_ptr ParquetBufferContextHandle::data() const { + if (!arrow_) { + return nullptr; + } + auto pin = std::make_shared(*this); + if (!pin->arrow_) { + return nullptr; + } + // Alias every Arrow buffer to the cache pin while sharing payload bytes. + arrow::ArrayVector chunks; + chunks.reserve(arrow_->num_chunks()); + for (const auto &chunk : arrow_->chunks()) { + chunks.emplace_back( + arrow::MakeArray(CloneWithPinnedBuffers(chunk->data(), pin))); + } + return std::make_shared(std::move(chunks), + arrow_->type()); +} + bool detail::ParquetBufferLoader::load(const ParquetBufferID &buffer_id, ParquetBufferPayload &payload, size_t &size) { @@ -97,18 +179,9 @@ bool detail::ParquetBufferLoader::load(const ParquetBufferID &buffer_id, size = 0; payload.arrow_refs.clear(); + std::unordered_set seen; for (auto &array : payload.arrow->chunks()) { - auto &buffers = array->data()->buffers; - for (size_t buf_idx = 0; buf_idx < buffers.size(); ++buf_idx) { - if (buffers[buf_idx] == nullptr) { - continue; - } - payload.arrow_refs.emplace_back(buffers[buf_idx]); - size += buffers[buf_idx]->capacity(); - std::shared_ptr hijacked_buffer(buffers[buf_idx].get(), - [](arrow::Buffer *) {}); - buffers[buf_idx] = hijacked_buffer; - } + RetainAndDetachArrowBuffers(array->data(), &payload, &seen, &size); } return true; diff --git a/src/db/index/storage/parquet_buffer_pool.h b/src/db/index/storage/parquet_buffer_pool.h index f04871cf2..730168ed3 100644 --- a/src/db/index/storage/parquet_buffer_pool.h +++ b/src/db/index/storage/parquet_buffer_pool.h @@ -34,7 +34,7 @@ struct ParquetBufferID { int column{0}; int row_group{0}; uint64_t file_id{0}; - long mtime{0}; + int64_t mtime{0}; ParquetBufferID() = default; ParquetBufferID(const std::string &filename, int column, int row_group); @@ -44,10 +44,14 @@ struct ParquetBufferID { struct ParquetBufferIDHash { size_t operator()(const ParquetBufferID &buffer_id) const { - size_t hash = std::hash{}(1); - hash = hash ^ (std::hash{}(buffer_id.file_id)); - hash = hash * 31 + std::hash{}(buffer_id.column); - hash = hash * 31 + std::hash{}(buffer_id.row_group); + size_t hash = std::hash{}(buffer_id.filename); + const auto combine = [&hash](size_t value) { + hash ^= value + 0x9e3779b9U + (hash << 6) + (hash >> 2); + }; + combine(std::hash{}(buffer_id.file_id)); + combine(std::hash{}(buffer_id.mtime)); + combine(std::hash{}(buffer_id.column)); + combine(std::hash{}(buffer_id.row_group)); return hash; } }; @@ -103,9 +107,8 @@ class ParquetBufferContextHandle { ~ParquetBufferContextHandle(); - std::shared_ptr data() { - return arrow_; - } + //! Return an Arrow view whose buffers retain the cache pin. + std::shared_ptr data() const; private: ParquetBufferID buffer_id_; @@ -129,6 +132,8 @@ class ParquetBufferPool { ParquetBufferPool &operator=(ParquetBufferPool &&) = delete; private: + static constexpr size_t kMaxConcurrentLoads = 4; + friend class ParquetBufferContextHandle; using Cache = @@ -136,7 +141,8 @@ class ParquetBufferPool { detail::ParquetBufferLoader, ParquetBufferIDHash, ParquetBufferIDEqual>; - ParquetBufferPool() = default; + ParquetBufferPool() + : cache_(detail::ParquetBufferLoader{}, kMaxConcurrentLoads) {} std::shared_ptr retain(ParquetBufferID buffer_id); diff --git a/src/db/sqlengine/planner/query_planner.cc b/src/db/sqlengine/planner/query_planner.cc index 951f588a8..ec5f3438a 100644 --- a/src/db/sqlengine/planner/query_planner.cc +++ b/src/db/sqlengine/planner/query_planner.cc @@ -432,6 +432,10 @@ Result QueryPlanner::make_physical_plan( // multi segment logic ailego::ThreadPool *pool = GlobalResource::Instance().query_thread_pool(); + if (pool == nullptr) { + return tl::make_unexpected( + Status::InternalError("Query thread pool initialization failed")); + } auto recall_node = std::make_shared(std::move(segment_plans), pool); auto source_node_options = diff --git a/src/include/zvec/ailego/buffer/block_eviction_queue.h b/src/include/zvec/ailego/buffer/block_eviction_queue.h index b93a62bd6..fec154291 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,251 @@ 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; + } + 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..65c7ff347 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,420 @@ 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; 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 +847,28 @@ class ZVEC_AILEGO_API VecBufferPoolHandle { char *get_single_page(size_t file_offset, size_t len, size_t &out_page_id); + 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 +880,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/src/include/zvec/core/framework/index_segment_storage.h b/src/include/zvec/core/framework/index_segment_storage.h index 06b17779a..bf350fc52 100644 --- a/src/include/zvec/core/framework/index_segment_storage.h +++ b/src/include/zvec/core/framework/index_segment_storage.h @@ -49,6 +49,17 @@ class IndexSegmentStorage : public IndexStorage { data_crc_(segment.data_crc()), parent_(parent->clone()) {} + //! Constructor (for clone) + Segment(const IndexStorage::Segment::Pointer &cloned_parent, + size_t data_offset, size_t data_size, size_t padding_size, + uint32_t data_crc) + : data_offset_(data_offset), + data_size_(data_size), + padding_size_(padding_size), + region_size_(data_size + padding_size), + data_crc_(data_crc), + parent_(cloned_parent) {} + //! Destructor ~Segment(void) override {} @@ -114,9 +125,11 @@ class IndexSegmentStorage : public IndexStorage { return; } - //! Clone the segment + //! Clone the segment (each clone gets an independent parent buffer + //! for thread safety — concurrent reads require separate buffers). IndexStorage::Segment::Pointer clone(void) override { - return shared_from_this(); + return std::make_shared(parent_->clone(), data_offset_, + data_size_, padding_size_, data_crc_); } private: diff --git a/src/include/zvec/core/framework/index_storage.h b/src/include/zvec/core/framework/index_storage.h index 049c79917..5d2479ca4 100644 --- a/src/include/zvec/core/framework/index_storage.h +++ b/src/include/zvec/core/framework/index_storage.h @@ -15,6 +15,8 @@ #pragma once #include +#include +#include #include #include #include @@ -39,7 +41,7 @@ class IndexStorage : public IndexModule { MBT_HEAP_SCRATCH = 3, }; - MemoryBlock() {} + MemoryBlock() = default; MemoryBlock(ailego::VecBufferPoolHandle *buffer_pool_handle, size_t block_id, void *data) : type_(MemoryBlockType::MBT_BUFFERPOOL) { @@ -47,14 +49,17 @@ class IndexStorage : public IndexModule { buffer_block_id_ = block_id; data_ = data; } + MemoryBlock( + const std::shared_ptr &buffer_pool_handle, + size_t block_id, void *data) + : type_(MemoryBlockType::MBT_BUFFERPOOL), + data_(data), + buffer_pool_handle_owner_(buffer_pool_handle), + buffer_pool_handle_(buffer_pool_handle.get()), + buffer_block_id_(block_id) {} MemoryBlock(void *data) : type_(MemoryBlockType::MBT_MMAP), data_(data) {} - //! Build an HEAP_SCRATCH MemoryBlock that owns `owned` (allocated via - //! ailego_malloc / ailego_aligned_malloc). `size` is the byte length of - //! the buffer and is required so that copy construction / copy - //! assignment can deep-copy the buffer instead of aliasing it (a shallow - //! copy would result in use-after-free once the original block is - //! destructed and frees the buffer). + //! Build an owned heap block; size enables safe deep copies. static MemoryBlock MakeOwned(void *owned, size_t size) { MemoryBlock mb; mb.type_ = MemoryBlockType::MBT_HEAP_SCRATCH; @@ -63,19 +68,34 @@ class IndexStorage : public IndexModule { return mb; } + //! Build a non-owning view over caller-managed memory (e.g. a query-level + //! scratch arena). The block frees and pins nothing on destruction, so the + //! backing buffer must outlive every copy of this block. Uses the MMAP + //! representation whose destructor is a no-op. + static MemoryBlock MakeBorrowedView(void *data) { + MemoryBlock mb; + mb.type_ = MemoryBlockType::MBT_MMAP; + mb.data_ = data; + return mb; + } + MemoryBlock(const MemoryBlock &rhs) { switch (rhs.type_) { case MemoryBlockType::MBT_MMAP: this->reset(rhs.data_); break; case MemoryBlockType::MBT_BUFFERPOOL: - this->reset(rhs.buffer_pool_handle_, rhs.buffer_block_id_, rhs.data_); + if (rhs.buffer_pool_handle_owner_) { + this->reset(rhs.buffer_pool_handle_owner_, rhs.buffer_block_id_, + rhs.data_); + } else { + this->reset(rhs.buffer_pool_handle_, rhs.buffer_block_id_, + rhs.data_); + } buffer_pool_handle_->acquire_one(buffer_block_id_); break; case MemoryBlockType::MBT_HEAP_SCRATCH: - // Deep copy: each owner must hold its own buffer, otherwise the - // first destructor frees the buffer and leaves the surviving - // copies dangling. + // Heap blocks do not share ownership. deep_copy_from(rhs); break; default: @@ -83,15 +103,19 @@ class IndexStorage : public IndexModule { } } - MemoryBlock(MemoryBlock &&rhs) { + MemoryBlock(MemoryBlock &&rhs) noexcept { switch (rhs.type_) { case MemoryBlockType::MBT_MMAP: - this->reset(std::move(rhs.data_)); + this->reset(rhs.data_); break; case MemoryBlockType::MBT_BUFFERPOOL: - this->reset(std::move(rhs.buffer_pool_handle_), - std::move(rhs.buffer_block_id_), std::move(rhs.data_)); + type_ = MemoryBlockType::MBT_BUFFERPOOL; + data_ = rhs.data_; + buffer_pool_handle_owner_ = std::move(rhs.buffer_pool_handle_owner_); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; rhs.buffer_pool_handle_ = nullptr; + rhs.data_ = nullptr; rhs.type_ = MemoryBlockType::MBT_UNKNOWN; break; case MemoryBlockType::MBT_HEAP_SCRATCH: @@ -114,8 +138,13 @@ class IndexStorage : public IndexModule { this->reset(rhs.data_); break; case MemoryBlockType::MBT_BUFFERPOOL: - this->reset(rhs.buffer_pool_handle_, rhs.buffer_block_id_, - rhs.data_); + if (rhs.buffer_pool_handle_owner_) { + this->reset(rhs.buffer_pool_handle_owner_, rhs.buffer_block_id_, + rhs.data_); + } else { + this->reset(rhs.buffer_pool_handle_, rhs.buffer_block_id_, + rhs.data_); + } buffer_pool_handle_->acquire_one(buffer_block_id_); break; case MemoryBlockType::MBT_HEAP_SCRATCH: @@ -123,22 +152,29 @@ class IndexStorage : public IndexModule { deep_copy_from(rhs); break; default: + release_current(); break; } } return *this; } - MemoryBlock &operator=(MemoryBlock &&rhs) { + MemoryBlock &operator=(MemoryBlock &&rhs) noexcept { if (this != &rhs) { switch (rhs.type_) { case MemoryBlockType::MBT_MMAP: - this->reset(std::move(rhs.data_)); + this->reset(rhs.data_); break; case MemoryBlockType::MBT_BUFFERPOOL: - this->reset(std::move(rhs.buffer_pool_handle_), - std::move(rhs.buffer_block_id_), std::move(rhs.data_)); + release_current(); + type_ = MemoryBlockType::MBT_BUFFERPOOL; + data_ = rhs.data_; + buffer_pool_handle_owner_ = + std::move(rhs.buffer_pool_handle_owner_); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; rhs.buffer_pool_handle_ = nullptr; + rhs.data_ = nullptr; rhs.type_ = MemoryBlockType::MBT_UNKNOWN; break; case MemoryBlockType::MBT_HEAP_SCRATCH: @@ -151,6 +187,7 @@ class IndexStorage : public IndexModule { rhs.type_ = MemoryBlockType::MBT_UNKNOWN; break; default: + release_current(); break; } } @@ -180,23 +217,37 @@ class IndexStorage : public IndexModule { return data_; } + void reset() { + release_current(); + } + void reset(ailego::VecBufferPoolHandle *buffer_pool_handle, size_t block_id, void *data) { - if (type_ == MemoryBlockType::MBT_BUFFERPOOL) { - buffer_pool_handle_->release_one(buffer_block_id_); - } else if (type_ == MemoryBlockType::MBT_HEAP_SCRATCH) { - release_owned(); - } + release_current(); type_ = MemoryBlockType::MBT_BUFFERPOOL; buffer_pool_handle_ = buffer_pool_handle; buffer_block_id_ = block_id; data_ = data; } + void reset( + const std::shared_ptr &buffer_pool_handle, + size_t block_id, void *data) { + release_current(); + type_ = MemoryBlockType::MBT_BUFFERPOOL; + buffer_pool_handle_owner_ = buffer_pool_handle; + buffer_pool_handle_ = buffer_pool_handle.get(); + buffer_block_id_ = block_id; + data_ = data; + } + void reset(void *data) { if (type_ == MemoryBlockType::MBT_BUFFERPOOL) { - buffer_pool_handle_->release_one(buffer_block_id_); + if (buffer_pool_handle_) { + buffer_pool_handle_->release_one(buffer_block_id_); + } buffer_pool_handle_ = nullptr; + buffer_pool_handle_owner_.reset(); } else if (type_ == MemoryBlockType::MBT_HEAP_SCRATCH) { release_owned(); } @@ -206,11 +257,10 @@ class IndexStorage : public IndexModule { MemoryBlockType type_{MBT_UNKNOWN}; void *data_{nullptr}; + std::shared_ptr buffer_pool_handle_owner_{}; mutable ailego::VecBufferPoolHandle *buffer_pool_handle_{nullptr}; size_t buffer_block_id_{0}; - //! Byte size of the heap-scratch buffer pointed to by `data_`; only used - //! when type_ == MBT_HEAP_SCRATCH. Required for safe deep-copy on - //! copy-construction / copy-assignment of HEAP_SCRATCH blocks. + //! Byte size used to copy heap scratch blocks. size_t scratch_size_{0}; private: @@ -222,10 +272,7 @@ class IndexStorage : public IndexModule { scratch_size_ = 0; } - //! Drop whatever the current MemoryBlock holds, regardless of type, so - //! that the slot is ready to receive new ownership. Mirrors what the - //! destructor would do (minus zeroing data_) but leaves the type alone - //! for the caller to overwrite immediately afterwards. + //! Release the current representation and reset the block. void release_current() { switch (type_) { case MemoryBlockType::MBT_BUFFERPOOL: @@ -233,6 +280,7 @@ class IndexStorage : public IndexModule { buffer_pool_handle_->release_one(buffer_block_id_); buffer_pool_handle_ = nullptr; } + buffer_pool_handle_owner_.reset(); break; case MemoryBlockType::MBT_HEAP_SCRATCH: release_owned(); @@ -241,18 +289,21 @@ class IndexStorage : public IndexModule { break; } data_ = nullptr; + scratch_size_ = 0; type_ = MemoryBlockType::MBT_UNKNOWN; } - //! Allocate a fresh buffer of the same size as `rhs.scratch_size_`, - //! memcpy `rhs.data_` into it, and become the new owner. Used by the - //! HEAP_SCRATCH copy ctor / copy assignment so the original and the - //! copy each free their own buffer independently. + //! Deep-copy heap scratch ownership. void deep_copy_from(const MemoryBlock &rhs) { type_ = MemoryBlockType::MBT_HEAP_SCRATCH; scratch_size_ = rhs.scratch_size_; if (scratch_size_ > 0 && rhs.data_) { data_ = ailego_malloc(scratch_size_); + if (data_ == nullptr) { + scratch_size_ = 0; + type_ = MemoryBlockType::MBT_UNKNOWN; + throw std::bad_alloc(); + } std::memcpy(data_, rhs.data_, scratch_size_); } else { data_ = nullptr; @@ -268,6 +319,9 @@ class IndexStorage : public IndexModule { SegmentData(size_t off, size_t len) : offset(off), length(len), data(nullptr) {} + SegmentData(size_t off, size_t len, const void *ptr) + : offset(off), length(len), data(ptr) {} + //! Members size_t offset; size_t length; @@ -280,6 +334,31 @@ class IndexStorage : public IndexModule { //! Index Storage Pointer typedef std::shared_ptr Pointer; + //! Cache admission/eviction hint. Backends without an evictable cache + //! ignore it; page-backed storage maps it to its eviction queues. + enum class CachePriority : uint8_t { + kLow = 0, + kNormal = 1, + kHigh = 2, + }; + + //! One bounded-lifetime read in a batch. Requests may reference different + //! segments owned by the same storage so backends can merge their page + //! misses into one I/O submission. + struct BorrowedRead { + BorrowedRead(Segment *segment_arg, size_t offset_arg, size_t length_arg, + MemoryBlock *block_arg) + : segment(segment_arg), + offset(offset_arg), + length(length_arg), + block(block_arg) {} + + Segment *segment; + size_t offset; + size_t length; + MemoryBlock *block; + }; + //! Destructor virtual ~Segment(void) {} @@ -308,6 +387,62 @@ class IndexStorage : public IndexModule { virtual size_t read(size_t offset, MemoryBlock &data, size_t len) = 0; + //! Read bytes that the caller guarantees will never be modified after + //! publication. Page-backed writable storage may safely return a pinned + //! direct view instead of copying the range for snapshot isolation. + virtual size_t read_immutable(size_t offset, MemoryBlock &data, + size_t len) { + return read(offset, data, len); + } + + //! Borrowed read; release the block before this Segment. The default keeps + //! the owning read behavior. + virtual size_t read_borrowed(size_t offset, MemoryBlock &data, size_t len) { + return read(offset, data, len); + } + + //! Borrowed-handle variant of read_immutable(). + virtual size_t read_borrowed_immutable(size_t offset, MemoryBlock &data, + size_t len) { + return read_immutable(offset, data, len); + } + + //! Whether batching is currently preferable to scalar borrowed reads. + virtual bool prefer_borrowed_batch() const { + return false; + } + + //! Batch borrowed reads. The default preserves compatibility by issuing + //! scalar reads; page-backed implementations may override this to batch + //! misses while keeping each returned MemoryBlock pinned independently. + virtual bool read_borrowed_batch(BorrowedRead *reads, size_t count) { + if (count == 0) { + return true; + } + if (reads == nullptr) { + return false; + } + for (size_t i = 0; i < count; ++i) { + if (reads[i].segment == nullptr || reads[i].block == nullptr) { + return false; + } + } + for (size_t i = 0; i < count; ++i) { + reads[i].block->reset(); + } + for (size_t i = 0; i < count; ++i) { + BorrowedRead &request = reads[i]; + if (request.segment->read_borrowed(request.offset, *request.block, + request.length) != request.length) { + for (size_t j = 0; j < count; ++j) { + reads[j].block->reset(); + } + return false; + } + } + return true; + } + virtual bool read(SegmentData *, size_t) { return false; } @@ -324,14 +459,84 @@ class IndexStorage : public IndexModule { //! Clone the segment virtual Pointer clone(void) = 0; - //! Retrieve the stable base data pointer if the storage backend supports - //! it (e.g. mmap-backed storage). Returns nullptr for backends with - //! mutable/evictable buffers (e.g. BufferStorage). When non-null the - //! caller may compute element addresses as base_data() + offset directly, - //! avoiding the full pointer chain through chunk->read(). + //! Return a stable base address, or nullptr for evictable storage. virtual const uint8_t *base_data(void) const { return nullptr; } + + virtual size_t abs_data_offset(void) const { + return 0; + } + + virtual void prefetch(size_t offset, size_t len, + CachePriority priority = CachePriority::kLow) { + (void)offset; + (void)len; + (void)priority; + } + + //! Apply ordered writes to this segment. Kept at the end of the vtable so + //! existing method slots remain stable. Backends may share pins/latches; + //! the default preserves the scalar write contract. + virtual bool write_batch(const SegmentData *writes, size_t count) { + if (count == 0) { + return true; + } + if (writes == nullptr) { + return false; + } + for (size_t i = 0; i < count; ++i) { + if (writes[i].length == 0) { + continue; + } + if (writes[i].data == nullptr || + write(writes[i].offset, writes[i].data, writes[i].length) != + writes[i].length) { + return false; + } + } + return true; + } + + //! Size-aware batch preference. The default retains the backend's existing + //! policy; page-backed writable storage can account for cross-page cost. + virtual bool prefer_borrowed_batch_for(size_t value_size) const { + (void)value_size; + return prefer_borrowed_batch(); + } + + //! Immutable counterpart of read_borrowed_batch(). Writable page-backed + //! storage may safely batch and pin ranges that will never be modified + //! after publication. + virtual bool read_borrowed_batch_immutable(BorrowedRead *reads, + size_t count) { + if (count == 0) { + return true; + } + if (reads == nullptr) { + return false; + } + for (size_t i = 0; i < count; ++i) { + if (reads[i].segment == nullptr || reads[i].block == nullptr) { + return false; + } + } + for (size_t i = 0; i < count; ++i) { + reads[i].block->reset(); + } + for (size_t i = 0; i < count; ++i) { + BorrowedRead &request = reads[i]; + if (request.segment->read_borrowed_immutable( + request.offset, *request.block, request.length) != + request.length) { + for (size_t j = 0; j < count; ++j) { + reads[j].block->reset(); + } + return false; + } + } + return true; + } }; //! Destructor @@ -386,6 +591,11 @@ class IndexStorage : public IndexModule { return MemoryBlock::MBT_MMAP; } + //! Return the shared page cache when this storage is backed by VecBufferPool. + virtual std::shared_ptr vec_buffer_pool(void) const { + return nullptr; + } + //! Test if the storage has unflushed data virtual bool is_dirty(void) const { return false; diff --git a/src/include/zvec/db/config.h b/src/include/zvec/db/config.h index ff17570bf..5180cd75d 100644 --- a/src/include/zvec/db/config.h +++ b/src/include/zvec/db/config.h @@ -14,6 +14,7 @@ #pragma once #include +#include #include #include #include @@ -127,91 +128,124 @@ class ZVEC_API GlobalConfig : public ailego::Singleton { uint64_t memory_limit_bytes() const noexcept; const LogConfig &log_config() const noexcept { - return *config_.log_config; + auto config = config_snapshot(); + return *config->log_config; } std::string log_type() const noexcept { - return config_.log_config->GetLoggerType(); + auto config = config_snapshot(); + return config->log_config->GetLoggerType(); } LogLevel log_level() const noexcept { - return config_.log_config->level; + auto config = config_snapshot(); + return config->log_config->level; } // File log specific accessors (only valid when using FileLogConfig) const std::string &log_dir() const noexcept { + auto config = config_snapshot(); const FileLogConfig *file_config = - dynamic_cast(config_.log_config.get()); + dynamic_cast(config->log_config.get()); static const std::string empty_string = ""; return file_config ? file_config->dir : empty_string; } const std::string &log_file_basename() const noexcept { + auto config = config_snapshot(); const FileLogConfig *file_config = - dynamic_cast(config_.log_config.get()); + dynamic_cast(config->log_config.get()); static const std::string empty_string = ""; return file_config ? file_config->basename : empty_string; } uint32_t log_file_size() const noexcept { + auto config = config_snapshot(); const FileLogConfig *file_config = - dynamic_cast(config_.log_config.get()); + dynamic_cast(config->log_config.get()); return file_config ? file_config->file_size : 0; } uint32_t log_overdue_days() const noexcept { + auto config = config_snapshot(); const FileLogConfig *file_config = - dynamic_cast(config_.log_config.get()); + dynamic_cast(config->log_config.get()); return file_config ? file_config->overdue_days : 0; } //! Query thread count uint32_t query_thread_count() const noexcept { - return config_.query_thread_count; + return config_snapshot()->query_thread_count; } //! Query thread binding bool query_thread_binding() const noexcept { - return config_.query_thread_binding; + return config_snapshot()->query_thread_binding; } //! Invert to forward scan ratio float invert_to_forward_scan_ratio() const noexcept { - return config_.invert_to_forward_scan_ratio; + return config_snapshot()->invert_to_forward_scan_ratio; } //! Brute force by keys ratio float brute_force_by_keys_ratio() const noexcept { - return config_.brute_force_by_keys_ratio; + return config_snapshot()->brute_force_by_keys_ratio; } //! FTS brute force by keys ratio (independent from brute_force_by_keys_ratio //! because FTS per-candidate cost is higher). float fts_brute_force_by_keys_ratio() const noexcept { - return config_.fts_brute_force_by_keys_ratio; + return config_snapshot()->fts_brute_force_by_keys_ratio; } //! Optimize thread count uint32_t optimize_thread_count() const noexcept { - return config_.optimize_thread_count; + return config_snapshot()->optimize_thread_count; } //! Optimize thread binding bool optimize_thread_binding() const noexcept { - return config_.optimize_thread_binding; + return config_snapshot()->optimize_thread_binding; } //! Effective jieba dict dir. Thread-safe. std::string jieba_dict_dir() const; private: - // Configuration data - ConfigData config_; + enum class InitializationState : uint8_t { + kUninitialized, + kInitializing, + kInitialized, + kFailed, + }; - // Atomic flag to ensure initialization happens only once - std::atomic initialized_{false}; + std::shared_ptr config_snapshot() const noexcept { + return std::atomic_load_explicit(&config_, std::memory_order_acquire); + } - // Guards config_ fields that may be written outside Initialize(). + // Readers atomically acquire an immutable snapshot, so Initialize() and the + // language-SDK jieba setter can publish whole configurations without data + // races or mixed-field observations. + std::shared_ptr config_{ + std::make_shared()}; + + // The legacy logging accessors return references. Keep replaced LogConfig + // objects alive for this GlobalConfig's lifetime so a reference acquired + // concurrently with the one-time snapshot publication cannot dangle. + std::shared_ptr retained_log_config_; + + // Initialize() can be called concurrently by language bindings. Publish a + // terminal state only after every initialization stage has completed, and + // make followers observe the same result instead of returning early while + // the winning thread is still working. + InitializationState initialization_state_{ + InitializationState::kUninitialized}; + Status initialization_status_{}; + std::condition_variable initialization_cv_; + std::mutex initialization_mutex_; + + // Serializes immutable snapshot writers and retained_log_config_. mutable std::mutex mutex_; }; 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..73d18d2ec --- /dev/null +++ b/tests/ailego/buffer/vector_page_table_test.cc @@ -0,0 +1,2061 @@ +// 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, 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(); + } + EXPECT_TRUE(table.force_evict_block(/*block_id=*/0)); +} + +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()); + 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) { + 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 second-chance mechanism must have spared at least some pages. + EXPECT_GT(s.second_chance, 0u); +} + +// --------------------------------------------------------------------------- +// 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__ diff --git a/tests/core/algorithm/diskann/diskann_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 752759e9d..99034a6f8 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include "diskann_context.h" #include "diskann_holder.h" using namespace zvec::core; @@ -98,6 +99,38 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) { ASSERT_EQ(0UL, stats.discarded_count()); ASSERT_GT(stats.trained_costtime(), 0UL); ASSERT_GT(stats.built_costtime(), 0UL); + + IndexStreamer::Pointer streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(nullptr, streamer); + Params search_params; + search_params.set("zvec.diskann.searcher.list_size", 100); + ASSERT_EQ(0, streamer->init(*_index_meta_ptr, search_params)); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer->open(storage)); + + auto context = streamer->create_context(); + ASSERT_NE(nullptr, context); + context->set_topk(10); + Params query_params; + query_params.set("zvec.diskann.searcher.list_size", 37); + ASSERT_EQ(0, context->update(query_params)); + auto *diskann_context = dynamic_cast(context.get()); + ASSERT_NE(nullptr, diskann_context); + ASSERT_EQ(37u, diskann_context->list_size()); + auto *const original_context = context.get(); + NumericalVector query(dim, 3.1f); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + + ASSERT_EQ(0, streamer->search_impl(query.data(), qmeta, context)); + EXPECT_EQ(original_context, context.get()); + EXPECT_EQ(37u, diskann_context->list_size()); + ASSERT_EQ(0, streamer->search_impl(query.data(), qmeta, context)); + EXPECT_EQ(original_context, context.get()); + EXPECT_EQ(37u, diskann_context->list_size()); } // Regression test: building a small DiskAnn index must complete quickly. diff --git a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc index a416b0623..f5b0766eb 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include "diskann_file_reader.h" #if defined(__linux) || defined(__linux__) @@ -37,6 +38,7 @@ int execute_io_libaio(io_context_t &ctx, int fd, } // namespace zvec using namespace zvec::core; +namespace ailego = zvec::ailego; namespace { @@ -147,6 +149,10 @@ class TemporaryFile { return fd_; } + const char *path() const { + return path_; + } + private: char path_[64] = "DiskAnnLinuxAioTest.XXXXXX"; int fd_; @@ -258,4 +264,159 @@ TEST(DiskAnnLinuxAioTest, DrainsAllCompletionsBeforePreadFallback) { std::free(output); } +TEST(DiskAnnBufferPoolFileReaderTest, + ReadsScatteredRequestsThroughOnePinnedPageBatch) { + if (ailego::kVectorPageSize != DiskAnnUtil::kSectorSize) { + GTEST_SKIP() << "DiskAnn sectors require one native buffer-pool page"; + } + + constexpr size_t kPageCount = 4; + const size_t page_size = ailego::kVectorPageSize; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kPageCount * page_size); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * page_size, static_cast(page + 1), + page_size); + } + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + ASSERT_EQ(::fsync(file.fd()), 0); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + 0, memory_pool.init( + 8 * page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(kPageCount))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + EXPECT_FALSE(reader.requires_io_context()); + reader.open(file.path()); + + void *output = nullptr; + ASSERT_EQ(::posix_memalign(&output, page_size, 4 * page_size), 0); + ASSERT_NE(output, nullptr); + std::memset(output, 0, 4 * page_size); + std::vector requests; + requests.emplace_back(3 * page_size, page_size, output); + requests.emplace_back(page_size, 2 * page_size, + static_cast(output) + page_size); + requests.emplace_back(3 * page_size, page_size, + static_cast(output) + 3 * page_size); + + IOContext unused{}; + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output, source.data() + 3 * page_size, page_size), 0); + EXPECT_EQ(std::memcmp(static_cast(output) + page_size, + source.data() + page_size, 2 * page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 3 * page_size, + source.data() + 3 * page_size, page_size), + 0); + EXPECT_EQ(pool->stats().miss, 3u); + + pool->page_table_.force_evict_all_loaded(); + std::free(output); +} + +TEST(DiskAnnBufferPoolFileReaderTest, + BypassesColdMissUnderPressureAndFansOutDuplicates) { + if (ailego::kVectorPageSize != DiskAnnUtil::kSectorSize) { + GTEST_SKIP() << "DiskAnn sectors require one native buffer-pool page"; + } + + constexpr size_t kPageCount = 4; + const size_t page_size = ailego::kVectorPageSize; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kPageCount * page_size); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * page_size, static_cast(page + 1), + page_size); + } + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + ASSERT_EQ(::fsync(file.fd()), 0); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + 0, memory_pool.init( + page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(kPageCount))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + reader.open(file.path()); + + char *seed = pool->acquire_buffer(0, 10); + ASSERT_NE(seed, nullptr); + + void *output = nullptr; + ASSERT_EQ(::posix_memalign(&output, page_size, 4 * page_size), 0); + ASSERT_NE(output, nullptr); + std::vector requests; + requests.emplace_back(page_size, 3 * page_size, output); + requests.emplace_back(2 * page_size, page_size, + static_cast(output) + 3 * page_size); + IOContext unused{}; + + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output, source.data() + page_size, page_size), 0); + EXPECT_EQ(std::memcmp(static_cast(output) + page_size, + source.data() + 2 * page_size, page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 2 * page_size, + source.data() + 3 * page_size, page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 3 * page_size, + source.data() + 2 * page_size, page_size), + 0); + EXPECT_FALSE(pool->is_page_resident(1)); + EXPECT_FALSE(pool->is_page_resident(2)); + EXPECT_FALSE(pool->is_page_resident(3)); + EXPECT_EQ(pool->stats().admission_rejected, 3u); + EXPECT_EQ(pool->stats().bypass_reads, 1u); + EXPECT_EQ(pool->stats().bypass_bytes, 3 * page_size); + EXPECT_EQ(pool->stats().bypass_io_requests, 1u); + EXPECT_EQ(pool->stats().bypass_rechecks, 3u); + EXPECT_EQ(pool->stats().bypass_cache_joins, 0u); + + pool->page_table_.release_block(0); + pool->page_table_.force_evict_all_loaded(); + EXPECT_EQ(destroy_io_ctx(unused), 0); + std::free(output); +} + +TEST(DiskAnnBufferPoolFileReaderTest, RejectsNonPageAlignedRequests) { + if (ailego::kVectorPageSize <= 512) { + GTEST_SKIP() << "test requires a native page larger than 512 bytes"; + } + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(2 * ailego::kVectorPageSize, 0x5a); + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ(0, memory_pool.init( + 4 * ailego::kVectorPageSize + + ailego::VecBufferPool::metadata_bytes_for_page_count(2))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + + void *output = allocate_aligned(512); + ASSERT_NE(output, nullptr); + std::vector requests; + requests.emplace_back(512, 512, output); + IOContext unused{}; + EXPECT_EQ(reader.read(requests, unused), IndexError_InvalidArgument); + std::free(output); +} + #endif // __linux__ diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc index 17c3c8704..522e94b9d 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_buffer_test.cc @@ -3,7 +3,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -46,6 +48,20 @@ void HnswStreamerTest::TearDown(void) { zvec::test_util::RemoveTestPath(dir_); } +TEST_F(HnswStreamerTest, MaxDegreeIsNeighborCountNotSerializedBytes) { + IndexStreamer::Stats stats; + HnswBufferPoolStreamerEntity entity(stats); + entity.set_l0_neighbor_cnt(192); + entity.set_upper_neighbor_cnt(96); + + EXPECT_EQ(192U, entity.max_degree(0)); + EXPECT_EQ(96U, entity.max_degree(1)); + EXPECT_EQ(sizeof(NeighborsHeader) + 192U * sizeof(node_id_t), + entity.neighbors_size()); + EXPECT_EQ(sizeof(NeighborsHeader) + 96U * sizeof(node_id_t), + entity.upper_neighbors_size()); +} + TEST_F(HnswStreamerTest, TestHnswSearch) { MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL); IndexStreamer::Pointer write_streamer = @@ -88,8 +104,83 @@ TEST_F(HnswStreamerTest, TestHnswSearch) { ASSERT_EQ(0, read_storage->init(stg_params)); ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSearch", false)); ASSERT_EQ(0, read_streamer->open(read_storage)); + + auto pool = read_storage->vec_buffer_pool(); + ASSERT_TRUE(pool); + const size_t page_size = kVectorPageSize; + + // The complete upper graph is a small, universal search hotset and should + // be resident at high priority immediately after open. + auto upper_chunk = read_storage->get("HnswT4S0"); + ASSERT_TRUE(upper_chunk); + ASSERT_GT(upper_chunk->data_size(), 0U); + const size_t upper_first = upper_chunk->data_offset() / page_size; + const size_t upper_last = + (upper_chunk->data_offset() + upper_chunk->data_size() - 1) / page_size; + for (size_t page = upper_first; page <= upper_last; ++page) { + EXPECT_TRUE(pool->is_page_resident(page)); + EXPECT_EQ(VecBufferPool::kHighPriority, + pool->page_table_.eviction_priority(page)); + } + + // The entry node record (vector + key + L0 adjacency) is also shared by + // every query and must retain high priority. + auto header_chunk = read_storage->get("HnswT1S0"); + ASSERT_TRUE(header_chunk); + HNSWHeader hnsw_header; + ASSERT_EQ(sizeof(hnsw_header), + header_chunk->fetch(0, &hnsw_header, sizeof(hnsw_header))); + ASSERT_NE(kInvalidNodeId, hnsw_header.entry_point()); + auto first_node_chunk = read_storage->get("HnswT3S0"); + ASSERT_TRUE(first_node_chunk); + ASSERT_GT(hnsw_header.graph.node_size, 0U); + const size_t nodes_per_chunk = + first_node_chunk->data_size() / hnsw_header.graph.node_size; + ASSERT_GT(nodes_per_chunk, 0U); + const size_t entry_chunk_id = hnsw_header.entry_point() / nodes_per_chunk; + const size_t entry_local_id = hnsw_header.entry_point() % nodes_per_chunk; + auto entry_chunk = + read_storage->get("HnswT3S" + std::to_string(entry_chunk_id)); + ASSERT_TRUE(entry_chunk); + const size_t entry_offset = + entry_local_id * static_cast(hnsw_header.graph.node_size); + const size_t entry_first = + (entry_chunk->data_offset() + entry_offset) / page_size; + const size_t entry_last = (entry_chunk->data_offset() + entry_offset + + hnsw_header.graph.node_size - 1) / + page_size; + for (size_t page = entry_first; page <= entry_last; ++page) { + EXPECT_TRUE(pool->is_page_resident(page)); + EXPECT_EQ(VecBufferPool::kHighPriority, + pool->page_table_.eviction_priority(page)); + } + size_t topk = 3; auto provider = read_streamer->create_provider(); + + // Query execution must not repeat range-prefetch priority scans after the + // one-time open protection above. With no eviction pressure, the number of + // normal-priority pages therefore remains unchanged. + size_t normal_before = 0; + for (size_t page = 0; page < pool->page_table_.entry_num(); ++page) { + normal_before += pool->is_page_resident(page) && + pool->page_table_.eviction_priority(page) == + VecBufferPool::kNormalPriority; + } + NumericalVector hotset_query(dim); + for (size_t j = 0; j < dim; ++j) { + hotset_query[j] = cnt / 2; + } + ctx->set_topk(topk); + ASSERT_EQ(0, read_streamer->search_impl(hotset_query.data(), qmeta, ctx)); + size_t normal_after = 0; + for (size_t page = 0; page < pool->page_table_.entry_num(); ++page) { + normal_after += pool->is_page_resident(page) && + pool->page_table_.eviction_priority(page) == + VecBufferPool::kNormalPriority; + } + EXPECT_EQ(normal_after, normal_before); + for (size_t i = 0; i < cnt; i += 1) { NumericalVector vec(dim); for (size_t j = 0; j < dim; ++j) { @@ -294,6 +385,80 @@ TEST_F(HnswStreamerTest, TestHnswSearchBuffer) { cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl; } +TEST_F(HnswStreamerTest, TestWideMConcurrentBuildBuffer) { + MemoryLimitPool::get_instance().init(512UL * 1024UL * 1024UL); + auto streamer = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + + Params params; + params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 96U); + params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 96U); + params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 128U); + params.set(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, 128UL * 1024UL * 1024UL); + params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); + ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params)); + + auto storage = IndexFactory::CreateStorage("BufferStorage"); + ASSERT_NE(nullptr, storage); + Params storage_params; + ASSERT_EQ(0, storage->init(storage_params)); + ASSERT_EQ(0, + storage->open(dir_ + "Test/WideMConcurrentBuildBuffer", true)); + ASSERT_EQ(0, streamer->open(storage)); + + constexpr size_t kThreadCount = 8; + constexpr size_t kVectorsPerThread = 1000; + auto add_vectors = [&streamer](size_t first) { + auto context = streamer->create_context(); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vector(dim); + size_t added = 0; + for (size_t i = 0; i < kVectorsPerThread; ++i) { + const size_t key = first + i; + for (size_t j = 0; j < dim; ++j) { + const uint32_t bits = static_cast( + key * 2654435761ULL + j * 2246822519ULL); + vector[j] = static_cast(bits) / 4294967295.0f; + } + added += streamer->add_impl(key, vector.data(), query_meta, context) == 0; + } + return added; + }; + + std::vector> workers; + for (size_t thread = 0; thread < kThreadCount; ++thread) { + workers.emplace_back(std::async(std::launch::async, add_vectors, + thread * kVectorsPerThread)); + } + for (auto &worker : workers) { + EXPECT_EQ(kVectorsPerThread, worker.get()); + } + + auto provider = streamer->create_provider(); + ASSERT_NE(nullptr, provider); + EXPECT_EQ(kThreadCount * kVectorsPerThread, provider->count()); + + auto search_context = streamer->create_context(); + search_context->set_topk(1); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector query(dim); + for (size_t key = 0; key < kThreadCount * kVectorsPerThread; key += 157) { + for (size_t j = 0; j < dim; ++j) { + const uint32_t bits = + static_cast(key * 2654435761ULL + j * 2246822519ULL); + query[j] = static_cast(bits) / 4294967295.0f; + } + ASSERT_EQ(0, + streamer->search_impl(query.data(), query_meta, search_context)); + const auto &result = search_context->result(); + ASSERT_EQ(1U, result.size()); + EXPECT_EQ(key, result[0].key()); + } + + ASSERT_EQ(0, streamer->flush(0)); + ASSERT_EQ(0, streamer->close()); +} + TEST_F(HnswStreamerTest, TestHnswSearchBufferMMap) { MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL); IndexStreamer::Pointer write_streamer = @@ -334,7 +499,8 @@ TEST_F(HnswStreamerTest, TestHnswSearchBufferMMap) { auto read_storage = IndexFactory::CreateStorage("MMapFileStorage"); ASSERT_NE(nullptr, read_storage); ASSERT_EQ(0, read_storage->init(stg_params)); - ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false)); + ASSERT_EQ(0, + read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false)); ASSERT_EQ(0, read_streamer->open(read_storage)); size_t topk = 3; auto provider = read_streamer->create_provider(); @@ -514,4 +680,4 @@ TEST_F(HnswStreamerTest, TestHnswSearchMMap) { #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop -#endif \ No newline at end of file +#endif diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 26f199c49..99f72955f 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -561,6 +561,73 @@ TEST(IndexInterface, BufferGeneral) { .build()); } +TEST(IndexInterface, HnswBufferPoolSearchWithEviction) { + constexpr uint32_t kDimension = 768; + constexpr uint32_t kDocCount = 512; + constexpr size_t kBufferBudget = 1024 * 1024; + const std::string index_name{"test_hnsw_buffer_eviction.index"}; + zvec::test_util::RemoveTestFiles(index_name + "*"); + + auto param = HNSWIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithEFConstruction(100) + .Build(); + + { + auto write_index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, write_index); + ASSERT_EQ(0, write_index->Open(index_name, + {StorageOptions::StorageType::kMMAP, true})); + + std::vector vector(kDimension); + VectorData vector_data; + for (uint32_t id = 0; id < kDocCount; ++id) { + std::fill(vector.begin(), vector.end(), static_cast(id)); + vector_data.vector = DenseVector{vector.data()}; + ASSERT_EQ(0, write_index->Add(vector_data, id)); + } + ASSERT_EQ(0, write_index->Flush()); + ASSERT_EQ(0, write_index->Close()); + } + + // A 768-D FP32 HNSW node is larger than 3 KiB, so this index is larger + // than the 1 MiB pool and many vector reads cross a 4 KiB page boundary. + ASSERT_EQ(0, + zvec::ailego::MemoryLimitPool::get_instance().init(kBufferBudget)); + { + auto read_index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, read_index); + ASSERT_EQ(0, read_index->Open( + index_name, + {StorageOptions::StorageType::kBufferPool, false, true})); + auto *hnsw_index = dynamic_cast(read_index.get()); + ASSERT_NE(nullptr, hnsw_index); + ASSERT_EQ("buffer_pool", hnsw_index->storage_mode()); + + auto query_param = + HNSWQueryParamBuilder().with_topk(1).with_ef_search(100).build(); + std::vector query_vector(kDimension); + VectorData query; + for (uint32_t id : {0U, 63U, 127U, 255U, 383U, 511U}) { + std::fill(query_vector.begin(), query_vector.end(), + static_cast(id)); + query.vector = DenseVector{query_vector.data()}; + SearchResult result; + ASSERT_EQ(0, read_index->Search(query, query_param, &result)); + ASSERT_EQ(1U, result.doc_list_.size()); + ASSERT_EQ(id, result.doc_list_[0].key()); + } + ASSERT_EQ(0, read_index->Close()); + } + + zvec::test_util::RemoveTestFiles(index_name + "*"); + ASSERT_EQ( + 0, zvec::ailego::MemoryLimitPool::get_instance().init(100 * 1024 * 1024)); +} + TEST(IndexInterface, SparseGeneral) { constexpr uint32_t kSparseCount = 3; @@ -750,10 +817,9 @@ TEST(IndexInterface, Merge) { auto index3 = create_index_func(param_target, index_name + "3"); ASSERT_NE(nullptr, index3); MergeOptions merge_options; - merge_options.write_concurrency = - (std::numeric_limits::max)(); - ASSERT_TRUE(0 == index3->Merge({index1, index2}, IndexFilter(), - merge_options)); + merge_options.write_concurrency = (std::numeric_limits::max)(); + ASSERT_TRUE( + 0 == index3->Merge({index1, index2}, IndexFilter(), merge_options)); ASSERT_TRUE(3 == index3->GetDocCount()); { VectorDataBuffer fetched_vector_data; @@ -784,11 +850,9 @@ TEST(IndexInterface, Merge) { filter.set([](uint64_t key) { return key == 0; }); // TODO: uint32? zvec::ailego::ThreadPool pool(1, false); MergeOptions merge_options; - merge_options.write_concurrency = - (std::numeric_limits::max)(); + merge_options.write_concurrency = (std::numeric_limits::max)(); merge_options.pool = &pool; - ASSERT_TRUE(0 == - index3->Merge({index1, index2}, filter, merge_options)); + ASSERT_TRUE(0 == index3->Merge({index1, index2}, filter, merge_options)); ASSERT_TRUE(2 == index3->GetDocCount()); { VectorDataBuffer fetched_vector_data; @@ -2314,8 +2378,8 @@ TEST(IndexInterface, ExternalVectorFastSearchRecallRegression) { exact_results.reserve(kNumVectors); for (uint32_t i = 0; i < kNumVectors; ++i) { const float *vector = all_vectors.data() + i * kDimension; - const float score = std::inner_product( - query_vector.begin(), query_vector.end(), vector, 0.0f); + const float score = std::inner_product(query_vector.begin(), + query_vector.end(), vector, 0.0f); exact_results.emplace_back(score, i); } std::sort(exact_results.begin(), exact_results.end(), @@ -2340,13 +2404,11 @@ TEST(IndexInterface, ExternalVectorFastSearchRecallRegression) { .Build(); auto index = IndexFactory::CreateAndInitIndex(*param); ASSERT_NE(nullptr, index); - ASSERT_EQ(0, - index->Open(index_name, - {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ( + 0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true})); for (uint32_t i = 0; i < kNumVectors; ++i) { - VectorData vector_data{ - DenseVector{all_vectors.data() + i * kDimension}}; + VectorData vector_data{DenseVector{all_vectors.data() + i * kDimension}}; ASSERT_EQ(0, index->AddWithSource(vector_data, i, source)); } diff --git a/tests/core/utility/buffer_read_storage_test.cc b/tests/core/utility/buffer_read_storage_test.cc new file mode 100644 index 000000000..4bad7f7fe --- /dev/null +++ b/tests/core/utility/buffer_read_storage_test.cc @@ -0,0 +1,332 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include "utility/utility_params.h" + +using namespace zvec; +using namespace zvec::core; + +namespace { + +constexpr size_t kPoolSize = 16UL * 1024UL * 1024UL; + +class BufferReadStorageTest : public testing::Test { + protected: + void SetUp() override { + ailego::MemoryLimitPool::get_instance().init(kPoolSize); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(dumper, nullptr); + ASSERT_EQ(0, dumper->create(file_path_)); + + // Four native pages force read_range's bulk cold-read path on every + // platform (macOS commonly uses 16 KiB pages, Linux 4 KiB). + payload_.resize(4UL * ailego::kVectorPageSize); + for (size_t i = 0; i < payload_.size(); ++i) { + payload_[i] = static_cast(i % 251); + } + ASSERT_EQ(payload_.size(), dumper->write(payload_.data(), payload_.size())); + ASSERT_EQ(0, dumper->append("payload", payload_.size(), 0, 0)); + ASSERT_EQ(0, dumper->close()); + } + + void TearDown() override { + std::remove(file_path_.c_str()); + } + + IndexStorage::Pointer CreateStorage(const std::string &warmup_mode) { + auto storage = IndexFactory::CreateStorage("BufferReadStorage"); + EXPECT_NE(storage, nullptr); + if (!storage) { + return nullptr; + } + + ailego::Params params; + params.set(BUFFER_READ_STORAGE_WARMUP_MODE, warmup_mode); + EXPECT_EQ(0, storage->init(params)); + return storage; + } + + std::string file_path_{"buffer_read_storage_warmup_test_file"}; + std::string payload_; +}; + +TEST_F(BufferReadStorageTest, NoneDefersPagePopulationUntilFirstRead) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + + auto &pool = ailego::MemoryLimitPool::get_instance(); + EXPECT_EQ(0u, pool.stats().page_used); + EXPECT_NE(nullptr, storage->vec_buffer_pool()); + + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + std::string actual(payload_.size(), '\0'); + ASSERT_EQ(actual.size(), segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(payload_, actual); + EXPECT_GT(pool.stats().page_used, 0u); +} + +TEST_F(BufferReadStorageTest, SequentialPreservesExistingWarmupBehavior) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + + EXPECT_GT(ailego::MemoryLimitPool::get_instance().stats().page_used, 0u); + + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + std::string actual(payload_.size(), '\0'); + ASSERT_EQ(actual.size(), segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(payload_, actual); +} + +TEST_F(BufferReadStorageTest, RejectsUnknownWarmupMode) { + auto storage = IndexFactory::CreateStorage("BufferReadStorage"); + ASSERT_NE(storage, nullptr); + + ailego::Params params; + params.set(BUFFER_READ_STORAGE_WARMUP_MODE, "unknown"); + EXPECT_EQ(IndexError_InvalidArgument, storage->init(params)); +} + +TEST_F(BufferReadStorageTest, PoolSmallerThanOnePageFallsBackToBypass) { + const size_t kTooSmall = ailego::kVectorPageSize - 1; + ASSERT_EQ(0, ailego::MemoryLimitPool::get_instance().init(kTooSmall)); + + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().metadata_used); + EXPECT_EQ(nullptr, storage->vec_buffer_pool()); + + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + std::string actual(payload_.size(), '\0'); + ASSERT_EQ(actual.size(), segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(payload_, actual); +} + +TEST_F(BufferReadStorageTest, PoolWithoutRoomForMetadataFallsBackToBypass) { + ASSERT_EQ( + 0, ailego::MemoryLimitPool::get_instance().init(ailego::kVectorPageSize)); + + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().metadata_used); + + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + IndexStorage::MemoryBlock block; + ASSERT_EQ(64u, segment->read(0, block, 64)); + EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); +} + +TEST_F(BufferReadStorageTest, CachePressureFallsBackToOwnedRead) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + + auto &pool = ailego::MemoryLimitPool::get_instance(); + const size_t external_charge = pool.available(); + ASSERT_GT(external_charge, 0u); + ASSERT_TRUE(pool.try_charge_external(external_charge)); + + IndexStorage::MemoryBlock block; + ASSERT_EQ(64u, segment->read(0, block, 64)); + EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + pool.release_external(external_charge); +} + +TEST_F(BufferReadStorageTest, MemoryBlockKeepsPoolAliveAfterStorageClose) { + IndexStorage::MemoryBlock block; + IndexStorage::MemoryBlock copy; + { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + ASSERT_EQ(64u, segment->read(0, block, 64)); + ASSERT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, block.type_); + EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); + copy = block; + + ASSERT_EQ(0, storage->close()); + segment.reset(); + storage.reset(); + EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); + } + block.reset(); + EXPECT_EQ(0, std::memcmp(payload_.data(), copy.data(), 64)); + copy.reset(); + EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().page_used); +} + +TEST_F(BufferReadStorageTest, BorrowedReadAvoidsOwningHandleOnResidentPage) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + + IndexStorage::MemoryBlock block; + ASSERT_EQ(64u, segment->read_borrowed(0, block, 64)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, block.type_); + EXPECT_EQ(nullptr, block.buffer_pool_handle_owner_); + EXPECT_NE(nullptr, block.buffer_pool_handle_); + EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); + + // Borrowed blocks must be released while their Segment/storage owner is + // still alive. + block.reset(); +} + +TEST_F(BufferReadStorageTest, PointerReadPinsUntilNextPointerRead) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + const void *data = nullptr; + ASSERT_EQ(64u, segment->read(0, &data, 64)); + ASSERT_NE(data, nullptr); + EXPECT_EQ(0, std::memcmp(payload_.data(), data, 64)); + auto &pool = ailego::MemoryLimitPool::get_instance(); + const size_t external_charge = pool.available(); + ASSERT_GT(external_charge, 0u); + ASSERT_TRUE(pool.try_charge_external(external_charge)); + EXPECT_FALSE(pool.try_charge_external(1)); + + IndexStorage::SegmentData copied(0, 64); + ASSERT_TRUE(segment->read(&copied, 1)); + EXPECT_EQ(0, std::memcmp(payload_.data(), copied.data, 64)); + EXPECT_TRUE(pool.try_charge_external(1)); + pool.release_external(1); + pool.release_external(external_charge); +} + +TEST_F(BufferReadStorageTest, PointerReadsUsePerThreadScratchBuffers) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + + constexpr size_t kReadSize = 64; + constexpr size_t kOffsets[] = {0, 127}; + std::atomic ready{0}; + std::atomic failed{false}; + std::array readers; + for (size_t i = 0; i < readers.size(); ++i) { + readers[i] = std::thread([&, i] { + const void *data = nullptr; + if (segment->read(kOffsets[i], &data, kReadSize) != kReadSize || + data == nullptr) { + failed.store(true, std::memory_order_release); + ready.fetch_add(1, std::memory_order_release); + return; + } + ready.fetch_add(1, std::memory_order_release); + while (ready.load(std::memory_order_acquire) != readers.size()) { + std::this_thread::yield(); + } + if (std::memcmp(data, payload_.data() + kOffsets[i], kReadSize) != 0) { + failed.store(true, std::memory_order_release); + } + }); + } + for (auto &reader : readers) { + reader.join(); + } + EXPECT_FALSE(failed.load(std::memory_order_acquire)); +} + +TEST_F(BufferReadStorageTest, MissingFileReturnsErrorWithoutThrowing) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + + const std::string missing_path = file_path_ + ".missing"; + std::remove(missing_path.c_str()); + EXPECT_EQ(IndexError_OpenFile, storage->open(missing_path, false)); + EXPECT_TRUE(storage->file_path().empty()); + EXPECT_FALSE(storage->has("payload")); +} + +TEST_F(BufferReadStorageTest, FailedReopenPreservesPublishedState) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + + auto published_segment = storage->get("payload"); + ASSERT_NE(published_segment, nullptr); + const std::string missing_path = file_path_ + ".missing"; + std::remove(missing_path.c_str()); + EXPECT_EQ(IndexError_OpenFile, storage->open(missing_path, false)); + + EXPECT_EQ(file_path_, storage->file_path()); + EXPECT_TRUE(storage->has("payload")); + std::array actual{}; + EXPECT_EQ(actual.size(), + published_segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(0, std::memcmp(payload_.data(), actual.data(), actual.size())); +} + +TEST_F(BufferReadStorageTest, RejectsOutOfRangeContainerOffset) { + auto storage = IndexFactory::CreateStorage("BufferReadStorage"); + ASSERT_NE(storage, nullptr); + + ailego::Params params; + params.set(BUFFER_READ_STORAGE_WARMUP_MODE, BUFFER_READ_STORAGE_WARMUP_NONE); + params.set(BUFFER_READ_STORAGE_HEADER_OFFSET, + std::numeric_limits::min()); + ASSERT_EQ(0, storage->init(params)); + EXPECT_EQ(IndexError_InvalidArgument, storage->open(file_path_, false)); +} + +TEST_F(BufferReadStorageTest, RangeChecksDoNotOverflow) { + auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(file_path_, false)); + auto segment = storage->get("payload"); + ASSERT_NE(segment, nullptr); + + std::string actual(payload_.size() - 1, '\0'); + EXPECT_EQ(actual.size(), segment->fetch(1, actual.data(), + std::numeric_limits::max())); + EXPECT_EQ(payload_.substr(1), actual); + + IndexStorage::SegmentData invalid(std::numeric_limits::max(), 2); + EXPECT_FALSE(segment->read(&invalid, 1)); + segment->prefetch(std::numeric_limits::max(), + std::numeric_limits::max()); +} + +} // namespace diff --git a/tests/core/utility/buffer_storage_write_test.cc b/tests/core/utility/buffer_storage_write_test.cc index 894c68f1e..78a52f325 100644 --- a/tests/core/utility/buffer_storage_write_test.cc +++ b/tests/core/utility/buffer_storage_write_test.cc @@ -42,7 +42,9 @@ class BufferStorageWriteTest : public ::testing::Test { ailego::File::MakePath("buffer_storage_write_test_dir"); } - void TearDown() override { ailego::File::Delete(file_path_); } + void TearDown() override { + ailego::File::Delete(file_path_); + } // Open BufferStorage in writable mode (create_if_missing=true) IndexStorage::Pointer OpenWritable() { @@ -69,7 +71,8 @@ class BufferStorageWriteTest : public ::testing::Test { // ===== Basic Write Tests ===== -// Test: Create new index via BufferStorage, append segment, write data, read back +// Test: Create new index via BufferStorage, append segment, write data, read +// back TEST_F(BufferStorageWriteTest, WriteBasicCreateAndWrite) { auto storage = OpenWritable(); ASSERT_TRUE(storage); @@ -184,6 +187,52 @@ TEST_F(BufferStorageWriteTest, WriteOverwrite) { EXPECT_EQ(0, storage->close()); } +TEST_F(BufferStorageWriteTest, WriteBatchSamePagePersists) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + + ASSERT_EQ(0, storage->append("batch_seg", 8192)); + auto segment = storage->get("batch_seg"); + ASSERT_TRUE(segment); + + const size_t base_in_page = segment->data_offset() % ailego::kVectorPageSize; + const size_t offset = base_in_page + 128 <= ailego::kVectorPageSize + ? 0 + : ailego::kVectorPageSize - base_in_page; + const uint32_t payload = 0x12345678u; + const uint32_t published_count = 7; + const IndexStorage::SegmentData writes[] = { + {offset + 64, sizeof(payload), &payload}, + {offset, sizeof(published_count), &published_count}, + }; + ASSERT_TRUE(segment->write_batch(writes, 2)); + EXPECT_EQ(offset + 64 + sizeof(payload), segment->data_size()); + + uint32_t got_count = 0; + uint32_t got_payload = 0; + EXPECT_EQ(sizeof(got_count), + segment->fetch(offset, &got_count, sizeof(got_count))); + EXPECT_EQ(sizeof(got_payload), + segment->fetch(offset + 64, &got_payload, sizeof(got_payload))); + EXPECT_EQ(published_count, got_count); + EXPECT_EQ(payload, got_payload); + + ASSERT_EQ(0, storage->close()); + storage = OpenReadOnly(); + ASSERT_TRUE(storage); + segment = storage->get("batch_seg"); + ASSERT_TRUE(segment); + got_count = 0; + got_payload = 0; + EXPECT_EQ(sizeof(got_count), + segment->fetch(offset, &got_count, sizeof(got_count))); + EXPECT_EQ(sizeof(got_payload), + segment->fetch(offset + 64, &got_payload, sizeof(got_payload))); + EXPECT_EQ(published_count, got_count); + EXPECT_EQ(payload, got_payload); + EXPECT_EQ(0, storage->close()); +} + // ===== Boundary / Error Tests ===== // Test: Write exceeding segment capacity returns 0 @@ -275,8 +324,7 @@ TEST_F(BufferStorageWriteTest, WriteMultipleFlushCycles) { EXPECT_EQ(0, storage->flush()); // Second write at a different offset + flush - EXPECT_EQ(data2.size(), - seg->write(200, data2.data(), data2.size())); + EXPECT_EQ(data2.size(), seg->write(200, data2.data(), data2.size())); EXPECT_EQ(0, storage->flush()); EXPECT_EQ(0, storage->close()); } @@ -353,8 +401,7 @@ TEST_F(BufferStorageWriteTest, WriteReadOnlyNoOp) { std::string new_data = "overwrite_attempt"; // Should return len (silent no-op) - EXPECT_EQ(new_data.size(), - seg->write(0, new_data.data(), new_data.size())); + EXPECT_EQ(new_data.size(), seg->write(0, new_data.data(), new_data.size())); // Data should remain unchanged (still "initial") std::vector buf(7); @@ -881,7 +928,8 @@ TEST_F(BufferStorageWriteTest, CR_ConcurrentWriteAndResize) { // chain split. After reopen, ALL segments must be findable. // (Tests fix for reserve()-induced dangling pointer in append_segment.) TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { - const int kNumSegments = 50; // Enough to trigger chain split with default 4096 meta capacity + const int kNumSegments = + 50; // Enough to trigger chain split with default 4096 meta capacity { auto storage = OpenWritable(); @@ -892,7 +940,8 @@ TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { ASSERT_EQ(0, storage->append(name, 4096)) << "Failed to append segment " << i; auto seg = storage->get(name); - ASSERT_TRUE(seg) << "Failed to get segment " << name << " right after append"; + ASSERT_TRUE(seg) << "Failed to get segment " << name + << " right after append"; // Write a marker so we can verify on reopen std::string marker = "marker_" + std::to_string(i); EXPECT_EQ(marker.size(), seg->write(0, marker.data(), marker.size())); @@ -908,7 +957,8 @@ TEST_F(BufferStorageWriteTest, CR_ChainSplitAllSegmentsAccessible) { for (int i = 0; i < kNumSegments; ++i) { std::string name = "chain_seg_" + std::to_string(i); auto seg = storage->get(name); - ASSERT_TRUE(seg) << "Segment " << name << " missing after reopen (chain-split bug?)"; + ASSERT_TRUE(seg) << "Segment " << name + << " missing after reopen (chain-split bug?)"; std::string expected = "marker_" + std::to_string(i); std::vector buf(expected.size()); EXPECT_EQ(expected.size(), seg->fetch(0, buf.data(), buf.size())); @@ -991,13 +1041,8 @@ TEST_F(BufferStorageWriteTest, CR_CrossPageWriteAndRead) { EXPECT_EQ(kWriteLen, seg->fetch(kWriteOffset, fetch_buf.data(), kWriteLen)); EXPECT_EQ(write_data, fetch_buf); - // Read back via read(MemoryBlock&) - exercises the cross-page alloc path. - // Scope the MemoryBlock so it is destroyed BEFORE storage->close(): - // when the read happens to land on a single page (e.g. macOS arm64 with - // 16KB pages, where [2000, 7000) fits in one page) the returned block - // is MBT_BUFFERPOOL holding a raw pointer to buffer_pool_handle_. Once - // close_index() resets buffer_pool_handle_/buffer_pool_, that raw - // pointer dangles and ~MemoryBlock()'s release_one() segfaults. + // Read back via read(MemoryBlock&) - exercises the cross-page alloc path on + // 4K hosts and the writable-pool snapshot path on large-page hosts. { IndexStorage::MemoryBlock mb; EXPECT_EQ(kWriteLen, seg->read(kWriteOffset, mb, kWriteLen)); @@ -1007,6 +1052,264 @@ TEST_F(BufferStorageWriteTest, CR_CrossPageWriteAndRead) { EXPECT_EQ(0, storage->close()); } +TEST_F(BufferStorageWriteTest, ImmutableReadPinsWritablePageWithoutCopy) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + + ASSERT_EQ(0, storage->append("immutable_read_seg", 8192)); + auto segment = storage->get("immutable_read_seg"); + ASSERT_TRUE(segment); + const size_t page_offset = segment->data_offset() % ailego::kVectorPageSize; + const size_t offset = page_offset + 256 <= ailego::kVectorPageSize + ? 0 + : ailego::kVectorPageSize - page_offset; + const std::string expected = "immutable-vector-bytes"; + ASSERT_EQ(expected.size(), + segment->write(offset, expected.data(), expected.size())); + + IndexStorage::MemoryBlock snapshot; + ASSERT_EQ(expected.size(), segment->read(offset, snapshot, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, snapshot.type_); + + IndexStorage::MemoryBlock pinned; + ASSERT_EQ(expected.size(), + segment->read_immutable(offset, pinned, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, pinned.type_); + EXPECT_TRUE(pinned.buffer_pool_handle_owner_); + EXPECT_EQ(expected, std::string(static_cast(pinned.data()), + expected.size())); + + IndexStorage::MemoryBlock borrowed; + ASSERT_EQ(expected.size(), segment->read_borrowed_immutable(offset, borrowed, + expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, borrowed.type_); + EXPECT_FALSE(borrowed.buffer_pool_handle_owner_); + + borrowed.reset(); + pinned.reset(); + snapshot.reset(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, ImmutableReadCopiesCrossPageRange) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + + const size_t page_size = ailego::kVectorPageSize; + ASSERT_EQ(0, storage->append("immutable_cross_page", 3 * page_size)); + auto segment = storage->get("immutable_cross_page"); + ASSERT_TRUE(segment); + + const size_t aligned_offset = + (page_size - segment->data_offset() % page_size) % page_size; + const size_t offset = aligned_offset + page_size - 128; + std::vector expected(3072); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = static_cast((i * 17 + 11) % 251); + } + ASSERT_EQ(expected.size(), + segment->write(offset, expected.data(), expected.size())); + + IndexStorage::MemoryBlock snapshot; + ASSERT_EQ(expected.size(), + segment->read_immutable(offset, snapshot, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, snapshot.type_); + EXPECT_EQ(0, std::memcmp(expected.data(), snapshot.data(), expected.size())); + snapshot.reset(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, ImmutableBatchReadUsesWritableCache) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + const size_t page_size = ailego::kVectorPageSize; + ASSERT_EQ(0, storage->append("immutable_batch", 3 * page_size)); + auto segment = storage->get("immutable_batch"); + ASSERT_TRUE(segment); + + EXPECT_FALSE(segment->prefer_borrowed_batch_for( + std::max(1, page_size / 8) - 1)); + EXPECT_TRUE( + segment->prefer_borrowed_batch_for(std::max(1, page_size / 8))); + + std::vector payload(3 * page_size); + for (size_t i = 0; i < payload.size(); ++i) { + payload[i] = static_cast((i * 31 + 7) % 251); + } + ASSERT_EQ(payload.size(), segment->write(0, payload.data(), payload.size())); + + const size_t aligned_offset = + (page_size - segment->data_offset() % page_size) % page_size; + const size_t cross_offset = aligned_offset + page_size - 64; + IndexStorage::MemoryBlock blocks[2]; + IndexStorage::Segment::BorrowedRead reads[] = { + {segment.get(), aligned_offset, 128, &blocks[0]}, + {segment.get(), cross_offset, 128, &blocks[1]}, + }; + ASSERT_TRUE(segment->read_borrowed_batch_immutable(reads, 2)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, blocks[0].type_); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_MMAP, blocks[1].type_); + EXPECT_EQ(0, std::memcmp(payload.data() + aligned_offset, blocks[0].data(), + reads[0].length)); + EXPECT_EQ(0, std::memcmp(payload.data() + cross_offset, blocks[1].data(), + reads[1].length)); + blocks[0].reset(); + blocks[1].reset(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, ReadOnlyBatchFallsBackWhenPinsExceedBudget) { + constexpr size_t kSegmentBytes = 68UL * 1024UL * 1024UL; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("batch_pressure", kSegmentBytes)); + auto segment = storage->get("batch_pressure"); + ASSERT_TRUE(segment); + const char marker = 'Z'; + ASSERT_EQ(1U, segment->write(kSegmentBytes - 1, &marker, 1)); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto segment = storage->get("batch_pressure"); + ASSERT_TRUE(segment); + + const size_t page_size = ailego::kVectorPageSize; + const size_t first = + (page_size - segment->data_offset() % page_size) % page_size; + const size_t count = (segment->data_size() - first) / page_size; + ASSERT_GT(count * page_size, 64UL * 1024UL * 1024UL); + + // Repeat the final (necessarily cold) page to verify fallback is deduplicated + // by page rather than issued once per vector occurrence. + constexpr size_t kDuplicateTailReads = 2; + std::vector blocks(count + kDuplicateTailReads); + std::vector reads; + reads.reserve(blocks.size()); + for (size_t i = 0; i < count; ++i) { + reads.emplace_back(segment.get(), first + i * page_size, 1, &blocks[i]); + } + for (size_t i = 0; i < kDuplicateTailReads; ++i) { + reads.emplace_back(segment.get(), first + (count - 1) * page_size, 1, + &blocks[count + i]); + } + + auto pool = storage->vec_buffer_pool(); + ASSERT_TRUE(pool); + const auto before = pool->stats(); + ASSERT_TRUE(segment->read_borrowed_batch_immutable(reads.data(), + reads.size())); + size_t cached_unique = 0; + size_t bypassed_unique = 0; + for (size_t i = 0; i < count; ++i) { + const auto &block = blocks[i]; + if (block.type_ == IndexStorage::MemoryBlock::MBT_BUFFERPOOL) { + ++cached_unique; + } else { + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_MMAP, block.type_); + ++bypassed_unique; + } + ASSERT_NE(nullptr, block.data()); + EXPECT_EQ(0, *static_cast(block.data())); + } + EXPECT_GT(cached_unique, 0U); + EXPECT_GT(bypassed_unique, 0U); + for (size_t i = 0; i < kDuplicateTailReads; ++i) { + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_MMAP, blocks[count + i].type_); + ASSERT_NE(nullptr, blocks[count + i].data()); + EXPECT_EQ(0, + *static_cast(blocks[count + i].data())); + } + const auto after = pool->stats(); + EXPECT_EQ(bypassed_unique, after.bypass_reads - before.bypass_reads); + blocks.clear(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, ReadOnlyBatchPreservesSharedCacheReserve) { + constexpr size_t kSegmentBytes = 64UL * 1024UL * 1024UL; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("batch_reserve", kSegmentBytes)); + auto segment = storage->get("batch_reserve"); + ASSERT_TRUE(segment); + const char marker = 'R'; + ASSERT_EQ(1U, segment->write(kSegmentBytes - 1, &marker, 1)); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto segment = storage->get("batch_reserve"); + ASSERT_TRUE(segment); + + const size_t page_size = ailego::kVectorPageSize; + const size_t first = + (page_size - segment->data_offset() % page_size) % page_size; + constexpr size_t kSharedReserveBytes = 4UL * 1024UL * 1024UL; + const size_t available_before = + ailego::MemoryLimitPool::get_instance().available(); + ASSERT_GT(available_before, kSharedReserveBytes); + const size_t count = + (available_before - kSharedReserveBytes) / page_size + 1; + ASSERT_LT(first + count * page_size, segment->data_size()); + + std::vector blocks(count); + std::vector reads; + reads.reserve(count); + for (size_t i = 0; i < count; ++i) { + reads.emplace_back(segment.get(), first + i * page_size, 1, &blocks[i]); + } + + ASSERT_TRUE(segment->read_borrowed_batch_immutable(reads.data(), + reads.size())); + EXPECT_GE(ailego::MemoryLimitPool::get_instance().available() + page_size, + kSharedReserveBytes); + blocks.clear(); + EXPECT_EQ(0, storage->close()); +} + +// Repeated legacy pointer reads from a writable cached page must reuse the +// pinned page. Retaining a separate 4K-aligned snapshot for every read grows +// memory until close() and makes long Optimize workloads consume gigabytes. +TEST_F(BufferStorageWriteTest, CR_WritableLegacyPointerReadReusesCachedPage) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + + ASSERT_EQ(0, storage->append("legacy_pointer_seg", 8192)); + auto seg = storage->get("legacy_pointer_seg"); + ASSERT_TRUE(seg); + + constexpr size_t kReadLen = 64; + const size_t data_offset_in_page = + seg->data_offset() % ailego::kVectorPageSize; + const size_t read_offset = + (ailego::kVectorPageSize - data_offset_in_page) % ailego::kVectorPageSize; + ASSERT_LE(read_offset + kReadLen, seg->capacity()); + + std::vector expected(kReadLen, 'P'); + ASSERT_EQ(kReadLen, + seg->write(read_offset, expected.data(), expected.size())); + + const void *first = nullptr; + ASSERT_EQ(kReadLen, seg->read(read_offset, &first, kReadLen)); + ASSERT_NE(nullptr, first); + EXPECT_EQ(0, std::memcmp(expected.data(), first, expected.size())); + + for (size_t i = 0; i < 128; ++i) { + const void *again = nullptr; + ASSERT_EQ(kReadLen, seg->read(read_offset, &again, kReadLen)); + EXPECT_EQ(first, again); + } + + EXPECT_EQ(0, storage->close()); +} + // Dirty flag race: write() after flush_index() must re-set the dirty flag. // If the write lands between CAS(dirty, false) and the end of flush, // the next flush must still persist it. Verified by close→reopen→read. @@ -1047,8 +1350,9 @@ TEST_F(BufferStorageWriteTest, CR_DirtyFlagNotLostAfterFlush) { } } -// Stress test: Concurrent flush + write interleaving to expose dirty flag races. -// All writes that return successfully MUST be visible after final close+reopen. +// Stress test: Concurrent flush + write interleaving to expose dirty flag +// races. All writes that return successfully MUST be visible after final +// close+reopen. TEST_F(BufferStorageWriteTest, CR_ConcurrentFlushWriteDirtyFlagStress) { auto storage = OpenWritable(); ASSERT_TRUE(storage); @@ -1112,7 +1416,8 @@ TEST_F(BufferStorageWriteTest, CR_PointerStabilityAcrossAppend) { // Write initial data std::string initial = "before_append"; - EXPECT_EQ(initial.size(), seg_first->write(0, initial.data(), initial.size())); + EXPECT_EQ(initial.size(), + seg_first->write(0, initial.data(), initial.size())); // Append many more segments (may trigger internal rehash/resize) for (int i = 0; i < 20; ++i) { @@ -1178,3 +1483,294 @@ TEST_F(BufferStorageWriteTest, CR_ConcurrentWriteAndCrcUpdate) { (void)seg->data_crc(); EXPECT_EQ(0, storage->close()); } + +TEST_F(BufferStorageWriteTest, CR_ConcurrentSamePageReadSeesWholeWrite) { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg1", 2 * ailego::kVectorPageSize)); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + + constexpr size_t kPayloadSize = 256; + const size_t page_offset = seg->data_offset() % ailego::kVectorPageSize; + const size_t offset = page_offset + kPayloadSize <= ailego::kVectorPageSize + ? 0 + : ailego::kVectorPageSize - page_offset; + std::vector pattern_a(kPayloadSize, static_cast(0x55)); + std::vector pattern_b(kPayloadSize, static_cast(0xAA)); + ASSERT_EQ(kPayloadSize, + seg->write(offset, pattern_a.data(), pattern_a.size())); + + std::atomic start{false}; + std::atomic failed{false}; + std::thread writer([&] { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + for (size_t i = 0; i < 10000; ++i) { + const auto &pattern = (i & 1U) == 0 ? pattern_b : pattern_a; + if (seg->write(offset, pattern.data(), pattern.size()) != kPayloadSize) { + failed.store(true, std::memory_order_release); + return; + } + } + }); + std::thread reader([&] { + std::vector observed(kPayloadSize); + start.store(true, std::memory_order_release); + for (size_t i = 0; i < 10000; ++i) { + if (seg->fetch(offset, observed.data(), observed.size()) != + kPayloadSize || + (observed != pattern_a && observed != pattern_b)) { + failed.store(true, std::memory_order_release); + return; + } + } + }); + writer.join(); + reader.join(); + EXPECT_FALSE(failed.load(std::memory_order_acquire)); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, CR_MemoryBlockOutlivesReadOnlyStorage) { + const std::string expected = "lifetime"; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg1", 4096)); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + ASSERT_EQ(expected.size(), seg->write(0, expected.data(), expected.size())); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + IndexStorage::MemoryBlock block; + { + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + ASSERT_EQ(expected.size(), seg->read(0, block, expected.size())); + ASSERT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, block.type_); + ASSERT_EQ(0, storage->close()); + seg.reset(); + storage.reset(); + EXPECT_EQ(0, std::memcmp(expected.data(), block.data(), expected.size())); + } + block.reset(); +} + +TEST_F(BufferStorageWriteTest, CR_BorrowedReadAvoidsOwningHandle) { + const std::string expected = "borrowed"; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg1", 4096)); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + ASSERT_EQ(expected.size(), seg->write(0, expected.data(), expected.size())); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + IndexStorage::MemoryBlock block; + ASSERT_EQ(expected.size(), seg->read_borrowed(0, block, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, block.type_); + EXPECT_EQ(nullptr, block.buffer_pool_handle_owner_); + EXPECT_NE(nullptr, block.buffer_pool_handle_); + EXPECT_EQ(0, std::memcmp(expected.data(), block.data(), expected.size())); + + block.reset(); + ASSERT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, BorrowedReadUsesOneColdPageLoadSequence) { + const std::string expected = "borrowed pressure fallback"; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg1", 4096)); + auto segment = storage->get("seg1"); + ASSERT_TRUE(segment); + ASSERT_EQ(expected.size(), + segment->write(0, expected.data(), expected.size())); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto segment = storage->get("seg1"); + ASSERT_TRUE(segment); + + auto &pool = ailego::MemoryLimitPool::get_instance(); + const size_t external_charge = pool.available(); + ASSERT_GT(external_charge, 0u); + ASSERT_TRUE(pool.try_charge_external(external_charge)); + const uint64_t high_watermark_hits_before = pool.stats().high_watermark_hits; + + IndexStorage::MemoryBlock block; + ASSERT_EQ(expected.size(), segment->read_borrowed(0, block, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + EXPECT_EQ(0, std::memcmp(expected.data(), block.data(), expected.size())); + // get_single_page() performs one initial attempt plus at most 50 retries. + EXPECT_LE(pool.stats().high_watermark_hits - high_watermark_hits_before, 51u); + + pool.release_external(external_charge); + ASSERT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, ReadOnlyPrefetchPreservesCachePriority) { + const size_t page_size = ailego::kVectorPageSize; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("hot", 3 * page_size)); + auto segment = storage->get("hot"); + ASSERT_TRUE(segment); + ASSERT_EQ(3 * page_size, segment->resize(3 * page_size)); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto segment = storage->get("hot"); + ASSERT_TRUE(segment); + auto pool = storage->vec_buffer_pool(); + ASSERT_TRUE(pool); + + const size_t first_page = segment->data_offset() / page_size; + segment->prefetch(0, 1, IndexStorage::Segment::CachePriority::kHigh); + EXPECT_TRUE(pool->is_page_resident(first_page)); + EXPECT_EQ(ailego::VecBufferPool::kHighPriority, + pool->page_table_.eviction_priority(first_page)); + + // A later, colder hint must not demote an already protected page. + segment->prefetch(0, 1, IndexStorage::Segment::CachePriority::kNormal); + EXPECT_EQ(ailego::VecBufferPool::kHighPriority, + pool->page_table_.eviction_priority(first_page)); + + const size_t second_page_offset = + page_size - (segment->data_offset() % page_size); + const size_t second_page = + (segment->data_offset() + second_page_offset) / page_size; + segment->prefetch(second_page_offset, 1, + IndexStorage::Segment::CachePriority::kNormal); + EXPECT_TRUE(pool->is_page_resident(second_page)); + EXPECT_EQ(ailego::VecBufferPool::kNormalPriority, + pool->page_table_.eviction_priority(second_page)); + + ASSERT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, BatchBorrowedReadAcrossSegments) { + const size_t page_size = ailego::kVectorPageSize; + std::vector payload_a(3 * page_size); + std::vector payload_b(2 * page_size); + for (size_t i = 0; i < payload_a.size(); ++i) { + payload_a[i] = static_cast((i * 17 + 3) % 251); + } + for (size_t i = 0; i < payload_b.size(); ++i) { + payload_b[i] = static_cast((i * 29 + 11) % 251); + } + + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg_a", payload_a.size())); + ASSERT_EQ(0, storage->append("seg_b", payload_b.size())); + auto seg_a = storage->get("seg_a"); + auto seg_b = storage->get("seg_b"); + ASSERT_TRUE(seg_a); + ASSERT_TRUE(seg_b); + ASSERT_EQ(payload_a.size(), + seg_a->write(0, payload_a.data(), payload_a.size())); + ASSERT_EQ(payload_b.size(), + seg_b->write(0, payload_b.data(), payload_b.size())); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto seg_a = storage->get("seg_a"); + auto seg_b = storage->get("seg_b"); + ASSERT_TRUE(seg_a); + ASSERT_TRUE(seg_b); + + const size_t a_page_aligned_offset = + (page_size - seg_a->data_offset() % page_size) % page_size; + const size_t a_cross_page_offset = + (2 * page_size - 32 - seg_a->data_offset() % page_size) % page_size; + const size_t b_page_aligned_offset = + (page_size - seg_b->data_offset() % page_size) % page_size; + ASSERT_LE(a_page_aligned_offset + 64, payload_a.size()); + ASSERT_LE(a_cross_page_offset + 128, payload_a.size()); + ASSERT_LE(b_page_aligned_offset + 96, payload_b.size()); + + IndexStorage::MemoryBlock blocks[4]; + IndexStorage::Segment::BorrowedRead reads[] = { + {seg_a.get(), a_page_aligned_offset, 64, &blocks[0]}, + {seg_a.get(), a_page_aligned_offset + 8, 32, &blocks[1]}, + {seg_a.get(), a_cross_page_offset, 128, &blocks[2]}, + {seg_b.get(), b_page_aligned_offset, 96, &blocks[3]}, + }; + ASSERT_TRUE(seg_a->read_borrowed_batch(reads, 4)); + + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, blocks[0].type_); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, blocks[1].type_); + // Cross-page reads are served from the reused thread-local scratch arena as + // non-owning views (ZVEC_CROSS_ARENA default on); the reassembled bytes are + // still validated below. + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_MMAP, blocks[2].type_); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_BUFFERPOOL, blocks[3].type_); + EXPECT_EQ(0, std::memcmp(payload_a.data() + a_page_aligned_offset, + blocks[0].data(), 64)); + EXPECT_EQ(0, std::memcmp(payload_a.data() + a_page_aligned_offset + 8, + blocks[1].data(), 32)); + EXPECT_EQ(0, std::memcmp(payload_a.data() + a_cross_page_offset, + blocks[2].data(), 128)); + EXPECT_EQ(0, std::memcmp(payload_b.data() + b_page_aligned_offset, + blocks[3].data(), 96)); + + for (auto &block : blocks) { + block.reset(); + } + ASSERT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, CR_ReadOnlyMetadataPressureFallsBackToBypass) { + const std::string expected = "bypass"; + { + auto storage = OpenWritable(); + ASSERT_TRUE(storage); + ASSERT_EQ(0, storage->append("seg1", 4096)); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + ASSERT_EQ(expected.size(), seg->write(0, expected.data(), expected.size())); + ASSERT_EQ(0, storage->flush()); + ASSERT_EQ(0, storage->close()); + } + + auto &pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ(0, pool.init(ailego::kVectorPageSize)); + { + auto storage = OpenReadOnly(); + ASSERT_TRUE(storage); + auto seg = storage->get("seg1"); + ASSERT_TRUE(seg); + IndexStorage::MemoryBlock block; + ASSERT_EQ(expected.size(), seg->read(0, block, expected.size())); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + EXPECT_EQ(0, std::memcmp(expected.data(), block.data(), expected.size())); + ASSERT_EQ(0, storage->close()); + } + ASSERT_EQ(0, pool.init(64UL * 1024UL * 1024UL)); +} diff --git a/tests/core/utility/memory_block_view_test.cc b/tests/core/utility/memory_block_view_test.cc new file mode 100644 index 000000000..b694de0be --- /dev/null +++ b/tests/core/utility/memory_block_view_test.cc @@ -0,0 +1,61 @@ +// 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 +#include + +using namespace zvec; +using namespace zvec::core; + +// MakeBorrowedView backs the cross-page scratch-arena path in BufferStorage: +// cross-page vectors are copied into a reused thread-local arena and exposed +// as non-owning views. The block must therefore free nothing and pin nothing +// on destruction, and copies/moves must keep aliasing the same buffer without +// taking ownership. A double-free here would trap under ASan. +TEST(MemoryBlockBorrowedView, IsNonOwning) { + auto *buf = new char[64]; + std::memset(buf, 0xAB, 64); + { + auto view = IndexStorage::MemoryBlock::MakeBorrowedView(buf); + EXPECT_EQ(static_cast(buf), view.data()); + + // Copy keeps a non-owning alias to the same buffer. + auto copy = view; + EXPECT_EQ(static_cast(buf), copy.data()); + + // Move keeps the alias too. + auto moved = std::move(copy); + EXPECT_EQ(static_cast(buf), moved.data()); + } // every view destroyed here; the backing buffer must NOT be freed. + + // Still readable, and safe to free exactly once by the sole owner. + EXPECT_EQ(static_cast(buf[0]), 0xABu); + delete[] buf; +} + +// A borrowed view over a slice of a shared buffer (as the arena hands out) +// must alias the exact slice and never free the shared backing storage. +TEST(MemoryBlockBorrowedView, AliasesArenaSlice) { + std::vector arena(256, 0); + arena[128] = 0x5A; + auto view = IndexStorage::MemoryBlock::MakeBorrowedView(arena.data() + 128); + ASSERT_EQ(static_cast(arena.data() + 128), view.data()); + EXPECT_EQ(0x5A, *static_cast(view.data())); + // Destroying the view leaves the arena intact for reuse. + view.reset(); + EXPECT_EQ(0x5A, arena[128]); +} diff --git a/tests/core/utility/visit_filter_test.cc b/tests/core/utility/visit_filter_test.cc new file mode 100644 index 000000000..d71c930e5 --- /dev/null +++ b/tests/core/utility/visit_filter_test.cc @@ -0,0 +1,51 @@ +// 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 "utility/visit_filter.h" +#include + +using zvec::core::VisitFilter; + +TEST(VisitFilterTest, DestroyIsIdempotentAndAllowsReinitialization) { + VisitFilter filter; + ASSERT_EQ(0, filter.init(VisitFilter::ByteMap, 128, 128, 0.01f)); + filter.set_visited(7); + EXPECT_TRUE(filter.visited(7)); + + filter.destroy(); + filter.destroy(); + EXPECT_EQ(VisitFilter::Default, filter.get_mode()); + + ASSERT_EQ(0, filter.init(VisitFilter::BitMap, 128, 128, 0.01f)); + filter.set_visited(11); + EXPECT_TRUE(filter.visited(11)); +} + +TEST(VisitFilterTest, DestructorOwnsEveryModeContext) { + for (int mode : + {VisitFilter::BloomFilter, VisitFilter::BitMap, VisitFilter::ByteMap}) { + VisitFilter filter; + ASSERT_EQ(0, filter.init(mode, 1024, 512, 0.01f)); + filter.set_visited(17); + EXPECT_TRUE(filter.visited(17)); + } +} + +TEST(VisitFilterTest, RejectsInvalidModeWithoutRetainingContext) { + VisitFilter filter; + EXPECT_EQ(zvec::core::IndexError_InvalidArgument, + filter.init(VisitFilter::Default, 128, 128, 0.01f)); + EXPECT_EQ(VisitFilter::Default, filter.get_mode()); + filter.destroy(); +} diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 5d9a3e43e..c3a10d30a 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -2901,6 +2901,9 @@ TEST_F(CollectionTest, Feature_Optimize_General) { stats = collection->Stats().value(); ASSERT_EQ(stats.doc_count, doc_count); ASSERT_EQ(stats.index_completeness["dense_fp32"], 1); + auto storage_mode = collection->DebugGetHnswStorageMode("dense_fp32"); + ASSERT_TRUE(storage_mode.has_value()); + ASSERT_EQ(enable_mmap ? "mmap" : "buffer_pool", storage_mode.value()); check_doc(); std::cout << "check success 2" << std::endl; @@ -2910,6 +2913,10 @@ TEST_F(CollectionTest, Feature_Optimize_General) { ASSERT_TRUE(result.has_value()); collection = std::move(result.value()); + storage_mode = collection->DebugGetHnswStorageMode("dense_fp32"); + ASSERT_TRUE(storage_mode.has_value()); + ASSERT_EQ(enable_mmap ? "mmap" : "buffer_pool", storage_mode.value()); + check_doc(); std::cout << "check success 3" << std::endl; }; @@ -2920,6 +2927,80 @@ TEST_F(CollectionTest, Feature_Optimize_General) { } } +TEST_F(CollectionTest, Feature_BufferStorage_OnlineWriteAndOptimize) { + constexpr uint64_t kDocCount = 64; + constexpr uint32_t kDimension = 16; + + auto schema = std::make_shared("buffer_write"); + schema->set_max_doc_count_per_segment(MAX_DOC_COUNT_PER_SEGMENT); + schema->add_field(std::make_shared( + "dense_fp32", DataType::VECTOR_FP32, kDimension, false, + std::make_shared(MetricType::L2, 16, 100))); + + auto options = CollectionOptions{false, false, 64 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc( + col_path, *schema, options, 0, kDocCount, false); + ASSERT_NE(nullptr, collection); + + auto check_query = [&](uint64_t doc_id) { + auto query_doc = TestHelper::CreateDoc(doc_id, *schema); + auto vector = query_doc.get>("dense_fp32"); + ASSERT_TRUE(vector.has_value()); + + SearchQuery query; + query.topk_ = 1; + query.target_.field_name_ = "dense_fp32"; + query.target_.set_vector( + std::string(reinterpret_cast(vector.value().data()), + vector.value().size() * sizeof(float))); + auto result = collection->Query(query); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(1U, result.value().size()); + ASSERT_EQ(TestHelper::MakePK(doc_id), result.value()[0]->pk()); + }; + + // New writes are immediately searchable through the BufferStorage-backed + // active FLAT index, before HNSW has been built. + check_query(7); + auto pool_stats = ailego::MemoryLimitPool::get_instance().stats(); + ASSERT_GT(pool_stats.page_used, 0U); + ASSERT_LE(pool_stats.used, pool_stats.pool_size); + ASSERT_TRUE(collection->Flush().ok()); + check_query(7); + + ASSERT_TRUE(collection->Optimize().ok()); + auto storage_mode = collection->DebugGetHnswStorageMode("dense_fp32"); + ASSERT_TRUE(storage_mode.has_value()); + ASSERT_EQ("buffer_pool", storage_mode.value()); + pool_stats = ailego::MemoryLimitPool::get_instance().stats(); + ASSERT_GT(pool_stats.page_used, 0U); + ASSERT_LE(pool_stats.used, pool_stats.pool_size); + check_query(7); + + // The second optimize compacts an immutable HNSW segment with a new delta. + // Its target must also be built through writable BufferStorage. + ASSERT_TRUE( + TestHelper::CollectionInsertDoc(collection, kDocCount, kDocCount + 16) + .ok()); + check_query(kDocCount + 3); + ASSERT_TRUE(collection->Optimize().ok()); + storage_mode = collection->DebugGetHnswStorageMode("dense_fp32"); + ASSERT_TRUE(storage_mode.has_value()); + ASSERT_EQ("buffer_pool", storage_mode.value()); + check_query(7); + check_query(kDocCount + 3); + + collection.reset(); + auto reopened = Collection::Open(col_path, options); + ASSERT_TRUE(reopened.has_value()); + collection = std::move(reopened.value()); + storage_mode = collection->DebugGetHnswStorageMode("dense_fp32"); + ASSERT_TRUE(storage_mode.has_value()); + ASSERT_EQ("buffer_pool", storage_mode.value()); + check_query(7); + check_query(kDocCount + 3); +} + TEST_F(CollectionTest, Feature_Optimize_Concurrent_ReadWrite_NonBlocking) { // Regression: Optimize() no longer blocks writes/reads for its whole // duration. Insert, Fetch and Query must all make progress during the diff --git a/tests/db/common/config_test.cc b/tests/db/common/config_test.cc index be5f6ed0b..742f6d7ff 100644 --- a/tests/db/common/config_test.cc +++ b/tests/db/common/config_test.cc @@ -14,6 +14,7 @@ #include "zvec/db/config.h" #include +#include "db/common/global_resource.h" #include "zvec/db/status.h" using namespace zvec; @@ -129,6 +130,41 @@ TEST_F(ConfigTest, ValidateConfigWithInvalidMemoryLimit) { std::string::npos); } +TEST_F(ConfigTest, InvalidInitializeCanBeCorrected) { + GlobalConfig config_instance; + GlobalConfig::ConfigData invalid; + invalid.memory_limit_bytes = 0; + auto invalid_status = config_instance.Initialize(invalid); + ASSERT_FALSE(invalid_status.ok()); + ASSERT_EQ(StatusCode::INVALID_ARGUMENT, invalid_status.code()); + + GlobalConfig::ConfigData valid; + auto valid_status = config_instance.Initialize(valid); + ASSERT_TRUE(valid_status.ok()) << valid_status.message(); + EXPECT_EQ(valid.memory_limit_bytes, config_instance.memory_limit_bytes()); +} + +TEST_F(ConfigTest, FailedResourceInitializationDoesNotPublishConfig) { + ASSERT_EQ(0, GlobalResource::Instance().initialize()); + + GlobalConfig config_instance; + const uint32_t original_query_threads = config_instance.query_thread_count(); + GlobalConfig::ConfigData requested; + const auto &published = GlobalConfig::Instance(); + requested.memory_limit_bytes = published.memory_limit_bytes(); + requested.query_thread_count = published.query_thread_count() + 1; + if (requested.query_thread_count == original_query_threads) { + ++requested.query_thread_count; + } + requested.query_thread_binding = published.query_thread_binding(); + requested.optimize_thread_count = published.optimize_thread_count(); + requested.optimize_thread_binding = published.optimize_thread_binding(); + + const auto status = config_instance.Initialize(requested); + ASSERT_FALSE(status.ok()); + EXPECT_EQ(original_query_threads, config_instance.query_thread_count()); +} + TEST_F(ConfigTest, ValidateConfigWithInvalidQueryThreadCount) { GlobalConfig::ConfigData config; config.query_thread_count = 0; // Invalid value @@ -234,6 +270,31 @@ TEST_F(ConfigTest, LogConfigPolymorphism) { ASSERT_EQ(file_config->GetLoggerType(), FILE_LOG_TYPE_NAME); } +TEST_F(ConfigTest, InitializePublishesAnImmutableLogConfigSnapshot) { + GlobalConfig config_instance; + const GlobalConfig::LogConfig &original_log = config_instance.log_config(); + const auto original_level = original_log.level; + + GlobalConfig::ConfigData requested; + const auto &process_config = GlobalConfig::Instance(); + requested.memory_limit_bytes = process_config.memory_limit_bytes(); + requested.query_thread_count = process_config.query_thread_count(); + requested.query_thread_binding = process_config.query_thread_binding(); + requested.optimize_thread_count = process_config.optimize_thread_count(); + requested.optimize_thread_binding = process_config.optimize_thread_binding(); + auto caller_owned_log = std::make_shared( + GlobalConfig::LogLevel::kInfo); + requested.log_config = caller_owned_log; + + const auto status = config_instance.Initialize(requested); + ASSERT_TRUE(status.ok()) << status.message(); + EXPECT_EQ(GlobalConfig::LogLevel::kInfo, config_instance.log_level()); + + caller_owned_log->level = GlobalConfig::LogLevel::kFatal; + EXPECT_EQ(GlobalConfig::LogLevel::kInfo, config_instance.log_level()); + EXPECT_EQ(original_level, original_log.level); +} + // jieba_dict_dir is the only ConfigData field that can be written outside // of Initialize() — language SDKs call set_default_jieba_dict_dir() at // module-load to register the dict path they bundled. The setter is diff --git a/tests/db/index/column/vector_column_indexer_test.cc b/tests/db/index/column/vector_column_indexer_test.cc index b0bfd9c81..9dba30fe1 100644 --- a/tests/db/index/column/vector_column_indexer_test.cc +++ b/tests/db/index/column/vector_column_indexer_test.cc @@ -2144,7 +2144,7 @@ TEST(VectorColumnIndexerTest, Failure) { std::make_shared(MetricType::IP))); ASSERT_TRUE(indexer); ASSERT_TRUE( - indexer->Open(vector_column_params::ReadOptions{true, true, false}) + indexer->Open(vector_column_params::ReadOptions{false, true, false}) .ok()); // Insert some data first auto data = vector_column_params::VectorData{ diff --git a/tests/db/index/segment/segment_test.cc b/tests/db/index/segment/segment_test.cc index 9582d2bf3..6b90bbdff 100644 --- a/tests/db/index/segment/segment_test.cc +++ b/tests/db/index/segment/segment_test.cc @@ -35,16 +35,43 @@ #include "db/index/common/id_map.h" #include "db/index/common/version_manager.h" #include "db/index/storage/wal/wal_file.h" -#include "segment_test_fixture.h" #include "utils/utils.h" #include "zvec/db/options.h" +#include "segment_test_fixture.h" using namespace zvec; +namespace { + +class ExternalChargeGuard { + public: + ExternalChargeGuard(zvec::ailego::MemoryLimitPool &pool, size_t bytes) + : pool_(pool), + bytes_(bytes), + charged_(pool_.try_charge_external(bytes_)) {} + + ~ExternalChargeGuard() { + if (charged_) { + pool_.release_external(bytes_); + } + } + + bool charged() const { + return charged_; + } + + private: + zvec::ailego::MemoryLimitPool &pool_; + size_t bytes_; + bool charged_; +}; + +} // namespace + TEST_P(SegmentTest, EmptySchema) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); EXPECT_EQ(segment->id(), 0); @@ -56,8 +83,8 @@ TEST_P(SegmentTest, General) { options_.max_buffer_size_ = 1 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 25); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 25); ASSERT_TRUE(segment != nullptr); auto combined_reader = segment->scan({LOCAL_ROW_ID, "id", "name", "age"}); @@ -106,8 +133,8 @@ TEST_P(SegmentTest, General) { TEST_P(SegmentTest, InsertMoreData) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); uint64_t MAX_DOC = 1000; @@ -137,14 +164,14 @@ TEST_P(SegmentTest, InsertScalarTypes) { auto invert_params = std::make_shared(false); schema_->add_field(std::make_shared("binary", DataType::BINARY, - false, invert_params)); + false, invert_params)); schema_->add_field(std::make_shared( "array_binary", DataType::ARRAY_BINARY, false, invert_params)); auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); } @@ -167,8 +194,8 @@ TEST_P(SegmentTest, InsertVectorTypes) { { Version v = version_manager_->get_current_version(); auto result = - Segment::Open(col_path_, *tmp_schema, *v.writing_segment_meta(), id_map_, - delete_store_, version_manager_, options_); + Segment::Open(col_path_, *tmp_schema, *v.writing_segment_meta(), + id_map_, delete_store_, version_manager_, options_); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -179,8 +206,8 @@ TEST_P(SegmentTest, InsertVectorTypes) { TEST_P(SegmentTest, FetchByGlobalDocID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 1); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 1); ASSERT_TRUE(segment != nullptr); auto ret_doc = segment->Fetch(0); @@ -192,8 +219,8 @@ TEST_P(SegmentTest, FetchByGlobalDocID) { TEST_P(SegmentTest, FetchSingleRow) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto func = [&](int index) -> void { @@ -219,8 +246,8 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { int doc_count = 1000; { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); } @@ -229,9 +256,9 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { Version v = version_manager_->get_current_version(); SegmentOptions open_options; open_options.read_only_ = false; - auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), - id_map_, delete_store_, version_manager_, - open_options); + auto result = + Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), id_map_, + delete_store_, version_manager_, open_options); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -259,8 +286,8 @@ TEST_P(SegmentTest, FetchSingleRowWithPersistStore) { TEST_P(SegmentTest, FetchSingleRowWithUserID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({USER_ID, "id", "name"}, 2); @@ -276,8 +303,8 @@ TEST_P(SegmentTest, FetchSingleRowWithUserID) { TEST_P(SegmentTest, FetchSingleRowWithGlobalDocID) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({GLOBAL_DOC_ID, "id", "name"}, 4); @@ -293,8 +320,8 @@ TEST_P(SegmentTest, FetchSingleRowWithGlobalDocID) { TEST_P(SegmentTest, FetchSingleRowWithNegativeIndex) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, -1); @@ -303,8 +330,8 @@ TEST_P(SegmentTest, FetchSingleRowWithNegativeIndex) { TEST_P(SegmentTest, FetchSingleRowWithOutOfRangeIndex) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, 15); @@ -313,8 +340,8 @@ TEST_P(SegmentTest, FetchSingleRowWithOutOfRangeIndex) { TEST_P(SegmentTest, FetchSingleRowWithInvalidColumn) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "invalid_column"}, 0); @@ -323,8 +350,8 @@ TEST_P(SegmentTest, FetchSingleRowWithInvalidColumn) { TEST_P(SegmentTest, FetchSingleRowWithEmptyColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({}, 0); @@ -333,8 +360,8 @@ TEST_P(SegmentTest, FetchSingleRowWithEmptyColumns) { TEST_P(SegmentTest, FetchSingleRowFromEmptySegment) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"id", "name"}, 0); @@ -343,8 +370,8 @@ TEST_P(SegmentTest, FetchSingleRowFromEmptySegment) { TEST_P(SegmentTest, FetchSingleRowWithBinaryFields) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); ExecBatchPtr batch = segment->fetch({"binary", "array_binary"}, 1); @@ -368,8 +395,8 @@ TEST_P(SegmentTest, Recover) { int doc_count = 100; { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); } @@ -424,9 +451,9 @@ TEST_P(SegmentTest, Recover) { Version v = version_manager_->get_current_version(); SegmentOptions open_options; open_options.read_only_ = false; - auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), - id_map_, delete_store_, version_manager_, - open_options); + auto result = + Segment::Open(col_path_, *schema_, *v.writing_segment_meta(), id_map_, + delete_store_, version_manager_, open_options); ASSERT_TRUE(result.has_value()); auto segment = result.value(); @@ -452,8 +479,8 @@ TEST_P(SegmentTest, Recover) { TEST_P(SegmentTest, UpdateDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // before update @@ -484,8 +511,8 @@ TEST_P(SegmentTest, UpdateDoc) { TEST_P(SegmentTest, UpdateDocBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update uint64_t count = segment->doc_count(delete_store_->make_filter()); @@ -513,8 +540,8 @@ TEST_P(SegmentTest, UpdateDocBatch) { TEST_P(SegmentTest, DeleteDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // before update @@ -541,8 +568,8 @@ TEST_P(SegmentTest, DeleteDoc) { TEST_P(SegmentTest, DeleteBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update @@ -562,8 +589,8 @@ TEST_P(SegmentTest, DeleteBatch) { TEST_P(SegmentTest, UpsertDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); // before update @@ -603,8 +630,8 @@ TEST_P(SegmentTest, UpsertDoc) { TEST_P(SegmentTest, UpsertDocBatch) { int doc_count = 10; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); // before update @@ -648,8 +675,8 @@ TEST_P(SegmentTest, UpsertDocBatch) { TEST_P(SegmentTest, Flush) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Flush the segment @@ -659,8 +686,8 @@ TEST_P(SegmentTest, Flush) { TEST_P(SegmentTest, FlushAfterInsert) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Flush the segment @@ -686,8 +713,8 @@ TEST_P(SegmentTest, FlushAfterInsert) { TEST_P(SegmentTest, Dump) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); // Dump the segment @@ -701,8 +728,8 @@ TEST_P(SegmentTest, Dump) { TEST_P(SegmentTest, DocCount) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 50); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 50); ASSERT_TRUE(segment != nullptr); // Get document count @@ -760,8 +787,8 @@ TEST_P(SegmentTest, CombinedVectorColumnIndexer) { options_.max_buffer_size_ = 10 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); @@ -947,8 +974,8 @@ TEST_P(SegmentTest, CombinedVectorColumnIndexerQueryWithPks) { TEST_P(SegmentTest, ConcurrentInsertOperations) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); const int num_threads = 4; @@ -980,8 +1007,8 @@ TEST_P(SegmentTest, ConcurrentInsertOperations) { TEST_P(SegmentTest, ConcurrentMixedOperations) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 100); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 100); ASSERT_TRUE(segment != nullptr); std::vector threads; @@ -1022,8 +1049,8 @@ TEST_P(SegmentTest, ConcurrentMixedOperations) { // corner cases TEST_P(SegmentTest, DuplicateInsert) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 0); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 0); ASSERT_TRUE(segment != nullptr); Doc doc1 = test::TestHelper::CreateDoc(0, *schema_); @@ -1051,8 +1078,8 @@ TEST_P(SegmentTest, DuplicateInsert) { TEST_P(SegmentTest, DuplicateDelete) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_2"); @@ -1068,8 +1095,8 @@ TEST_P(SegmentTest, DuplicateDelete) { TEST_P(SegmentTest, DeleteNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_999"); @@ -1078,8 +1105,8 @@ TEST_P(SegmentTest, DeleteNonExistentDoc) { TEST_P(SegmentTest, UpdateNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); Doc doc = test::TestHelper::CreateDoc(999, *schema_); @@ -1091,8 +1118,8 @@ TEST_P(SegmentTest, UpdateNonExistentDoc) { TEST_P(SegmentTest, UpsertNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); Doc doc = test::TestHelper::CreateDoc(999, *schema_); @@ -1109,8 +1136,8 @@ TEST_P(SegmentTest, UpsertNonExistentDoc) { TEST_P(SegmentTest, ScanWithEmptyColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto reader = segment->scan({}); @@ -1119,8 +1146,8 @@ TEST_P(SegmentTest, ScanWithEmptyColumns) { TEST_P(SegmentTest, ScanWithInvalidColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // Try to scan with invalid column name @@ -1130,8 +1157,8 @@ TEST_P(SegmentTest, ScanWithInvalidColumns) { TEST_P(SegmentTest, FetchNonExistentDoc) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto doc = segment->Fetch(999); @@ -1140,8 +1167,8 @@ TEST_P(SegmentTest, FetchNonExistentDoc) { TEST_P(SegmentTest, FetchWithInvalidSegmentDocIDs) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); std::vector invalid_segment_doc_ids = {999, 1000}; @@ -1152,8 +1179,8 @@ TEST_P(SegmentTest, FetchWithInvalidSegmentDocIDs) { TEST_P(SegmentTest, FetchWithInvalidColumns) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 10); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 10); ASSERT_TRUE(segment != nullptr); // Try to fetch with invalid column name @@ -1166,8 +1193,8 @@ TEST_P(SegmentTest, InsertEmptyDocWithNullableSchema) { auto nullable_schema = test::TestHelper::CreateNormalSchema(true, col_name_); auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *nullable_schema, 0, 0, id_map_, delete_store_, version_manager_, - options_, 0, 0); + col_path_, *nullable_schema, 0, 0, id_map_, delete_store_, + version_manager_, options_, 0, 0); ASSERT_TRUE(segment != nullptr); Doc empty_doc; @@ -1178,8 +1205,8 @@ TEST_P(SegmentTest, InsertEmptyDocWithNullableSchema) { TEST_P(SegmentTest, MultipleDuplicateDeletes) { auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, 5); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, 5); ASSERT_TRUE(segment != nullptr); auto status1 = segment->Delete("pk_1"); @@ -1202,8 +1229,8 @@ TEST_P(SegmentTest, FetchWithTwoVectorFields) { int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment.reset(); version_manager_.reset(); @@ -1262,8 +1289,8 @@ TEST_P(SegmentTest, FetchPerf) { int doc_count = 1000; options_.max_buffer_size_ = 100 * 1024; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1367,8 +1394,8 @@ TEST_P(SegmentTest, AddColumn) { options_.max_buffer_size_ = 10 * 1024 * 1024; int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->add_column( @@ -1600,8 +1627,8 @@ TEST_P(SegmentTest, AddNullableColumnWithoutExpressionMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1674,9 +1701,8 @@ TEST_P(SegmentTest, AddNullableColumnWithoutExpressionMultiBlock) { auto field_schema = std::make_shared(nullable_col_name, data_type, true); s = segment->add_column(field_schema, "", AddColumnOptions()); - ASSERT_TRUE(s.ok()) - << "Failed to add nullable column " << nullable_col_name << ": " - << s.message(); + ASSERT_TRUE(s.ok()) << "Failed to add nullable column " << nullable_col_name + << ": " << s.message(); auto combined_reader = segment->scan({"id", "name", "age", nullable_col_name}); @@ -1699,8 +1725,8 @@ TEST_P(SegmentTest, AddColumnWithExpressionMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1804,11 +1830,17 @@ TEST_P(SegmentTest, AddColumnWithExpressionMultiBlock) { } TEST_P(SegmentTest, AlterColumnMultiBlock) { + constexpr size_t kSharedCacheCapacity = MIN_MEMORY_LIMIT_BYTES; + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + if (!GetParam() && memory_pool.capacity() != kSharedCacheCapacity) { + ASSERT_EQ(0, memory_pool.init(kSharedCacheCapacity)); + } + options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1868,10 +1900,20 @@ TEST_P(SegmentTest, AlterColumnMultiBlock) { segment = std::move(result).value(); ASSERT_TRUE(segment != nullptr); + constexpr size_t kExternalCharge = 4096; + ExternalChargeGuard external_charge(memory_pool, kExternalCharge); + if (!GetParam()) { + ASSERT_TRUE(external_charge.charged()); + } + // Alter column type: int32 -> int64 on multi-block segment auto new_field = std::make_shared("id", DataType::INT64, false); s = segment->alter_column("id", new_field, AlterColumnOptions()); ASSERT_TRUE(s.ok()) << "alter_column failed: " << s.message(); + if (!GetParam()) { + EXPECT_EQ(kSharedCacheCapacity, memory_pool.capacity()); + EXPECT_GE(memory_pool.external_used(), kExternalCharge); + } auto combined_reader = segment->scan({"id", "name", "age"}); ASSERT_TRUE(combined_reader != nullptr); @@ -1888,11 +1930,17 @@ TEST_P(SegmentTest, AlterColumnMultiBlock) { } TEST_P(SegmentTest, DropColumnMultiBlock) { + constexpr size_t kSharedCacheCapacity = MIN_MEMORY_LIMIT_BYTES; + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + if (!GetParam() && memory_pool.capacity() != kSharedCacheCapacity) { + ASSERT_EQ(0, memory_pool.init(kSharedCacheCapacity)); + } + options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -1952,9 +2000,19 @@ TEST_P(SegmentTest, DropColumnMultiBlock) { segment = std::move(result).value(); ASSERT_TRUE(segment != nullptr); + constexpr size_t kExternalCharge = 4096; + ExternalChargeGuard external_charge(memory_pool, kExternalCharge); + if (!GetParam()) { + ASSERT_TRUE(external_charge.charged()); + } + // Drop column on multi-block segment s = segment->drop_column("id"); ASSERT_TRUE(s.ok()) << "drop_column failed: " << s.message(); + if (!GetParam()) { + EXPECT_EQ(kSharedCacheCapacity, memory_pool.capacity()); + EXPECT_GE(memory_pool.external_used(), kExternalCharge); + } auto combined_reader = segment->scan({"id"}); ASSERT_TRUE(combined_reader == nullptr); @@ -1978,8 +2036,8 @@ TEST_P(SegmentTest, AddNullableThenAlterDropMultiBlock) { options_.max_buffer_size_ = 1 * 1024; int doc_count = 100; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); segment->dump(); @@ -2087,8 +2145,8 @@ TEST_P(SegmentTest, AlterColumn) { // create segment int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->alter_column( @@ -2203,8 +2261,8 @@ TEST_P(SegmentTest, DropColumn) { // create segment int doc_count = 1000; auto segment = test::TestHelper::CreateSegmentWithDoc( - col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_, - 0, doc_count); + col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, + options_, 0, doc_count); ASSERT_TRUE(segment != nullptr); auto s = segment->drop_column("int32"); diff --git a/tests/db/index/storage/bufferpool_store_test.cc b/tests/db/index/storage/bufferpool_store_test.cc index 3ea9024c1..a603c0c72 100644 --- a/tests/db/index/storage/bufferpool_store_test.cc +++ b/tests/db/index/storage/bufferpool_store_test.cc @@ -21,7 +21,9 @@ #include #include #include +#include #include "db/index/storage/bufferpool_forward_store.h" +#include "db/index/storage/parquet_buffer_pool.h" #include "utils/utils.h" using namespace zvec; @@ -45,6 +47,35 @@ class BufferPoolStoreTest : public testing::Test { std::string parquet_path = "test.parquet"; }; +TEST_F(BufferPoolStoreTest, EscapedNestedScalarKeepsParquetCachePinned) { + auto store = std::make_shared(parquet_path); + ASSERT_TRUE(store->Open().ok()); + const int column = store->physic_schema()->GetFieldIndex("list_utf8"); + ASSERT_GE(column, 0); + + std::shared_ptr scalar; + { + auto handle = ParquetBufferPool::get_instance().acquire_buffer( + ParquetBufferID(parquet_path, column, /*row_group=*/0)); + auto data = handle.data(); + ASSERT_NE(nullptr, data); + auto scalar_result = data->GetScalar(0); + ASSERT_TRUE(scalar_result.ok()) << scalar_result.status().ToString(); + scalar = scalar_result.ValueOrDie(); + } + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + const size_t pinned_bytes = memory_pool.external_used(); + ASSERT_GT(pinned_bytes, 0u); + EXPECT_EQ(0u, ailego::BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(pinned_bytes, memory_pool.external_used()); + EXPECT_FALSE(scalar->ToString().empty()); + + scalar.reset(); + EXPECT_EQ(1u, ailego::BlockEvictionQueue::get_instance().batch_recycle(1)); + EXPECT_EQ(0u, memory_pool.external_used()); +} + TEST_F(BufferPoolStoreTest, ParquetFetch) { auto store = std::make_shared(parquet_path); diff --git a/tools/core/local_builder.cc b/tools/core/local_builder.cc index ff2d4d100..3e8fc38df 100644 --- a/tools/core/local_builder.cc +++ b/tools/core/local_builder.cc @@ -477,12 +477,9 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, uint32_t thread_count, RetrievalMode retrieval_mode, const IndexStorage::Pointer &storage = nullptr) { int ret; - ailego::ThreadPool pool(thread_count, false); - thread_count = static_cast(pool.count()); std::atomic finished{0}; - int errcode = 0; + std::atomic errcode{0}; std::mutex mutex; - std::atomic_bool error{false}; std::condition_variable cond{}; auto meta = streamer->meta(); @@ -528,6 +525,12 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, }; } + // Declare the pool after every object captured by worker tasks. Destruction + // is reversed, so an early return joins active workers before their captured + // state is destroyed. + ailego::ThreadPool pool(thread_count, false); + thread_count = static_cast(pool.count()); + auto do_build = [&](size_t idx) { AILEGO_DEFER([&]() { std::lock_guard latch(mutex); @@ -535,50 +538,57 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, }); auto ctx = streamer->create_context(); if (!ctx) { - if (!error.exchange(true)) { + int expected = 0; + if (errcode.compare_exchange_strong(expected, IndexError_NoMemory)) { LOG_ERROR("Failed to create streamer context"); - errcode = IndexError_NoMemory; } return; } std::string ovec; IndexQueryMeta ometa; - for (uint32_t id = idx; id < holder->count() && !stop_now; + for (uint32_t id = idx; id < holder->count() && !stop_now && + errcode.load(std::memory_order_acquire) == 0; id += thread_count) { uint64_t key = holder->get_key(id); + int task_ret = 0; if (retrieval_mode == RM_DENSE) { if (reformer) { - ret = reformer->convert(holder->get_vector_by_index(id), qmeta, &ovec, - &ometa); - if (ret != 0) { - LOG_ERROR("Failed to convert vector for %s", IndexError::What(ret)); - errcode = ret; + task_ret = reformer->convert(holder->get_vector_by_index(id), qmeta, + &ovec, &ometa); + if (task_ret != 0) { + int expected = 0; + if (errcode.compare_exchange_strong(expected, task_ret)) { + LOG_ERROR("Failed to convert vector for %s", + IndexError::What(task_ret)); + } return; } - ret = add_to_streamer(key, ovec.data(), ometa, ctx); + task_ret = add_to_streamer(key, ovec.data(), ometa, ctx); } else { - ret = + task_ret = add_to_streamer(key, holder->get_vector_by_index(id), qmeta, ctx); } } else { - LOG_ERROR("Retrieval mode not supported"); - errcode = IndexError_Unsupported; + int expected = 0; + if (errcode.compare_exchange_strong(expected, IndexError_Unsupported)) { + LOG_ERROR("Retrieval mode not supported"); + } return; } - if (ailego_unlikely(ret != 0)) { - if (!error.exchange(true)) { + if (ailego_unlikely(task_ret != 0)) { + int expected = 0; + if (errcode.compare_exchange_strong(expected, task_ret)) { LOG_ERROR("streamer add_impl failed"); - errcode = ret; } return; } if (id >= keep_docs) { - ret = streamer->remove_impl(holder->get_key(id - keep_docs), ctx); - if (ailego_unlikely(ret != 0)) { - if (!error.exchange(true)) { + task_ret = streamer->remove_impl(holder->get_key(id - keep_docs), ctx); + if (ailego_unlikely(task_ret != 0)) { + int expected = 0; + if (errcode.compare_exchange_strong(expected, task_ret)) { LOG_ERROR("streamer remove_impl failed"); - errcode = ret; } return; } @@ -596,16 +606,16 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, std::unique_lock lk(mutex); cond.wait_until( lk, std::chrono::system_clock::now() + std::chrono::seconds(15)); - if (error.load(std::memory_order_acquire)) { + if (errcode.load(std::memory_order_acquire) != 0) { LOG_ERROR("Failed to build index while waiting finish"); - return errcode; + return errcode.load(std::memory_order_relaxed); } LOG_INFO("Built cnt %zu, finished percent %.3f%%", finished.load(), finished.load() * 100.0f / holder->count()); } - if (error.load(std::memory_order_acquire)) { + if (errcode.load(std::memory_order_acquire) != 0) { LOG_ERROR("Failed to build index while waiting finish"); - return errcode; + return errcode.load(std::memory_order_relaxed); } pool.wait_finish(); @@ -669,9 +679,7 @@ int build_by_streamer(IndexStreamer::Pointer &streamer, LOG_DEBUG("thread count: %zu, retrieval mode: %s", thread_count, retrieval_mode == 1 ? "Dense" : "Sparse"); - do_build_by_streamer(streamer, thread_count, retrieval_mode, storage); - - return 0; + return do_build_by_streamer(streamer, thread_count, retrieval_mode, storage); } IndexSparseHolder::Pointer convert_sparse_holder(