diff --git a/src/ailego/CMakeLists.txt b/src/ailego/CMakeLists.txt index 29cf22cd1..1e524a1b2 100644 --- a/src/ailego/CMakeLists.txt +++ b/src/ailego/CMakeLists.txt @@ -16,6 +16,10 @@ set(EXTRA_LIBS ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS}) if(UNIX AND NOT APPLE) list(APPEND EXTRA_LIBS ${LIB_RT}) + find_library(LIB_AIO NAMES aio) + if(LIB_AIO) + list(APPEND EXTRA_LIBS ${LIB_AIO}) + endif() endif() if(NOT ANDROID AND AUTO_DETECT_ARCH) @@ -123,3 +127,8 @@ cc_library( LIBS ${EXTRA_LIBS} VERSION "${GIT_SRCS_VER}" ) + +if(LIB_AIO) + target_compile_definitions(zvec_ailego PUBLIC ZVEC_HAS_LIBAIO) + message(STATUS "Found libaio: ${LIB_AIO} (HNSW async prefetch enabled)") +endif() diff --git a/src/ailego/buffer/block_eviction_queue.cc b/src/ailego/buffer/block_eviction_queue.cc index eff930127..328ac24a5 100644 --- a/src/ailego/buffer/block_eviction_queue.cc +++ b/src/ailego/buffer/block_eviction_queue.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include namespace zvec { @@ -59,13 +60,44 @@ bool BlockEvictionQueue::evict_block(BlockType &item) { void BlockEvictionQueue::recycle() { BlockType item; + // Upper bound on iterations: the CLOCK second-chance path re-enqueues + // spared pages, so is_full() alone could spin for a while when most queued + // pages are hot. Cap the work a single foreground caller absorbs; the + // background evictor picks up the slack. + const size_t max_attempts = + evict_batch_size_ * 200 * CACHE_QUEUE_NUM + 16; + size_t attempts = 0; while (MemoryLimitPool::get_instance().is_full() && evict_block(item)) { - 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); + { + 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); + } } + if (++attempts >= max_attempts) { + break; + } + } +} + +size_t BlockEvictionQueue::batch_recycle(size_t count) { + std::shared_lock lock(valid_owners_mutex_); + size_t evicted = 0; + // Bound attempts so spared (second-chance) pages that get re-enqueued do + // not turn this into an unbounded loop when nothing is truly evictable. + const size_t max_attempts = count * 4 + 16; + size_t attempts = 0; + while (evicted < count && attempts < max_attempts) { + ++attempts; + BlockType item; + if (!evict_single_block(item)) break; + if (item.owner == nullptr || + valid_owners_.find(item.owner) == valid_owners_.end()) continue; + if (item.owner->is_dead_block(item.owner_key, item.version)) continue; + if (item.owner->evict_block(item.owner_key)) ++evicted; } + return evicted; } bool BlockEvictionQueue::add_single_block(const BlockType &block, @@ -78,29 +110,188 @@ bool BlockEvictionQueue::add_single_block(const BlockType &block, return true; } +MemoryLimitPool::~MemoryLimitPool() { + stop_background_evictor(); + drain_free_list(); +} + +void MemoryLimitPool::drain_free_list() { + // Clear every shard free-list, then release the backing slabs. The shard + // heads point into the slabs, so they must be dropped before ailego_free. + for (size_t i = 0; i < kNumFreeShards; ++i) { + std::lock_guard lk(free_shards_[i].mutex); + free_shards_[i].head = nullptr; + free_shards_[i].count.store(0, std::memory_order_relaxed); + } + std::lock_guard lock(slab_mutex_); + free_all_slabs_locked(); +} + +void MemoryLimitPool::free_all_slabs_locked() { + size_t n = slabs_.size(); + for (char *base : slabs_) { + ailego_free(base); + } + slabs_.clear(); + slab_cursor_ = nullptr; + slab_remaining_ = 0; + if (n > 0) { + LOG_INFO("MemoryLimitPool: released %zu slabs", n); + } +} + +size_t MemoryLimitPool::pick_shard() { + // Sticky per-thread shard assignment: threads keep reusing the same shard, + // maximizing free-list locality and minimizing cross-thread lock traffic. + thread_local size_t idx = + shard_seq_.fetch_add(1, std::memory_order_relaxed); + return idx % kNumFreeShards; +} + +char *MemoryLimitPool::carve_from_slab_locked(size_t buffer_size) { + static constexpr size_t kSlabAlign = 4096UL; + static constexpr size_t kSlabBytes = 2UL * 1024UL * 1024UL; + if (buffer_size == 0 || (buffer_size & (kSlabAlign - 1UL)) != 0) { + char *p = + static_cast(ailego_aligned_malloc(buffer_size, kSlabAlign)); + if (p) { + slabs_.push_back(p); + } + return p; + } + if (slab_remaining_ < buffer_size) { + size_t slab_size = (kSlabBytes < buffer_size) ? buffer_size : kSlabBytes; + slab_size = ((slab_size + buffer_size - 1UL) / buffer_size) * buffer_size; + char *base = + static_cast(ailego_aligned_malloc(slab_size, kSlabAlign)); + if (!base) { + return nullptr; + } + slabs_.push_back(base); + slab_cursor_ = base; + slab_remaining_ = slab_size; + } + char *p = slab_cursor_; + slab_cursor_ += buffer_size; + slab_remaining_ -= buffer_size; + return p; +} + int MemoryLimitPool::init(size_t pool_size) { + // 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_ = 0; BlockEvictionQueue::get_instance().recycle(); + drain_free_list(); pool_size_ = pool_size; LOG_INFO("MemoryLimitPool initialized with pool size: %lu", pool_size_); + if (pool_size_ > 0) { + start_background_evictor(); + } return 0; } +void MemoryLimitPool::start_background_evictor() { + bool expected = false; + if (!bg_running_.compare_exchange_strong(expected, true)) { + return; // already running + } + bg_thread_ = std::thread([this] { background_evict_loop(); }); +} + +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() || + used_size_.load() >= high_watermark(); + }); + } + if (!bg_running_.load()) break; + if (pool_size_ == 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) break; // queues empty or nothing evictable right now + 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_) { + // 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; } desired = expected + buffer_size; } while (!used_size_.compare_exchange_weak(expected, desired)); - buffer = (char *)ailego_aligned_malloc(buffer_size, 4096); + + // Fast path: pull from this thread's shard free-list. + size_t s = pick_shard(); + { + std::lock_guard lock(free_shards_[s].mutex); + if (free_shards_[s].head) { + buffer = free_shards_[s].head; + free_shards_[s].head = *reinterpret_cast(buffer); + free_shards_[s].count.fetch_sub(1, std::memory_order_relaxed); + alloc_from_freelist_.fetch_add(1, std::memory_order_relaxed); + return true; + } + } + // Steal path: our shard is empty, but buffers freed by other threads (most + // notably the background evictor, which releases into its own sticky shard) + // may be sitting idle in foreign shards. Reuse one of those before carving + // new slab memory, otherwise slab allocations grow without bound while freed + // buffers pile up unused, since used_size_ accounting stays within budget. + for (size_t i = 1; i < kNumFreeShards; ++i) { + size_t victim = (s + i) % kNumFreeShards; + std::lock_guard lock(free_shards_[victim].mutex); + if (free_shards_[victim].head) { + buffer = free_shards_[victim].head; + free_shards_[victim].head = *reinterpret_cast(buffer); + free_shards_[victim].count.fetch_sub(1, std::memory_order_relaxed); + alloc_from_freelist_.fetch_add(1, std::memory_order_relaxed); + return true; + } + } + // Cold path: all shards empty, carve a fresh buffer from the slab allocator. + { + std::lock_guard lock(slab_mutex_); + buffer = carve_from_slab_locked(buffer_size); + } if (!buffer) { used_size_.fetch_sub(buffer_size); return false; } + alloc_from_slab_.fetch_add(1, std::memory_order_relaxed); return true; } @@ -119,7 +310,11 @@ void MemoryLimitPool::release_buffer(char *buffer, const size_t buffer_size) { desired = expected - buffer_size; assert(expected >= buffer_size); } while (!used_size_.compare_exchange_weak(expected, desired)); - ailego_free(buffer); + size_t s = pick_shard(); + std::lock_guard lock(free_shards_[s].mutex); + *reinterpret_cast(buffer) = free_shards_[s].head; + free_shards_[s].head = buffer; + free_shards_[s].count.fetch_add(1, std::memory_order_relaxed); } void MemoryLimitPool::release_external(const size_t buffer_size) { @@ -135,5 +330,100 @@ bool MemoryLimitPool::is_full() { return used_size_.load() >= pool_size_; } +size_t MemoryLimitPool::batch_acquire_buffers(size_t buffer_size, + char **out, size_t count) { + if (count == 0) return 0; + size_t total_size = count * buffer_size; + size_t actual_count = count; + size_t expected, desired; + do { + expected = used_size_.load(); + if (expected >= pool_size_) return 0; + size_t avail = (pool_size_ - 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)); + + size_t from_list = 0; + size_t s = pick_shard(); + { + std::lock_guard lock(free_shards_[s].mutex); + while (from_list < actual_count && free_shards_[s].head) { + out[from_list] = free_shards_[s].head; + free_shards_[s].head = *reinterpret_cast(out[from_list]); + free_shards_[s].count.fetch_sub(1, std::memory_order_relaxed); + ++from_list; + } + } + if (from_list) { + alloc_from_freelist_.fetch_add(from_list, std::memory_order_relaxed); + } + // Steal from foreign shards before carving new slab memory (see the note in + // try_acquire_buffer): freed buffers stranded in other shards must be reused, + // otherwise slab allocations grow without bound under eviction churn. + if (from_list < actual_count) { + for (size_t i = 1; i < kNumFreeShards && from_list < actual_count; ++i) { + size_t victim = (s + i) % kNumFreeShards; + std::lock_guard lock(free_shards_[victim].mutex); + while (from_list < actual_count && free_shards_[victim].head) { + out[from_list] = free_shards_[victim].head; + free_shards_[victim].head = *reinterpret_cast(out[from_list]); + free_shards_[victim].count.fetch_sub(1, std::memory_order_relaxed); + alloc_from_freelist_.fetch_add(1, std::memory_order_relaxed); + ++from_list; + } + } + } + if (from_list < actual_count) { + std::lock_guard lock(slab_mutex_); + for (size_t i = from_list; i < actual_count; ++i) { + out[i] = carve_from_slab_locked(buffer_size); + if (!out[i]) { + size_t rollback = (actual_count - i) * buffer_size; + used_size_.fetch_sub(rollback); + actual_count = i; + break; + } + alloc_from_slab_.fetch_add(1, std::memory_order_relaxed); + } + } + return actual_count; +} + +MemoryLimitPool::PoolStats MemoryLimitPool::stats() const { + PoolStats s; + s.pool_size = pool_size_; + s.used = used_size_.load(std::memory_order_relaxed); + 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.alloc_from_freelist = alloc_from_freelist_.load(std::memory_order_relaxed); + s.alloc_from_slab = alloc_from_slab_.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.high_watermark_hits = high_watermark_hits_.load(std::memory_order_relaxed); + return s; +} + +void MemoryLimitPool::log_stats() const { + PoolStats s = stats(); + LOG_INFO( + "MemoryLimitPool stats: pool_size=%llu used=%llu free_buffers=%llu " + "alloc_from_freelist=%llu alloc_from_slab=%llu bg_evict_rounds=%llu " + "bg_evicted_buffers=%llu high_watermark_hits=%llu", + static_cast(s.pool_size), + static_cast(s.used), + static_cast(s.free_buffers), + static_cast(s.alloc_from_freelist), + static_cast(s.alloc_from_slab), + static_cast(s.bg_evict_rounds), + static_cast(s.bg_evicted_buffers), + static_cast(s.high_watermark_hits)); +} + } // namespace ailego } // namespace zvec diff --git a/src/ailego/buffer/vector_page_table.cc b/src/ailego/buffer/vector_page_table.cc index 3318db1ca..570e9efa4 100644 --- a/src/ailego/buffer/vector_page_table.cc +++ b/src/ailego/buffer/vector_page_table.cc @@ -13,12 +13,14 @@ // limitations under the License. #include +#include #include #include #include #include #include #include +#include #if defined(_MSC_VER) #ifndef NOMINMAX @@ -91,6 +93,7 @@ bool VectorPageTable::init(size_t entry_num) { 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].referenced.store(false); segments_[s][i].buffer = nullptr; segments_[s][i].file_offset = 0; } @@ -127,6 +130,7 @@ bool VectorPageTable::extend(size_t new_entry_num) { 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].referenced.store(false); segments_[s][i].buffer = nullptr; segments_[s][i].file_offset = 0; } @@ -143,17 +147,26 @@ bool VectorPageTable::extend(size_t new_entry_num) { char *VectorPageTable::acquire_block(block_id_t block_id) { assert(block_id < entry_num_.load(std::memory_order_relaxed)); Entry &e = entry_at(block_id); - while (true) { - int current_count = e.ref_count.load(std::memory_order_acquire); - if (current_count < 0) { + int old = e.ref_count.fetch_add(1, std::memory_order_acq_rel); + if (ailego_likely(old >= 0)) { + // CLOCK: record the access so the evictor grants this page a second + // chance, and count the hit for observability. + e.referenced.store(true, std::memory_order_relaxed); + hit_count_.fetch_add(1, std::memory_order_relaxed); + return e.buffer; + } + // Undo the increment: the entry is in eviction state. + // Bounded retry to prevent livelock under extreme contention. + int cur = old + 1; + for (int attempts = 0; attempts < 64; ++attempts) { + if (cur >= 0 || cur == std::numeric_limits::min()) break; + if (e.ref_count.compare_exchange_weak(cur, cur - 1, + std::memory_order_relaxed, + std::memory_order_relaxed)) { return nullptr; } - if (e.ref_count.compare_exchange_weak(current_count, current_count + 1, - std::memory_order_acq_rel, - std::memory_order_acquire)) { - return e.buffer; - } } + return nullptr; } void VectorPageTable::release_block(block_id_t block_id) { @@ -161,6 +174,9 @@ void VectorPageTable::release_block(block_id_t block_id) { Entry &e = entry_at(block_id); if (e.ref_count.fetch_sub(1, std::memory_order_release) == 1) { + if (e.in_evict_queue.load(std::memory_order_relaxed)) { + return; + } std::atomic_thread_fence(std::memory_order_acquire); bool expected = false; if (e.in_evict_queue.compare_exchange_strong(expected, true, @@ -170,54 +186,68 @@ void VectorPageTable::release_block(block_id_t block_id) { block.owner = this; block.owner_key = block_id; block.version = 0; - BlockEvictionQueue::get_instance().add_single_block(block, 0); + BlockEvictionQueue::get_instance().add_single_block( + block, static_cast(e.evict_priority)); } } } -void VectorPageTable::evict_block(block_id_t block_id) { +bool VectorPageTable::evict_block(block_id_t block_id) { 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; + static constexpr int kEvicting = std::numeric_limits::min() / 2; + bool evicted = false; if (e.ref_count.compare_exchange_strong(expected, kEvicting)) { + // CLOCK second chance: if the page was accessed since it entered the + // eviction queue, spare it once -- clear the bit, return it to the + // released state and re-enqueue it for a later eviction pass. This is + // what turns the underlying FIFO queue into an approximate-LRU policy. + if (e.referenced.load(std::memory_order_relaxed)) { + e.referenced.store(false, std::memory_order_relaxed); + second_chance_count_.fetch_add(1, std::memory_order_relaxed); + // Restore the released state. Keep in_evict_queue == true because the + // page is (re)inserted into the queue below; recycle() only consumed + // the previous slot. + e.ref_count.store(0, std::memory_order_release); + BlockEvictionQueue::BlockType block; + block.owner = this; + block.owner_key = block_id; + block.version = 0; + BlockEvictionQueue::get_instance().add_single_block( + block, static_cast(e.evict_priority)); + return false; // spared, not reclaimed + } 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); + dirty_flush_count_.fetch_add(1, std::memory_order_relaxed); } 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. + evict_count_.fetch_add(1, std::memory_order_relaxed); e.ref_count.store(std::numeric_limits::min(), std::memory_order_release); + evicted = true; } e.in_evict_queue.store(false, std::memory_order_relaxed); + return evicted; } 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; + static constexpr auto kHardTimeout = std::chrono::seconds(30); while (true) { int current_count = e.ref_count.load(std::memory_order_acquire); if (current_count >= 0) { @@ -228,19 +258,17 @@ char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, 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.referenced.store(false, std::memory_order_relaxed); + e.ever_loaded = true; 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) { @@ -248,7 +276,6 @@ char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, } 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)) { @@ -258,6 +285,17 @@ char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, static_cast(block_id)); warned = true; } + if (elapsed >= kHardTimeout) { + LOG_ERROR( + "set_block_acquired: hard timeout (%lld s) on block_id=%zu; " + "giving up to prevent indefinite thread hang", + static_cast( + std::chrono::duration_cast(elapsed) + .count()), + static_cast(block_id)); + MemoryLimitPool::get_instance().release_buffer(buffer, kVectorPageSize); + return nullptr; + } if (elapsed >= std::chrono::seconds(1) && (now - last_log) >= std::chrono::seconds(1)) { const auto secs = @@ -272,31 +310,84 @@ char *VectorPageTable::set_block_acquired(block_id_t block_id, char *buffer, } } -VecBufferPool::VecBufferPool(const std::string &filename, bool writable) { +VecBufferPool::VecBufferPool(const std::string &filename, bool writable, + bool enable_direct_io) { file_name_ = filename; writable_ = 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); + (void)enable_direct_io; // O_DIRECT not supported on this path +#else + int base_flags = writable_ ? O_RDWR : O_RDONLY; + // Metadata channel: always buffered IO. Serves the unaligned + // header/footer/segment_meta reads & writes and benefits from page cache. + meta_fd_ = ::open(filename.c_str(), base_flags, 0644); + // Page-data channel: optionally O_DIRECT; fall back to buffered open when + // the filesystem (tmpfs/overlayfs/...) rejects O_DIRECT. + int data_flags = base_flags; +#ifdef O_DIRECT + if (enable_direct_io) { + data_flags |= O_DIRECT; + } +#endif + fd_ = ::open(filename.c_str(), data_flags, 0644); +#ifdef O_DIRECT + if (fd_ < 0 && (data_flags & O_DIRECT)) { + 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_ = (data_flags & O_DIRECT) != 0; + } +#else + (void)enable_direct_io; +#endif +#endif + if (fd_ < 0 || meta_fd_ < 0) { + if (fd_ >= 0) { +#if defined(_MSC_VER) + _close(fd_); #else - int flags = writable_ ? O_RDWR : O_RDONLY; - fd_ = ::open(filename.c_str(), flags, 0644); + ::close(fd_); #endif - if (fd_ < 0) { + } + 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) || defined(__linux__) + if (direct_io_enabled_ && LibAioLoader::Instance().Load()) { + aio_ctx_ = nullptr; + if (LibAioLoader::Instance().io_setup(256, &aio_ctx_) == 0) { + aio_enabled_ = true; + } + } +#endif } int VecBufferPool::init() { @@ -320,20 +411,19 @@ int VecBufferPool::init() { // the backing file without needing to know about fd_ directly. if (writable_) { int fd = fd_; - const std::string &name = file_name_; - page_table_.set_flush_callback([fd, &name](block_id_t /*block_id*/, - char *buf, size_t sz, - size_t off) -> int { - ssize_t w = zvec_pwrite(fd, buf, sz, off); - if (w != static_cast(sz)) { - LOG_ERROR( - "Buffer pool flush failed: file[%s], offset[%zu], " - "expected[%zu], got[%zd]", - name.c_str(), off, sz, w); - return -1; - } - return 0; - }); + 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; + }); } return 0; } @@ -375,24 +465,49 @@ char *VecBufferPool::acquire_buffer(block_id_t page_id, int retry) { } 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); - } - 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; + // Cold path: the page is being loaded from disk (or zero-filled), i.e. a + // cache miss. Count it for the pool hit-rate metric. + miss_count_.fetch_add(1, std::memory_order_relaxed); + // Skip pread for pages created by extend_file (beyond the original file + // size at open time) that have never been loaded before. Their on-disk + // content is guaranteed to be zeros (ftruncate). After eviction the + // ever_loaded flag stays true so reloads correctly pread the flushed data. + if (writable_ && page_offset >= initial_file_size_ && + !page_table_.is_ever_loaded(page_id)) { + std::memset(buffer, 0, kVectorPageSize); + } else { + // O_DIRECT requires the IO length to be a multiple of the device block + // size. For files whose size is page-aligned (e.g. BufferStorage), reading + // a full page never reads past EOF. For files that are NOT page-aligned + // (e.g. IVF via BufferReadStorage), the last page may be a short read; + // we accept it and zero-pad the remainder. + 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); + } + 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); + return nullptr; + } + } } return page_table_.set_block_acquired(page_id, buffer, page_offset); } int VecBufferPool::get_meta(size_t offset, size_t length, char *buffer) { - ssize_t read_bytes = zvec_pread(fd_, buffer, length, offset); + 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], " @@ -446,7 +561,7 @@ 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); + 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,23 +576,72 @@ int VecBufferPool::flush_all() { if (!writable_) { return 0; } + const size_t total = page_table_.entry_num(); + if (total == 0) { + return 0; + } + + static constexpr size_t kBatchPages = 256; + const size_t kBatchSize = kBatchPages * kVectorPageSize; + char *batch_buf = + static_cast(ailego_aligned_malloc(kBatchSize, 4096)); + 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 i = 0; + + while (i < total) { + if (!page_table_.is_block_dirty(i)) { + ++i; + continue; + } + + const size_t run_start = i; + size_t run_count = 0; + const size_t limit = batch_buf ? kBatchPages : 1; + while (i < total && run_count < limit && page_table_.is_block_dirty(i)) { + char *buf = page_table_.get_block_buffer(i); + if (!buf) break; + if (batch_buf) { + std::memcpy(batch_buf + run_count * kVectorPageSize, buf, + kVectorPageSize); + } + ++run_count; + ++i; + } + if (run_count == 0) { + ++i; + continue; + } + total_dirty += run_count; + + bool ok = false; + if (batch_buf && run_count > 0) { + const size_t write_size = run_count * kVectorPageSize; + ssize_t w = + zvec_pwrite(fd_, batch_buf, write_size, run_start * kVectorPageSize); + ok = (w == static_cast(write_size)); + } + if (ok) { + for (size_t j = run_start; j < run_start + run_count; ++j) { + page_table_.clear_dirty(j); + } + } else { + for (size_t j = run_start; j < run_start + run_count; ++j) { + int r = page_table_.flush_block(j); + if (r != 0) { + rc = r; + ++fail_count; + } } } } + + if (batch_buf) { + ailego_free(batch_buf); + } 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. LOG_ERROR( "VecBufferPool::flush_all: %zu/%zu dirty page(s) failed to flush, " "file[%s] last_rc=%d -- on-disk data may be stale.", @@ -495,6 +659,10 @@ bool VecBufferPool::extend_file(size_t new_size) { if (new_size <= file_size_) { return true; } + // The backing file must stay page-aligned so that O_DIRECT full-page reads + // never read past EOF. All current callers pass page-aligned targets. + assert(new_size % kVectorPageSize == 0 && + "extend_file target must be page-aligned for O_DIRECT correctness"); // 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 @@ -563,23 +731,111 @@ bool VecBufferPoolHandle::read_range(size_t file_offset, size_t len, size_t last_page = (file_offset + len - 1) / kVectorPageSize; size_t remaining = len; size_t dst_cursor = 0; + + 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); + if (!page) return false; + 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, page + intra_offset, chunk); + pool_.page_table_.release_block(pid); + 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); + if (!page) 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); + std::memcpy(out + dst_cursor, page + intra_offset, chunk); + pool_.page_table_.release_block(static_cast(pg)); + dst_cursor += chunk; + remaining -= chunk; + continue; + } + + ssize_t got = zvec_pread(pool_.fd_, bulk_buf, run_bytes, run_file_off); + 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); - std::memcpy(out + dst_cursor, page + intra_offset, chunk); - pool_.page_table_.release_block(pg); - dst_cursor += chunk; - remaining -= chunk; + size_t actually_read = static_cast(got); + + 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); + pool_.page_table_.set_block_acquired( + pid, page_buf, run_file_off + j * kVectorPageSize); + pool_.page_table_.release_block(pid); + } + } + } + ailego_free(bulk_buf); + pg = run_end - 1; } return true; } @@ -617,5 +873,546 @@ void VecBufferPoolHandle::acquire_one(block_id_t block_id) { pool_.page_table_.acquire_block(block_id); } +void VecBufferPool::warmup() { + const size_t total_pages = page_table_.entry_num(); + // Read in large sequential chunks to minimize syscall overhead. + // Each chunk = 1024 pages = 4MB (maximize sequential I/O throughput). + 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; + + // 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(read_bytes)) break; + + // 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); + page_table_.set_block_acquired(page_id, buf, + file_offset + j * kVectorPageSize); + 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) { + size_t end_page = first_page + page_count; + if (end_page > page_table_.entry_num()) { + end_page = page_table_.entry_num(); + } + if (first_page >= end_page) return; + +#if defined(__linux) || defined(__linux__) + if (aio_enabled_) { + prefetch_pages_aio(first_page, end_page - first_page); + return; + } +#endif + + bool all_loaded = true; + for (size_t pg = first_page; pg < end_page; ++pg) { + if (!page_table_.is_loaded(pg)) { + all_loaded = false; + break; + } + } + if (all_loaded) return; + + 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)) { + ++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; + ssize_t got = zvec_pread(fd_, chunk_buf, read_bytes, file_off); + if (got != static_cast(read_bytes)) { + pg = run_end; + continue; + } + + for (size_t j = 0; j < run_pages; ++j) { + block_id_t pid = static_cast(run_start + j); + if (page_table_.is_loaded(pid)) 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_.set_block_acquired(pid, buf, file_off + j * kVectorPageSize); + page_table_.release_block(pid); + } + pg = run_end; + } + ailego_free(chunk_buf); +} + +void VecBufferPoolHandle::prefetch_range(size_t file_offset, size_t len) { + if (len == 0) return; + 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); +} + +#if defined(__linux) || defined(__linux__) +namespace { +// Dedicated thread-local AIO context for the *blocking* prefetch path +// (prefetch_pages_aio). Kept separate from the shared member aio_ctx_ so that +// each thread reaps only the events it submitted. Sharing one context across +// threads let thread B's io_getevents steal thread A's completions, which made +// prefetch release buffers whose kernel DMA was still in flight -- the DMA then +// overwrote the buffer's first 8 bytes (the MemoryLimitPool free-list next +// pointer), corrupting the list and crashing the next try_acquire_buffer. +struct ThreadLocalPrefetchAioCtx { + io_context_t ctx{nullptr}; + bool inited{false}; + bool ok{false}; + + bool ensure() { + if (inited) return ok; + inited = true; + if (!LibAioLoader::Instance().IsAvailable()) return false; + ctx = nullptr; + if (LibAioLoader::Instance().io_setup(256, &ctx) == 0) { + ok = true; + } + return ok; + } + ~ThreadLocalPrefetchAioCtx() { + if (ok && ctx) { + LibAioLoader::Instance().io_destroy(ctx); + } + } +}; +static thread_local ThreadLocalPrefetchAioCtx tl_prefetch_aio; +} // namespace +#endif + +void VecBufferPool::prefetch_pages_aio( + [[maybe_unused]] block_id_t first_page, + [[maybe_unused]] size_t page_count) { +#if defined(__linux) || defined(__linux__) + static constexpr size_t kMaxBatch = 128; + + // Use a thread-local AIO context: each thread waits only for its own + // completions, so a buffer is never returned to the free-list while the + // kernel is still DMA-ing into it. + if (!tl_prefetch_aio.ensure()) return; + io_context_t ctx = tl_prefetch_aio.ctx; + + size_t end_page = first_page + page_count; + if (end_page > page_table_.entry_num()) { + end_page = page_table_.entry_num(); + } + + size_t pg = first_page; + while (pg < end_page) { + std::vector miss_pages; + miss_pages.reserve(kMaxBatch); + while (pg < end_page && miss_pages.size() < kMaxBatch) { + if (!page_table_.is_loaded(pg)) { + miss_pages.push_back(static_cast(pg)); + } + ++pg; + } + if (miss_pages.empty()) continue; + + size_t count = miss_pages.size(); + std::vector cbs(count); + std::vector cb_ptrs(count); + std::vector buffers(count, nullptr); + + size_t submitted = 0; + for (size_t i = 0; i < count; ++i) { + 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) break; + buffers[submitted] = buf; + size_t offset = miss_pages[submitted] * kVectorPageSize; + io_prep_pread(&cbs[submitted], fd_, buf, kVectorPageSize, + static_cast(offset)); + // Record the submission index so out-of-order completions can be mapped + // back to the correct buffer/page. + cbs[submitted].data = reinterpret_cast(submitted); + cb_ptrs[submitted] = &cbs[submitted]; + ++submitted; + } + + if (submitted == 0) return; + + int ret = LibAioLoader::Instance().io_submit( + ctx, static_cast(submitted), cb_ptrs.data()); + if (ret <= 0) { + for (size_t i = 0; i < submitted; ++i) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + } + return; + } + + // io_submit may accept fewer than requested; the accepted requests are a + // prefix of cb_ptrs. The tail was never submitted (no in-flight DMA), so + // its buffers are safe to release immediately. + size_t accepted = static_cast(ret); + for (size_t i = accepted; i < submitted; ++i) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + buffers[i] = nullptr; + } + + // Block until every accepted I/O completes. The blocking wait (nullptr + // timeout) guarantees no DMA is still in flight when we touch the buffers + // below, so a buffer can never be written by the kernel after being + // returned to the free-list. + std::vector events(accepted); + size_t done = 0; + while (done < accepted) { + int n = LibAioLoader::Instance().io_getevents( + ctx, static_cast(accepted - done), + static_cast(accepted - done), events.data() + done, nullptr); + if (n <= 0) break; + done += static_cast(n); + } + + for (size_t i = 0; i < done; ++i) { + size_t idx = reinterpret_cast(events[i].data); + if (idx >= submitted || buffers[idx] == nullptr) continue; + block_id_t pid = miss_pages[idx]; + if (static_cast(events[i].res) == + static_cast(kVectorPageSize)) { + std::lock_guard lock( + block_mutexes_[pid % VecBufferPool::kMutexBucketCount]); + if (page_table_.is_loaded(pid)) { + MemoryLimitPool::get_instance().release_buffer(buffers[idx], + kVectorPageSize); + } else { + page_table_.set_block_acquired(pid, buffers[idx], + pid * kVectorPageSize); + page_table_.release_block(pid); + } + } else { + MemoryLimitPool::get_instance().release_buffer(buffers[idx], + kVectorPageSize); + } + buffers[idx] = nullptr; + } + + // If io_getevents failed before all events were harvested (done < + // accepted), the remaining buffers may still have in-flight DMA. We must + // NOT release them here (that would reintroduce the use-after-free); leave + // them owned by the pool accounting. Not expected under a blocking wait. + } +#endif +} + +#if defined(__linux) || defined(__linux__) +namespace { +struct ThreadLocalAioCtx { + io_context_t ctx{nullptr}; + bool inited{false}; + bool ok{false}; + + // Pending async AIO state + char *pending_bufs[128]; + block_id_t pending_pids[128]; + size_t pending_count{0}; // total submitted (slot high-water mark) + size_t harvested_count{0}; // how many completed & processed + VecBufferPool *pending_pool{nullptr}; + + ~ThreadLocalAioCtx() { + // Drain all pending AIO before destroying context (thread exit) + size_t in_flight = pending_count - harvested_count; + if (in_flight > 0 && ok) { + struct io_event events[128]; + // Must block-wait: kernel is still DMA-ing into our buffers + LibAioLoader::Instance().io_getevents( + ctx, static_cast(in_flight), + static_cast(in_flight), events, nullptr); + } + // Release ALL pending buffers (both harvested-but-not-released and in-flight) + for (size_t i = 0; i < pending_count; ++i) { + if (pending_bufs[i]) { + MemoryLimitPool::get_instance().release_buffer(pending_bufs[i], + kVectorPageSize); + } + } + pending_count = 0; + harvested_count = 0; + if (ok && ctx) { + LibAioLoader::Instance().io_destroy(ctx); + } + } + + bool ensure() { + if (inited) return ok; + inited = true; + if (!LibAioLoader::Instance().IsAvailable()) return false; + ctx = nullptr; + if (LibAioLoader::Instance().io_setup(128, &ctx) == 0) { + ok = true; + } + return ok; + } +}; + +static thread_local ThreadLocalAioCtx tl_aio; +} // namespace +#endif + +void VecBufferPool::submit_aio_async(const block_id_t *page_ids, size_t count) { + if (count == 0) return; + +#if defined(__linux) || defined(__linux__) + if (!direct_io_enabled_) return; + if (!tl_aio.ensure()) return; + + // Harvest any previously completed AIO (non-blocking) + if (tl_aio.pending_count > 0) { + harvest_aio(); + } + + // Determine how many slots are available for new submissions + size_t base = tl_aio.pending_count; // existing in-flight count + size_t max_new = 128 - base; + if (max_new == 0) return; // All slots occupied by in-flight I/O + + // Filter: skip already-loaded, in-flight, and deduplicate + block_id_t miss_pages[128]; + size_t miss_count = 0; + for (size_t i = 0; i < count && miss_count < max_new; ++i) { + if (page_ids[i] >= page_table_.entry_num()) continue; + if (page_table_.is_loaded(page_ids[i])) continue; + // Skip if already in pending (still in-flight, buf != nullptr) + bool in_flight = false; + for (size_t j = 0; j < base; ++j) { + if (tl_aio.pending_bufs[j] && tl_aio.pending_pids[j] == page_ids[i]) { + in_flight = true; break; + } + } + if (in_flight) continue; + bool dup = false; + for (size_t j = 0; j < miss_count; ++j) { + if (miss_pages[j] == page_ids[i]) { dup = true; break; } + } + if (!dup) miss_pages[miss_count++] = page_ids[i]; + } + if (miss_count == 0) return; + + BlockEvictionQueue::get_instance().batch_recycle(miss_count); + char *buffers[128]; + size_t submitted = MemoryLimitPool::get_instance().batch_acquire_buffers( + kVectorPageSize, buffers, miss_count); + if (submitted == 0) return; + + struct iocb cbs[128]; + struct iocb *cb_ptrs[128]; + for (size_t i = 0; i < submitted; ++i) { + size_t offset = miss_pages[i] * kVectorPageSize; + io_prep_pread(&cbs[i], fd_, buffers[i], kVectorPageSize, + static_cast(offset)); + // Index = base + i so harvest can map back to pending_bufs/pids + cbs[i].data = reinterpret_cast(base + i); + cb_ptrs[i] = &cbs[i]; + } + + int ret = LibAioLoader::Instance().io_submit( + tl_aio.ctx, static_cast(submitted), cb_ptrs); + if (ret <= 0) { + for (size_t i = 0; i < submitted; ++i) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + } + return; + } + + // Append to pending state + size_t actual = static_cast(ret); + for (size_t i = 0; i < actual; ++i) { + tl_aio.pending_bufs[base + i] = buffers[i]; + tl_aio.pending_pids[base + i] = miss_pages[i]; + } + tl_aio.pending_count = base + actual; + tl_aio.pending_pool = this; + + // Release buffers for requests that weren't submitted + for (size_t i = actual; i < submitted; ++i) { + MemoryLimitPool::get_instance().release_buffer(buffers[i], + kVectorPageSize); + } +#else + (void)page_ids; + (void)count; +#endif +} + +void VecBufferPool::harvest_aio() { +#if defined(__linux) || defined(__linux__) + if (!tl_aio.ok || tl_aio.pending_count == 0) return; + size_t in_flight = tl_aio.pending_count - tl_aio.harvested_count; + if (in_flight == 0) return; + + struct io_event events[128]; + struct timespec timeout = {0, 0}; // Non-blocking + int ret = LibAioLoader::Instance().io_getevents( + tl_aio.ctx, 0, static_cast(in_flight), events, &timeout); + + size_t completed = (ret > 0) ? static_cast(ret) : 0; + if (completed == 0) return; // Nothing ready yet + + VecBufferPool *pool = tl_aio.pending_pool; + for (size_t i = 0; i < completed; ++i) { + size_t idx = reinterpret_cast(events[i].data); + if (static_cast(events[i].res) != + static_cast(kVectorPageSize)) { + MemoryLimitPool::get_instance().release_buffer(tl_aio.pending_bufs[idx], + kVectorPageSize); + } else { + block_id_t pid = tl_aio.pending_pids[idx]; + std::lock_guard lock( + pool->block_mutexes_[pid % VecBufferPool::kMutexBucketCount]); + if (pool->page_table_.is_loaded(pid)) { + MemoryLimitPool::get_instance().release_buffer(tl_aio.pending_bufs[idx], + kVectorPageSize); + } else { + pool->page_table_.set_block_acquired(pid, tl_aio.pending_bufs[idx], + pid * kVectorPageSize); + pool->page_table_.release_block(pid); + } + } + tl_aio.pending_bufs[idx] = nullptr; // Mark as processed + } + + tl_aio.harvested_count += completed; + + // If all submitted are harvested, reset for reuse + if (tl_aio.harvested_count == tl_aio.pending_count) { + tl_aio.pending_count = 0; + tl_aio.harvested_count = 0; + tl_aio.pending_pool = nullptr; + } +#endif +} + +void VecBufferPool::wait_aio() { +#if defined(__linux) || defined(__linux__) + if (!tl_aio.ok || tl_aio.pending_count == 0) return; + + VecBufferPool *pool = tl_aio.pending_pool; + struct io_event events[128]; + // Block (NULL timeout) until every in-flight request has completed. Unlike + // harvest_aio() this never returns early, so after it the whole submitted + // batch is guaranteed resident -- the caller can resolve every page as a hit + // without a per-neighbour synchronous fallback. + while (tl_aio.harvested_count < tl_aio.pending_count) { + size_t in_flight = tl_aio.pending_count - tl_aio.harvested_count; + int ret = LibAioLoader::Instance().io_getevents( + tl_aio.ctx, static_cast(in_flight), + static_cast(in_flight), events, nullptr); + if (ret <= 0) break; + size_t completed = static_cast(ret); + for (size_t i = 0; i < completed; ++i) { + size_t idx = reinterpret_cast(events[i].data); + if (static_cast(events[i].res) != + static_cast(kVectorPageSize)) { + MemoryLimitPool::get_instance().release_buffer(tl_aio.pending_bufs[idx], + kVectorPageSize); + } else { + block_id_t pid = tl_aio.pending_pids[idx]; + std::lock_guard lock( + pool->block_mutexes_[pid % VecBufferPool::kMutexBucketCount]); + if (pool->page_table_.is_loaded(pid)) { + MemoryLimitPool::get_instance().release_buffer( + tl_aio.pending_bufs[idx], kVectorPageSize); + } else { + pool->page_table_.set_block_acquired(pid, tl_aio.pending_bufs[idx], + pid * kVectorPageSize); + pool->page_table_.release_block(pid); + } + } + tl_aio.pending_bufs[idx] = nullptr; + } + tl_aio.harvested_count += completed; + } + + tl_aio.pending_count = 0; + tl_aio.harvested_count = 0; + tl_aio.pending_pool = nullptr; +#endif +} + +void VecBufferPool::log_stats() const { + Stats s = stats(); + LOG_INFO( + "VecBufferPool stats: file[%s] hit=%llu miss=%llu hit_rate=%.4f " + "evict=%llu second_chance=%llu dirty_flush=%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)); +} + } // namespace ailego } // namespace zvec diff --git a/src/core/algorithm/flat/flat_searcher.cc b/src/core/algorithm/flat/flat_searcher.cc index ef4e7d4b0..6b60abda2 100644 --- a/src/core/algorithm/flat/flat_searcher.cc +++ b/src/core/algorithm/flat/flat_searcher.cc @@ -147,6 +147,9 @@ int FlatSearcher::load(IndexStorage::Pointer cntr, return IndexError_ReadData; } + // Keep the segment alive so that keys_ pointer remains valid + keys_segment_ = std::move(keys_segment); + for (size_t i = 0; i < keys_count; i++) { key_id_mapping_[keys_[i]] = i; } diff --git a/src/core/algorithm/flat/flat_searcher.h b/src/core/algorithm/flat/flat_searcher.h index 2b0d0d93d..e48a9ced6 100644 --- a/src/core/algorithm/flat/flat_searcher.h +++ b/src/core/algorithm/flat/flat_searcher.h @@ -50,6 +50,7 @@ class FlatSearcher : public IndexSearcher { container_ = nullptr; measure_ = nullptr; features_segment_ = nullptr; + keys_segment_ = nullptr; keys_ = nullptr; key_id_mapping_.clear(); return 0; @@ -172,6 +173,7 @@ class FlatSearcher : public IndexSearcher { IndexMetric::Pointer measure_{}; ailego::Params params_{}; IndexStorage::Segment::Pointer features_segment_{}; + IndexStorage::Segment::Pointer keys_segment_{}; mutable std::vector mapping_{}; mutable std::mutex mapping_mutex_{}; FlatDistanceMatrix distance_matrix_{}; diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc index 1fd6f6567..bed620c86 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.cc +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "hnsw_algorithm.h" +#include +#include +#include #include namespace zvec { @@ -81,8 +84,17 @@ int HnswAlgorithm::search(HnswContext *ctx) const { } dist_t dist = ctx->dist_calculator().dist(entry_point); + const auto &upper_entity = static_cast(ctx->get_entity()); + // Upper-level entry-point descent must never skip a cache-miss page: + // a missed vector makes select_entry_point pick a worse entry point into + // level 0, which drops recall (buffer vs mmap parity). Block-load on miss + // by granting an effectively unbounded I/O budget for the upper levels. + // (Restored: originally added in 3e0c044, dropped by the cb27887 AIO + // refactor, and only the level-0 budget was reinstated in 96722dc.) + upper_entity.reset_io_budget(INT32_MAX); for (level_t cur_level = maxLevel; cur_level >= 1; --cur_level) { select_entry_point(cur_level, &entry_point, &dist, ctx); + upper_entity.release_vectors(); } auto &topk_heap = ctx->topk_heap(); @@ -103,6 +115,11 @@ void HnswAlgorithm::select_entry_point(level_t level, HnswContext *ctx) const { const auto &entity = static_cast(ctx->get_entity()); HnswDistCalculator &dc = ctx->dist_calculator(); + uint32_t buf_cap = entity.max_degree(level); + std::vector neighbor_ids(buf_cap); + std::vector neighbor_vecs(buf_cap); + std::vector dists(buf_cap); + while (true) { const auto neighbors = entity.get_neighbors_typed(level, *entry_point); if (ailego_unlikely(ctx->debugging())) { @@ -113,31 +130,51 @@ void HnswAlgorithm::select_entry_point(level_t level, break; } - std::vector neighbor_vec_blocks; - int ret = entity.get_vector_typed(&neighbors[0], size, neighbor_vec_blocks); - if (ailego_unlikely(ctx->debugging())) { - (*ctx->mutable_stats_get_vector())++; + if (size > buf_cap) { + buf_cap = size; + neighbor_ids.resize(buf_cap); + neighbor_vecs.resize(buf_cap); + dists.resize(buf_cap); } - if (ailego_unlikely(ret != 0)) { - break; + for (uint32_t i = 0; i < size; ++i) { + neighbor_ids[i] = neighbors[i]; } - bool find_closer = false; + if (ailego_unlikely(entity.resolve_vectors(neighbor_ids.data(), size, + neighbor_vecs.data()) != 0)) { + break; + } + if (ailego_unlikely(ctx->debugging())) { + (*ctx->mutable_stats_get_vector())++; + } - std::vector dists(size); - std::vector neighbor_vecs(size); + // Partition: resolved (non-null) to front + uint32_t resolved = 0; for (uint32_t i = 0; i < size; ++i) { - neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + if (neighbor_vecs[i]) { + if (i != resolved) { + std::swap(neighbor_vecs[i], neighbor_vecs[resolved]); + std::swap(neighbor_ids[i], neighbor_ids[resolved]); + } + ++resolved; + } } - dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + if (resolved == 0) { + entity.release_vectors(); + break; // No vectors available, can't progress + } - for (uint32_t i = 0; i < size; ++i) { - dist_t cur_dist = dists[i]; + dc.batch_dist(neighbor_vecs.data(), resolved, dists.data()); + + // Release per-hop pages AFTER batch_dist to prevent UAF. + entity.release_vectors(); - if (cur_dist < *dist) { - *entry_point = neighbors[i]; - *dist = cur_dist; + bool find_closer = false; + for (uint32_t i = 0; i < resolved; ++i) { + if (dists[i] < *dist) { + *entry_point = neighbor_ids[i]; + *dist = dists[i]; find_closer = true; } } @@ -176,7 +213,10 @@ 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. +// fast_search_neighbors: level-0 unfiltered search for all storage +// modes (mmap/BufferPool/contiguous). Vector +// resolution delegated to entity via +// resolve_vectors()/release_vectors(). // Uses BlockHeap (AVX2) or LinearPool (scalar) // for visited tracking and top-k maintenance. // dual_heap_search_neighbors: CandidateHeap + TopkHeap + VisitFilter. @@ -184,20 +224,27 @@ void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, // search, upper levels, and BufferPool fallback. // ============================================================================ -// mmap/contiguous variant: resolve vectors via get_vector_ptr and use -// LinearPool or BlockHeap for visited tracking + top-k maintenance. -// HeapType must expose reset/set_visited/check_visited/push_block/has_next/pop. template void fast_search_neighbors(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) { - const uint32_t max_deg = entity.max_degree(0); // level 0 only + 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(); + // I/O budget: env ZVEC_IO_BUDGET controls (-1=blocking, 0=non-blocking, >0=limit) + { + int32_t budget = static_cast(ef / 2); // default: ef/2 + const char *env = std::getenv("ZVEC_IO_BUDGET"); + if (env && *env) { + budget = std::atoi(env); + } + entity.reset_io_budget(budget); + } + visit.set_visited(entry_point); pool.push_block(&entry_dist, &entry_point, 1); @@ -206,6 +253,18 @@ void fast_search_neighbors(const EntityType &entity, HeapType &pool, std::vector dists(buf_capacity); std::vector neighbor_vecs(buf_capacity); + // Warm start: prefetch pages for entry_point's neighbors + { + const auto neighbors = entity.get_neighbors_typed(0, entry_point); + node_id_t prefetch_ids_buf[64]; + uint32_t pc = 0; + for (uint32_t i = 0; i < neighbors.size() && pc < 64; ++i) { + if (!visit.visited(neighbors[i])) + prefetch_ids_buf[pc++] = neighbors[i]; + } + if (pc > 0) entity.submit_prefetch(prefetch_ids_buf, pc); + } + while (pool.has_next()) { auto current_node = pool.pop(); @@ -219,42 +278,80 @@ void fast_search_neighbors(const EntityType &entity, HeapType &pool, neighbor_vecs.resize(buf_capacity); } - const uint32_t po = - std::min(static_cast(neighbors.size()), prefetch_offset); uint32_t unvisited_count = 0; - uint32_t i = 0; - - // Phase 1: scan first `po` neighbors with prefetch. - for (; i < po; ++i) { + for (uint32_t i = 0; i < neighbors.size(); ++i) { node_id_t node = neighbors[i]; if (visit.visited(node)) continue; visit.set_visited(node); - const void *vec_ptr = entity.get_vector_ptr(node); - const char *p = reinterpret_cast(vec_ptr); - for (uint32_t cl = 0; cl < prefetch_lines; ++cl) { - ailego_prefetch(p + cl * 64); + neighbor_ids[unvisited_count++] = node; + } + + if (unvisited_count == 0) continue; + + // Harvest AIO from previous iteration (or warm start) + entity.harvest_prefetch(); + + if (ailego_unlikely(entity.resolve_vectors(neighbor_ids.data(), + unvisited_count, + neighbor_vecs.data()) != 0)) + break; + + // Partition: move resolved vectors (non-null) to front. + uint32_t resolved = 0; + for (uint32_t i = 0; i < unvisited_count; ++i) { + if (neighbor_vecs[i]) { + if (i != resolved) { + std::swap(neighbor_vecs[i], neighbor_vecs[resolved]); + std::swap(neighbor_ids[i], neighbor_ids[resolved]); + } + ++resolved; } - neighbor_ids[unvisited_count] = node; - neighbor_vecs[unvisited_count] = vec_ptr; - unvisited_count++; } - // Phase 2: scan remaining neighbors. - for (; i < neighbors.size(); ++i) { - node_id_t node = neighbors[i]; - if (visit.visited(node)) continue; - visit.set_visited(node); - neighbor_ids[unvisited_count] = node; - neighbor_vecs[unvisited_count] = entity.get_vector_ptr(node); - unvisited_count++; + // Submit AIO for miss pages: current unresolved + next hop peek-ahead + { + node_id_t prefetch_ids[128]; + uint32_t pc = 0; + // Current hop's unresolved nodes (their pages will help future hops) + for (uint32_t i = resolved; i < unvisited_count && pc < 64; ++i) { + prefetch_ids[pc++] = neighbor_ids[i]; + } + // Next hop peek-ahead + if (pool.has_next()) { + node_id_t next_node = pool.peek(); + const auto next_neighbors = entity.get_neighbors_typed(0, next_node); + for (uint32_t i = 0; i < next_neighbors.size() && pc < 128; ++i) { + if (!visit.visited(next_neighbors[i])) + prefetch_ids[pc++] = next_neighbors[i]; + } + } + if (pc > 0) entity.submit_prefetch(prefetch_ids, pc); } - if (unvisited_count == 0) continue; - dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + // CPU prefetch + distance computation + if (resolved > 0) { + const uint32_t po = std::min(prefetch_offset, resolved); + 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(), resolved, dists.data()); + } + // Unresolved vectors get FLT_MAX - they won't enter the candidate pool. + for (uint32_t i = resolved; i < unvisited_count; ++i) { + dists[i] = FLT_MAX; + } pool.push_block(dists.data(), neighbor_ids.data(), static_cast(unvisited_count)); + + entity.release_vectors(); } + + // Final harvest to clean up any pending AIO + entity.harvest_prefetch(); } // ============================================================================ @@ -379,9 +476,7 @@ void dual_heap_search_neighbors(const EntityType &entity, level_t level, // search_neighbors: Dispatch to fast or dual-heap path. // // - 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) +// - level-0 unfiltered search → fast_search_neighbors // ============================================================================ template void HnswAlgorithm::search_neighbors(level_t level, @@ -393,7 +488,6 @@ void HnswAlgorithm::search_neighbors(level_t level, HnswDistCalculator &dc = ctx->dist_calculator(); if (!use_pool || ctx->filter().is_valid() || level != 0) { - // Dual-heap path: add_node, filtered search, or upper-level scan. auto run_with_filter = [&](auto &&filter) { dual_heap_search_neighbors( entity, level, entry_point, dist, topk, ctx, dc, @@ -410,36 +504,24 @@ void HnswAlgorithm::search_neighbors(level_t level, run_with_filter(filter); } } else { - // Pool-based path for level-0 unfiltered search. - if constexpr (std::is_same_v) { - const uint32_t prefetch_lines = - ctx->pl() > 0 ? ctx->pl() : (entity.vector_size() + 63) / 64; - - // Fast path: direct pointer access via get_vector_ptr. - // BlockHeap (AVX2) or LinearPool (scalar) for top-k tracking. - 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(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(entity, lpool, visit, dc, topk_v, ef_v, - *entry_point, *dist, prefetch_lines, ctx->po()); - copy_pool_to_topk(lpool, topk); - } + 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(entity, bpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, ctx->po()); + copy_pool_to_topk(bpool, 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); + auto &lpool = ctx->pool(); + fast_search_neighbors(entity, lpool, visit, dc, topk_v, ef_v, + *entry_point, *dist, prefetch_lines, ctx->po()); + copy_pool_to_topk(lpool, topk); } } } diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 8bd03f276..16eb98c03 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -276,6 +276,11 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { if (ret != 0) { return ret; } + + if (entity_->storage_mode() == HnswStorageMode::kBufferPool) { + static_cast(entity_.get()) + ->mark_upper_level_pages(); + } IndexMeta index_meta; ret = entity_->get_index_meta(&index_meta); if (ret == IndexError_NoExist) { diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc index 1ad1ebd9c..84cb04651 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -811,6 +811,40 @@ 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 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 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_, node_chunk_bases_, + upper_neighbor_chunk_bases_); + 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 9c53aaf92..15e9e249f 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -14,7 +14,11 @@ #pragma once +#include +#include +#include #include +#include #include #include #include @@ -864,14 +868,24 @@ class HnswMmapStreamerEntity : public HnswStreamerEntity { return *reinterpret_cast(base + offset); } - //! Direct vector pointer access (no MemoryBlock wrapper). - //! For use in the merged search loop to avoid intermediate allocations. ailego_force_inline const void *get_vector_ptr(node_id_t id) const { uint32_t chunk_idx = id >> node_index_mask_bits_; uint32_t offset = (id & node_index_mask_) * node_size(); return get_node_chunk_base(chunk_idx) + offset; } + ailego_force_inline int resolve_vectors(const node_id_t *ids, uint32_t count, + const void **out) const { + for (uint32_t i = 0; i < count; ++i) out[i] = get_vector_ptr(ids[i]); + return 0; + } + + ailego_force_inline void release_vectors() const {} + void submit_prefetch(const node_id_t *, uint32_t) const {} // no-op + void harvest_prefetch() const {} // no-op + void wait_prefetch() const {} // no-op for mmap + void reset_io_budget(int32_t) const {} // no-op for mmap + protected: //! Get cached base address for a node chunk, syncing if needed ailego_force_inline const char *get_node_chunk_base( @@ -917,7 +931,6 @@ class HnswMmapStreamerEntity : public HnswStreamerEntity { mutable std::vector upper_neighbor_chunk_bases_{}; }; -//! Typed entity subclass for buffer pool mode. class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { public: using MemoryBlock = BufferPoolMemoryBlock; @@ -929,6 +942,8 @@ class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { return HnswStorageMode::kBufferPool; } + 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); @@ -944,6 +959,310 @@ class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { inline key_t get_key_typed(node_id_t id) const { return HnswStreamerEntity::get_key_typed(id); } + + //! Set I/O budget for the current search. + //! budget < 0 → blocking mode (always pread on miss) + //! budget = 0 → non-blocking mode (FLT_MAX on miss) + //! budget > 0 → allow up to N preads per search + void reset_io_budget(int32_t budget) const { io_budget_ = budget; } + + int resolve_vectors(const node_id_t *ids, uint32_t count, + const void **out) const { + ensure_pinned_pages(); + if (ailego_unlikely(!pinned_pages_.bound())) return -1; + const size_t vec_sz = vector_size(); + const size_t pg_sz = ailego::kVectorPageSize; + + cross_page_used_ = 0; + if (cross_page_arena_.size() < count * vec_sz) + cross_page_arena_.resize(count * vec_sz); + for (uint32_t i = 0; i < count; ++i) { + const size_t abs_off = get_vector_abs_offset(ids[i]); + const auto page_id = static_cast(abs_off / pg_sz); + const size_t intra = abs_off % pg_sz; + if (ailego_likely(intra + vec_sz <= pg_sz)) { + char *page = pinned_pages_.try_get_page(page_id); + if (!page) { + // Cache miss: check budget + if (io_budget_ != 0) { + page = pinned_pages_.get_page(page_id); + if (io_budget_ > 0) --io_budget_; + } + if (!page) { out[i] = nullptr; continue; } + } + out[i] = page + intra; + } else { + const size_t part1 = pg_sz - intra; + char *p1 = pinned_pages_.try_get_page(page_id); + char *p2 = pinned_pages_.try_get_page(page_id + 1); + if (!p1 && io_budget_ != 0) { + p1 = pinned_pages_.get_page(page_id); + if (io_budget_ > 0) --io_budget_; + } + if (!p2 && io_budget_ != 0) { + p2 = pinned_pages_.get_page(page_id + 1); + if (io_budget_ > 0) --io_budget_; + } + if (!p1 || !p2) { out[i] = nullptr; continue; } + char *scratch = cross_page_arena_.data() + cross_page_used_ * vec_sz; + ++cross_page_used_; + std::memcpy(scratch, p1 + intra, part1); + std::memcpy(scratch + part1, p2, vec_sz - part1); + out[i] = scratch; + } + } + return 0; + } + + void release_vectors() const { + pinned_pages_.release_all(); + } + + //! Submit non-blocking AIO prefetch for the given node ids' vector pages. + void submit_prefetch(const node_id_t *ids, uint32_t count) const { + auto *pool = vec_buffer_pool(); + if (!pool || !pool->aio_enabled()) return; + const size_t pg_sz = ailego::kVectorPageSize; + const size_t vec_sz = vector_size(); + ailego::block_id_t page_ids[128]; + uint32_t n = 0; + for (uint32_t i = 0; i < count && n < 126; ++i) { + const size_t abs_off = get_vector_abs_offset(ids[i]); + auto pid = static_cast(abs_off / pg_sz); + page_ids[n++] = pid; + // Cross-page: also prefetch the next page + const size_t intra = abs_off % pg_sz; + if (intra + vec_sz > pg_sz) { + page_ids[n++] = pid + 1; + } + } + if (n > 0) pool->submit_aio_async(page_ids, n); + } + + //! Harvest previously submitted async AIO results (non-blocking). + void harvest_prefetch() const { + auto *pool = vec_buffer_pool(); + if (pool) pool->harvest_aio(); + } + + //! Block until every page submitted via submit_prefetch() is resident. + //! After this a hop can resolve all its neighbour pages as cache hits, so + //! the read cost of a whole hop collapses to a single concurrent burst + //! instead of a chain of per-neighbour synchronous preads. + void wait_prefetch() const { + auto *pool = vec_buffer_pool(); + if (pool) pool->wait_aio(); + } + + void mark_upper_level_pages() { + auto *pool = vec_buffer_pool(); + if (!pool) return; + auto ep = entry_point(); + auto max_lvl = cur_max_level(); + if (ep == kInvalidNodeId || max_lvl == 0) return; + + const uint32_t n = doc_cnt(); + std::vector visited(n, false); + std::vector upper_nodes; + upper_nodes.reserve(n / scaling_factor() + 64); + upper_nodes.push_back(ep); + visited[ep] = true; + + for (level_t lvl = max_lvl; lvl >= 1; --lvl) { + for (size_t idx = 0; idx < upper_nodes.size(); ++idx) { + auto id = upper_nodes[idx]; + auto it = upper_neighbor_index_->find(id); + if (it == upper_neighbor_index_->end()) continue; + auto meta = + reinterpret_cast(&it->second); + if (lvl > static_cast(meta->bits.level)) continue; + auto neighbors = get_neighbors_typed(lvl, id); + for (uint32_t i = 0; i < neighbors.size(); ++i) { + auto nid = neighbors[i]; + if (nid < n && !visited[nid]) { + visited[nid] = true; + upper_nodes.push_back(nid); + } + } + } + } + + const size_t pg_sz = ailego::kVectorPageSize; + const size_t vec_sz = vector_size(); + std::vector page_ids; + page_ids.reserve(upper_nodes.size()); + for (auto id : upper_nodes) { + const size_t abs_off = get_vector_abs_offset(id); + page_ids.push_back(static_cast(abs_off / pg_sz)); + const size_t intra = abs_off % pg_sz; + if (intra + vec_sz > pg_sz) { + page_ids.push_back(static_cast(abs_off / pg_sz) + + 1); + } + } + std::sort(page_ids.begin(), page_ids.end()); + page_ids.erase(std::unique(page_ids.begin(), page_ids.end()), + page_ids.end()); + + size_t marked = 0; + for (auto pid : page_ids) { + pool->page_table_.set_evict_priority(pid, 2); + char *buf = pool->acquire_buffer(pid, 50); + if (buf) { + pool->page_table_.release_block(pid); + ++marked; + } + } + LOG_DEBUG( + "mark_upper_level_pages: marked %zu/%zu pages for %zu upper-level " + "nodes (maxLevel=%d, priority=2)", + marked, page_ids.size(), upper_nodes.size(), (int)max_lvl); + } + + private: + struct PinnedPageSet { + // Open-addressing set of pinned pages for a single search hop. It must be + // able to hold every distinct page a hop touches, otherwise get_page() + // returns nullptr on overflow and the corresponding neighbor is silently + // dropped -- which measurably lowers recall vs mmap. A hop resolves up to + // l0_neighbor_cnt() vectors and each vector may straddle a 4K page + // boundary (two pages), so the worst case is 2 * l0_neighbor_cnt() pages. + // The table is therefore sized from the entity's neighbor count at bind() + // time instead of a fixed constant (see reserve_for()). + static constexpr size_t kMinCapacity = 128; + static constexpr ailego::block_id_t kEmpty = + std::numeric_limits::max(); + + PinnedPageSet() = default; + ~PinnedPageSet() { + release_all(); + } + PinnedPageSet(const PinnedPageSet &) = delete; + PinnedPageSet &operator=(const PinnedPageSet &) = delete; + + //! Bind to a pool and size the table to hold up to "max_pages" distinct + //! entries without ever hitting the load-factor cap. + void bind(ailego::VecBufferPool *pool, size_t max_pages) { + pool_ = pool; + reserve_for(max_pages); + } + bool bound() const { + return pool_ != nullptr; + } + + //! Try to get a page WITHOUT triggering disk I/O. + //! Returns buffer if page is in PinnedPageSet or already in pool memory. + //! Returns nullptr if page would need a pread (cache miss) or set is full. + char *try_get_page(ailego::block_id_t page_id) { + size_t slot = static_cast(page_id) & mask_; + for (size_t probe = 0; probe < capacity_; ++probe) { + if (ids_[slot] == page_id) return bufs_[slot]; + if (ids_[slot] == kEmpty) { + if (ailego_unlikely(count_ >= max_load_)) { + return nullptr; + } + char *buf = pool_->try_acquire_buffer(page_id); + if (!buf) return nullptr; // page not in memory, skip + ids_[slot] = page_id; + bufs_[slot] = buf; + ++count_; + return buf; + } + slot = (slot + 1) & mask_; + } + return nullptr; + } + + //! Get a page WITH blocking pread if not in cache. + //! Returns nullptr only if PinnedPageSet is full or pool alloc fails. + char *get_page(ailego::block_id_t page_id) { + size_t slot = static_cast(page_id) & mask_; + for (size_t probe = 0; probe < capacity_; ++probe) { + if (ids_[slot] == page_id) return bufs_[slot]; + if (ids_[slot] == kEmpty) { + if (ailego_unlikely(count_ >= max_load_)) { + return nullptr; + } + char *buf = pool_->acquire_buffer(page_id, 50); + if (ailego_unlikely(!buf)) return nullptr; + ids_[slot] = page_id; + bufs_[slot] = buf; + ++count_; + return buf; + } + slot = (slot + 1) & mask_; + } + return nullptr; + } + + void release_all() { + if (!pool_ || count_ == 0) return; + for (size_t i = 0; i < capacity_; ++i) { + if (ids_[i] != kEmpty) { + pool_->page_table_.release_block(ids_[i]); + ids_[i] = kEmpty; + bufs_[i] = nullptr; + } + } + count_ = 0; + } + + private: + //! Size the table so its 75% load-factor cap covers "max_pages" entries. + //! Capacity is rounded up to a power of two so the mask-based probe works. + void reserve_for(size_t max_pages) { + size_t need = (max_pages * 4 + 2) / 3 + 1; // invert 3/4 load factor + size_t cap = kMinCapacity; + while (cap < need) cap <<= 1; + if (cap == capacity_ && !ids_.empty()) { + release_all(); + return; + } + capacity_ = cap; + mask_ = cap - 1; + max_load_ = cap * 3 / 4; + ids_.assign(cap, kEmpty); + bufs_.assign(cap, nullptr); + count_ = 0; + } + + ailego::VecBufferPool *pool_{nullptr}; + std::vector ids_{}; + std::vector bufs_{}; + size_t capacity_{0}; + size_t mask_{0}; + size_t max_load_{0}; + size_t count_{0}; + }; + + ailego::VecBufferPool *vec_buffer_pool() const { + if (broker_ && broker_->storage()) { + return broker_->storage()->vec_buffer_pool(); + } + return nullptr; + } + + size_t get_vector_abs_offset(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + return node_chunks_[loc.first]->abs_data_offset() + loc.second; + } + + void ensure_pinned_pages() const { + if (!pinned_pages_.bound()) { + auto *pool = vec_buffer_pool(); + if (pool) { + // A single hop resolves up to l0_neighbor_cnt() vectors, each of which + // may straddle a 4K page boundary (two pages). Size the set to the + // worst case so a hop never overflows and silently drops neighbors. + pinned_pages_.bind(pool, 2 * l0_neighbor_cnt() + 2); + } + } + } + + mutable PinnedPageSet pinned_pages_; + mutable std::vector cross_page_arena_; + mutable uint32_t cross_page_used_{0}; + mutable int32_t io_budget_{0}; // <0: blocking, 0: non-blocking, >0: limited }; //! Typed entity subclass for contiguous memory mode. @@ -1053,18 +1372,26 @@ class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity { return HnswMmapStreamerEntity::get_key_typed(id); } - //! Direct vector pointer from flat vector array (stride = vector_size). - //! For use in the merged search loop to avoid intermediate allocations. ailego_force_inline const void *get_vector_ptr(node_id_t id) const { if (ailego_likely(vector_base_ != nullptr)) { return vector_base_ + static_cast(id) * vector_size(); } - // Fallback to mmap chunk-based access uint32_t chunk_idx = id >> node_index_mask_bits_; uint32_t offset = (id & node_index_mask_) * node_size(); return get_node_chunk_base(chunk_idx) + offset; } + ailego_force_inline int resolve_vectors(const node_id_t *ids, uint32_t count, + const void **out) const { + for (uint32_t i = 0; i < count; ++i) out[i] = get_vector_ptr(ids[i]); + return 0; + } + + ailego_force_inline void release_vectors() const {} + void submit_prefetch(const node_id_t *, uint32_t) const {} // no-op + void harvest_prefetch() const {} // no-op + void reset_io_budget(int32_t) const {} // no-op for contiguous + protected: //! Custom deleter for contiguous memory (munmap / _aligned_free / free) //! Used by shared_ptr to properly release mmap'd memory. diff --git a/src/core/algorithm/ivf/ivf_dumper.cc b/src/core/algorithm/ivf/ivf_dumper.cc index c4edcfea5..4fd3b5649 100644 --- a/src/core/algorithm/ivf/ivf_dumper.cc +++ b/src/core/algorithm/ivf/ivf_dumper.cc @@ -93,13 +93,13 @@ int IVFDumper::dump_container_segment(const IndexStorage::Pointer &container, const size_t total_size = seg->data_size() + seg->padding_size(); size_t off = 0; while (off < total_size) { - const void *data = nullptr; + IndexStorage::MemoryBlock block; size_t rd_size = std::min(batch_size, total_size - off); - if (seg->read(off, &data, rd_size) != rd_size) { + if (seg->read(off, block, rd_size) != rd_size) { LOG_ERROR("Failed to read data, off=%zu size=%zu", off, rd_size); return IndexError_ReadData; } - if (dumper_->write(data, rd_size) != rd_size) { + if (dumper_->write(block.data(), rd_size) != rd_size) { LOG_ERROR("Failed to write data, size=%zu", rd_size); return IndexError_WriteData; } diff --git a/src/core/algorithm/ivf/ivf_entity.cc b/src/core/algorithm/ivf/ivf_entity.cc index 2c3b4b195..d9c46f7a0 100644 --- a/src/core/algorithm/ivf/ivf_entity.cc +++ b/src/core/algorithm/ivf/ivf_entity.cc @@ -12,8 +12,49 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "ivf_entity.h" +#include +#include #include +#include #include "ivf_utility.h" +#ifdef _OPENMP +#include +#endif + +namespace { +std::atomic g_active_batch_callers{0}; + +struct BatchCallerGuard { + BatchCallerGuard() { + g_active_batch_callers.fetch_add(1, std::memory_order_relaxed); + } + ~BatchCallerGuard() { + g_active_batch_callers.fetch_sub(1, std::memory_order_relaxed); + } +}; + +// 0=on-demand, 1=unconditional prefetch, 2=adaptive (default). Test switch. +inline int prefetch_mode() { + static int mode = []() { + const char *e = std::getenv("ZVEC_PREFETCH_MODE"); + return (e && *e) ? std::atoi(e) : 2; + }(); + return mode; +} + +#ifdef _OPENMP +inline int adaptive_thread_count(size_t query_count) { + if (query_count <= 4) { + return 1; + } + int hw = omp_get_max_threads(); + if (hw <= 0) hw = 1; + int active = g_active_batch_callers.load(std::memory_order_relaxed); + if (active <= 0) active = 1; + return std::max(1, hw / active); +} +#endif +} // namespace namespace zvec { namespace core { @@ -462,12 +503,14 @@ int IVFEntity::load_header(const IndexStorage::Pointer &container) { IVF_INVERTED_HEADER_SEG_ID.c_str()); return IndexError_InvalidFormat; } - const void *data = nullptr; - if (header->read(0, &data, header->data_size()) != header->data_size()) { + IndexStorage::MemoryBlock header_block; + if (header->read(0, header_block, header->data_size()) != + header->data_size()) { LOG_ERROR("Failed to read data, segment %s", IVF_INVERTED_HEADER_SEG_ID.c_str()); return IndexError_ReadData; } + const void *data = header_block.data(); std::memcpy(&header_, data, sizeof(header_)); if (header_.header_size < sizeof(header_) + header_.index_meta_size || header_.header_size > header->data_size()) { @@ -606,14 +649,17 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, IndexContext::Stats *context_stats) const { ailego_assert_with(inverted_list_id < header_.inverted_list_count, "invalid id"); - auto list_meta = this->inverted_list_meta(inverted_list_id); + IndexStorage::MemoryBlock meta_block; + auto list_meta = this->inverted_list_meta(inverted_list_id, meta_block); ivf_assert(list_meta, IndexError_ReadData); - const void *data = nullptr; + const size_t block_size = header_.block_size; + + IndexStorage::MemoryBlock data_block; + IndexStorage::MemoryBlock keys_block; const size_t block_vecs = header_.block_vector_count; std::vector distances(block_vecs); const size_t batch_size = kBatchBlocks; - const size_t block_size = header_.block_size; const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); for (size_t i = 0; i < list_meta->block_count; i += batch_size) { //! Read vecs @@ -622,15 +668,17 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, const size_t size = std::min(blocks * block_size, static_cast(header_.inverted_body_size - off)); - if (inverted_->read(off, &data, size) != size) { + if (inverted_->read(off, data_block, size) != size) { LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); return IndexError_ReadData; } + const void *data = data_block.data(); //! Read keys size_t items = std::min(blocks * block_vecs, list_meta->vector_count - (i * block_vecs)); - auto keys = get_keys(list_meta->id_offset + i * block_vecs, items); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items, + keys_block); if (!keys) { return IndexError_ReadData; } @@ -680,14 +728,17 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, IndexContext::Stats *context_stats) const { ailego_assert_with(inverted_list_id < header_.inverted_list_count, "invalid id"); - auto list_meta = inverted_list_meta(inverted_list_id); + IndexStorage::MemoryBlock meta_block; + auto list_meta = inverted_list_meta(inverted_list_id, meta_block); ivf_assert(list_meta, IndexError_ReadData); - const void *data = nullptr; + const size_t block_size = header_.block_size; + + IndexStorage::MemoryBlock data_block; + IndexStorage::MemoryBlock keys_block; const size_t block_vecs = header_.block_vector_count; std::vector distances(block_vecs); const size_t batch_size = kBatchBlocks; - const size_t block_size = header_.block_size; const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); for (size_t i = 0; i < list_meta->block_count; i += batch_size) { //! Read vecs @@ -696,15 +747,17 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, const size_t size = std::min(blocks * block_size, static_cast(header_.inverted_body_size - off)); - if (inverted_->read(off, &data, size) != size) { + if (inverted_->read(off, data_block, size) != size) { LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); return IndexError_ReadData; } + const void *data = data_block.data(); //! Read keys size_t items = std::min(blocks * block_vecs, list_meta->vector_count - (i * block_vecs)); - auto keys = get_keys(list_meta->id_offset + i * block_vecs, items); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items, + keys_block); if (!keys) { return IndexError_ReadData; } @@ -731,6 +784,213 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, return 0; } +//! Block-level batch search: scan cluster once, compute for all queries (with +//! filter) +int IVFEntity::search_batch(size_t inverted_list_id, const IndexFilter &filter, + BatchQueryItem *items, size_t query_count, + uint32_t *scan_count) const { + ailego_assert_with(inverted_list_id < header_.inverted_list_count, + "invalid id"); + IndexStorage::MemoryBlock meta_block; + auto list_meta = this->inverted_list_meta(inverted_list_id, meta_block); + ivf_assert(list_meta, IndexError_ReadData); + + const size_t block_size = header_.block_size; + // Adaptive prefetch admission: prefetch the whole cluster only if it fits in + // this caller's fair share of the pool's free space. Avoids partial prefetch + // under memory pressure (which evicts still-useful pages and lowers QPS). + // The divisor is the number of concurrent batch callers because the buffer + // pool is shared process-wide. + { + const size_t total_data_size = list_meta->block_count * block_size; + const int pmode = prefetch_mode(); + if (pmode == 1) { + inverted_->prefetch(list_meta->offset, total_data_size); + } else if (pmode == 2) { + int callers = g_active_batch_callers.load(std::memory_order_relaxed); + if (callers < 1) callers = 1; + const size_t budget = + inverted_->prefetch_budget() / static_cast(callers); + if (total_data_size <= budget) { + inverted_->prefetch(list_meta->offset, total_data_size); + } + } + } + + const void *data = nullptr; + const size_t block_vecs = header_.block_vector_count; + const size_t batch_size = kBatchBlocks; + const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); + BatchCallerGuard caller_guard; +#ifdef _OPENMP + const int omp_threads = adaptive_thread_count(query_count); +#endif + + IndexStorage::MemoryBlock data_block; + IndexStorage::MemoryBlock keys_block; + for (size_t i = 0; i < list_meta->block_count; i += batch_size) { + //! Read vecs - ONCE for all queries + const size_t off = list_meta->offset + i * block_size; + const size_t blocks = std::min(batch_size, list_meta->block_count - i); + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + if (inverted_->read(off, data_block, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + data = data_block.data(); + + //! Read keys - ONCE for all queries + size_t items_count = std::min(blocks * block_vecs, + list_meta->vector_count - (i * block_vecs)); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items_count, + keys_block); + if (!keys) { + return IndexError_ReadData; + } + + //! For each block, compute distances for ALL queries + for (size_t b = 0; b < blocks; ++b) { + const size_t vecs_count = + std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs); + auto block_keys = keys + b * block_vecs; + + size_t keeps = 0; + ailego_assert_with(block_vecs < sizeof(keeps) * 8, "bits overflow"); + for (size_t k = 0; k < vecs_count; ++k) { + if (!filter(block_keys[k])) { + keeps |= (1ULL << k); + } + } + size_t filtered = vecs_count - ailego_popcount64(keeps); + + if (keeps == 0) { + for (size_t q = 0; q < query_count; ++q) { + *(items[q].stats->mutable_filtered_count()) += filtered; + } + continue; + } + + const void *block_data = static_cast(data) + b * block_size; + uint32_t id_off = list_meta->id_offset + (i + b) * block_vecs; + +#ifdef _OPENMP +#pragma omp parallel for num_threads(omp_threads) \ + schedule(static) if (omp_threads > 1) +#endif + for (size_t q = 0; q < query_count; ++q) { + std::vector local_distances(block_vecs); + calculator_->query_features_distance( + items[q].query, block_data, vecs_count, local_distances.data()); + *(items[q].stats->mutable_dist_calced_count()) += vecs_count; + *(items[q].stats->mutable_filtered_count()) += filtered; + + for (size_t k = 0; k < vecs_count; ++k) { + if ((keeps & (1ULL << k)) && block_keys[k] != kInvalidKey) { + items[q].heap->emplace(block_keys[k], local_distances[k] * norm_val, + id_off + k); + } + } + } + } + } + + *scan_count = list_meta->vector_count; + return 0; +} + +//! Block-level batch search without filter +int IVFEntity::search_batch(size_t inverted_list_id, BatchQueryItem *items, + size_t query_count, uint32_t *scan_count) const { + ailego_assert_with(inverted_list_id < header_.inverted_list_count, + "invalid id"); + IndexStorage::MemoryBlock meta_block; + auto list_meta = inverted_list_meta(inverted_list_id, meta_block); + ivf_assert(list_meta, IndexError_ReadData); + + const size_t block_size = header_.block_size; + // Adaptive prefetch admission (see the filtered search_batch above). + { + const size_t total_data_size = list_meta->block_count * block_size; + const int pmode = prefetch_mode(); + if (pmode == 1) { + inverted_->prefetch(list_meta->offset, total_data_size); + } else if (pmode == 2) { + int callers = g_active_batch_callers.load(std::memory_order_relaxed); + if (callers < 1) callers = 1; + const size_t budget = + inverted_->prefetch_budget() / static_cast(callers); + if (total_data_size <= budget) { + inverted_->prefetch(list_meta->offset, total_data_size); + } + } + } + + const void *data = nullptr; + const size_t block_vecs = header_.block_vector_count; + const size_t batch_size = kBatchBlocks; + const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); + BatchCallerGuard caller_guard; +#ifdef _OPENMP + const int omp_threads = adaptive_thread_count(query_count); +#endif + + IndexStorage::MemoryBlock data_block; + IndexStorage::MemoryBlock keys_block; + for (size_t i = 0; i < list_meta->block_count; i += batch_size) { + //! Read vecs - ONCE for all queries + const size_t off = list_meta->offset + i * block_size; + const size_t blocks = std::min(batch_size, list_meta->block_count - i); + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + if (inverted_->read(off, data_block, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + data = data_block.data(); + + //! Read keys - ONCE for all queries + size_t items_count = std::min(blocks * block_vecs, + list_meta->vector_count - (i * block_vecs)); + auto keys = get_keys(list_meta->id_offset + i * block_vecs, items_count, + keys_block); + if (!keys) { + return IndexError_ReadData; + } + + //! For each block, compute distances for ALL queries + for (size_t b = 0; b < blocks; ++b) { + const size_t vecs_count = + std::min(block_vecs, list_meta->vector_count - (i + b) * block_vecs); + auto block_keys = keys + b * block_vecs; + const void *block_data = static_cast(data) + b * block_size; + uint32_t id_off = list_meta->id_offset + (i + b) * block_vecs; + +#ifdef _OPENMP +#pragma omp parallel for num_threads(omp_threads) \ + schedule(static) if (omp_threads > 1) +#endif + for (size_t q = 0; q < query_count; ++q) { + std::vector local_distances(block_vecs); + calculator_->query_features_distance( + items[q].query, block_data, vecs_count, local_distances.data()); + for (size_t k = 0; k < vecs_count; ++k) { + if (block_keys[k] != kInvalidKey) { + items[q].heap->emplace(block_keys[k], local_distances[k] * norm_val, + id_off + k); + } + } + *(items[q].stats->mutable_dist_calced_count()) += vecs_count; + } + } + } + + *scan_count = list_meta->vector_count; + return 0; +} + //! search all inverted list with filter int IVFEntity::search(const void *query, const IndexFilter &filter, IndexDocumentHeap *heap, @@ -762,48 +1022,53 @@ int IVFEntity::search(const void *query, IndexDocumentHeap *heap, const void *IVFEntity::get_vector(size_t id) const { if (features_) { - const void *data = nullptr; size_t element_size = features_->data_size() / vector_count(); size_t off = id * element_size; - if (features_->read(off, &data, element_size) != element_size) { + IndexStorage::MemoryBlock block; + if (features_->read(off, block, element_size) != element_size) { LOG_ERROR("Failed to read segment, off=%zu size=%zu", off, element_size); return nullptr; } - return data; + vector_.assign(static_cast(block.data()), element_size); + return vector_.data(); } - const void *data = nullptr; + IndexStorage::MemoryBlock off_block; size_t size = sizeof(InvertedVecLocation); - if (offsets_->read(id * size, &data, size) != size) { + if (offsets_->read(id * size, off_block, size) != size) { LOG_ERROR("Failed to read offsets segment, id=%zu", id); return nullptr; } - auto &loc = *reinterpret_cast(data); + InvertedVecLocation loc = + *reinterpret_cast(off_block.data()); if (loc.column_major) { vector_.resize(meta_.element_size()); auto unit_size = IndexMeta::AlignSizeof(meta_.data_type()); size_t cols = meta_.element_size() / unit_size; size_t step = block_vector_count() * unit_size; size_t rd_size = step * (cols - 1) + unit_size; - if (inverted_->read(loc.offset, &data, rd_size) != rd_size) { + IndexStorage::MemoryBlock block; + if (inverted_->read(loc.offset, block, rd_size) != rd_size) { LOG_ERROR("Failed to read data, off=%zu size=%zu", static_cast(loc.offset), rd_size); return nullptr; } + const char *data = static_cast(block.data()); for (size_t c = 0; c < cols; ++c) { - vector_.replace(c * unit_size, unit_size, - reinterpret_cast(data) + c * step, - unit_size); + vector_.replace(c * unit_size, unit_size, data + c * step, unit_size); } return vector_.data(); } else { - if (inverted_->read(loc.offset, &data, meta_.element_size()) != + IndexStorage::MemoryBlock block; + if (inverted_->read(loc.offset, block, meta_.element_size()) != meta_.element_size()) { LOG_ERROR("Failed to read data, off=%zu size=%u", static_cast(loc.offset), meta_.element_size()); return nullptr; } - return data; + vector_.assign(static_cast(block.data()), + meta_.element_size()); + return vector_.data(); } } @@ -825,23 +1090,23 @@ int IVFEntity::get_vector(size_t id, IndexStorage::MemoryBlock &block) const { LOG_ERROR("Failed to read offsets segment, id=%zu", id); return IndexError_Runtime; } - const void *data = data_block.data(); - auto &loc = *reinterpret_cast(data); + InvertedVecLocation loc = + *reinterpret_cast(data_block.data()); if (loc.column_major) { vector_.resize(meta_.element_size()); auto unit_size = IndexMeta::AlignSizeof(meta_.data_type()); size_t cols = meta_.element_size() / unit_size; size_t step = block_vector_count() * unit_size; size_t rd_size = step * (cols - 1) + unit_size; - if (inverted_->read(loc.offset, &data, rd_size) != rd_size) { + IndexStorage::MemoryBlock vec_block; + if (inverted_->read(loc.offset, vec_block, rd_size) != rd_size) { LOG_ERROR("Failed to read data, off=%zu size=%zu", static_cast(loc.offset), rd_size); return IndexError_Runtime; } + const char *data = static_cast(vec_block.data()); for (size_t c = 0; c < cols; ++c) { - vector_.replace(c * unit_size, unit_size, - reinterpret_cast(data) + c * step, - unit_size); + vector_.replace(c * unit_size, unit_size, data + c * step, unit_size); } block.reset(vector_.data()); return 0; @@ -860,23 +1125,23 @@ uint32_t IVFEntity::key_to_id(uint64_t key) const { //! Do binary search uint32_t start = 0UL; uint32_t end = vector_count(); - const void *data = nullptr; uint32_t idx = 0u; while (start < end) { idx = start + (end - start) / 2; - if (ailego_unlikely(mapping_->read(idx * sizeof(uint32_t), &data, + IndexStorage::MemoryBlock map_block; + if (ailego_unlikely(mapping_->read(idx * sizeof(uint32_t), map_block, sizeof(uint32_t)) != sizeof(uint32_t))) { LOG_ERROR("Failed to read mapping segment, idx=%u", idx); return std::numeric_limits::max(); } - const uint64_t *mkey; - uint32_t local_id = *reinterpret_cast(data); - if (ailego_unlikely(keys_->read(local_id * sizeof(uint64_t), - (const void **)(&mkey), + uint32_t local_id = *reinterpret_cast(map_block.data()); + IndexStorage::MemoryBlock key_block; + if (ailego_unlikely(keys_->read(local_id * sizeof(uint64_t), key_block, sizeof(uint64_t)) != sizeof(uint64_t))) { LOG_ERROR("Read key from segment failed"); return std::numeric_limits::max(); } + const uint64_t *mkey = static_cast(key_block.data()); if (*mkey < key) { start = idx + 1; } else if (*mkey > key) { diff --git a/src/core/algorithm/ivf/ivf_entity.h b/src/core/algorithm/ivf/ivf_entity.h index 50aa7edbe..50e6887ec 100644 --- a/src/core/algorithm/ivf/ivf_entity.h +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -53,6 +53,22 @@ class IVFEntity { int search(size_t inverted_list_id, const void *query, uint32_t *scan_count, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + //! Batch query item for block-level parallel search + struct BatchQueryItem { + const void *query; + IndexDocumentHeap *heap; + IndexContext::Stats *stats; + }; + + //! Block-level batch search: scan cluster once, compute for all queries + int search_batch(size_t inverted_list_id, const IndexFilter &filter, + BatchQueryItem *items, size_t query_count, + uint32_t *scan_count) const; + + //! Block-level batch search without filter + int search_batch(size_t inverted_list_id, BatchQueryItem *items, + size_t query_count, uint32_t *scan_count) const; + //! search all inverted list with filter int search(const void *query, const IndexFilter &filter, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; @@ -107,8 +123,10 @@ class IVFEntity { //! Retrieve a block of vectors const void *read_block(size_t inverted_list_id, size_t local_block_id, - size_t *vecs_count) const { - auto iv_meta = this->inverted_list_meta(inverted_list_id); + size_t *vecs_count, + IndexStorage::MemoryBlock &block) const { + IndexStorage::MemoryBlock meta_block; + auto iv_meta = this->inverted_list_meta(inverted_list_id, meta_block); if (!iv_meta || local_block_id >= iv_meta->block_count) { LOG_ERROR("Failed to read inverted list, listId=%zu blockIdx=%zu", inverted_list_id, local_block_id); @@ -122,66 +140,64 @@ class IVFEntity { "invalid vecs"); const size_t off = iv_meta->offset + local_block_id * header_.block_size; const size_t size = *vecs_count * meta_.element_size(); - const void *data = nullptr; - if (inverted_->read(off, &data, size) != size) { + if (inverted_->read(off, block, size) != size) { LOG_ERROR("Failed to read block off=%zu size=%zu", off, size); return nullptr; } - return data; + return block.data(); } //! Retrieve the inverted list meta - const InvertedListMeta *inverted_list_meta(size_t inverted_list_id) const { - const void *data = nullptr; + const InvertedListMeta *inverted_list_meta( + size_t inverted_list_id, IndexStorage::MemoryBlock &block) const { const size_t size = sizeof(InvertedListMeta); const size_t offset = inverted_list_id * size; - if (inverted_meta_->read(offset, &data, size) != size) { + if (inverted_meta_->read(offset, block, size) != size) { LOG_ERROR("Failed to read inverted meta, id=%zu, size=%zu", inverted_list_id, size); return nullptr; } - return static_cast(data); + return static_cast(block.data()); } - //! Retrieve the keys by consecutive local ids - const uint64_t *get_keys(size_t id, size_t count) const { - const void *data = nullptr; + //! Retrieve the keys with RAII page management (no permanent pin) + const uint64_t *get_keys(size_t id, size_t count, + IndexStorage::MemoryBlock &block) const { const size_t offset = id * sizeof(uint64_t); const size_t size = count * sizeof(uint64_t); - if (keys_->read(offset, &data, size) != size) { + if (keys_->read(offset, block, size) != size) { LOG_ERROR("Failed to read keys, id=%zu, size=%zu", id, size); return nullptr; } - - return static_cast(data); + return static_cast(block.data()); } //! Retrieve the key by local id uint64_t get_key(size_t id) const { - const void *data = nullptr; + IndexStorage::MemoryBlock block; const size_t offset = id * sizeof(uint64_t); const size_t size = sizeof(uint64_t); - if (keys_->read(offset, &data, size) != size) { + if (keys_->read(offset, block, size) != size) { LOG_ERROR("Failed to read key, id=%zu", id); return kInvalidKey; } - return *static_cast(data); + return *static_cast(block.data()); } //! Retrieve the key-order mapping (sorted rank -> local_id). //! mapping[rank] is the local_id of the vector with the rank-th smallest //! key. Returns nullptr if mapping segment is unavailable. - const uint32_t *get_key_order_mapping() const { + const uint32_t *get_key_order_mapping( + IndexStorage::MemoryBlock &block) const { if (!mapping_) return nullptr; - const void *data = nullptr; const size_t size = vector_count() * sizeof(uint32_t); - if (mapping_->read(0, &data, size) != size) { + if (mapping_->read(0, block, size) != size) { return nullptr; } - return static_cast(data); + return static_cast(block.data()); } //! Retrieve vector by local id @@ -221,15 +237,16 @@ class IVFEntity { // ailego_assert_with(integer_quantizer_params_, "nullptr"); if (integer_quantizer_params_ != nullptr) { - const void *data = nullptr; + IndexStorage::MemoryBlock block; size_t size = sizeof(InvertedIntegerQuantizerParams); size_t off = inverted_list_id * size; - if (integer_quantizer_params_->read(off, &data, size) != size) { + if (integer_quantizer_params_->read(off, block, size) != size) { LOG_ERROR("Failed to read data from segment, off=%zu", off); return 1.0f; } auto scale = - static_cast(data)->scale; + static_cast(block.data()) + ->scale; return this->convert_to_normalize_value(scale); } diff --git a/src/core/algorithm/ivf/ivf_index_provider.h b/src/core/algorithm/ivf/ivf_index_provider.h index 1b9e9fe25..3c461f303 100644 --- a/src/core/algorithm/ivf/ivf_index_provider.h +++ b/src/core/algorithm/ivf/ivf_index_provider.h @@ -74,7 +74,7 @@ class IVFIndexProvider : public IndexProvider { public: SortedIterator(const IVFEntity::Pointer &entity) : entity_(entity) { count_ = entity_->vector_count(); - mapping_ = entity_->get_key_order_mapping(); + mapping_ = entity_->get_key_order_mapping(mapping_block_); if (!mapping_) { // Fallback: compute sorting if mapping segment is unavailable fallback_.resize(count_); @@ -114,6 +114,7 @@ class IVFIndexProvider : public IndexProvider { //! Members IVFEntity::Pointer entity_; + IndexStorage::MemoryBlock mapping_block_; // owns the mapping segment data const uint32_t *mapping_{nullptr}; // points into mapping_ segment data std::vector fallback_; // used only if mapping_ unavailable size_t count_{0}; diff --git a/src/core/algorithm/ivf/ivf_searcher.cc b/src/core/algorithm/ivf/ivf_searcher.cc index e8cda56b9..9a1a360c1 100644 --- a/src/core/algorithm/ivf/ivf_searcher.cc +++ b/src/core/algorithm/ivf/ivf_searcher.cc @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "ivf_searcher.h" +#include +#include #include #include #include "ivf_centroid_index.h" @@ -208,7 +210,6 @@ int IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, ivf_check_error_code(ret); } - ctx->reset_results(count); auto &entity = ctx->entity(); auto &filter = ctx->filter(); @@ -221,9 +222,11 @@ int IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); ivf_check_with_msg(ret, "Failed to transform querys"); - for (size_t q = 0; q < count; ++q) { - auto ¢roids = centroid_index_ctx->result(q); - auto &context_stats = ctx->mutable_stats(q); + if (count == 1) { + // Single query: use original query-first path + ctx->reset_results(1); + auto ¢roids = centroid_index_ctx->result(0); + auto &context_stats = ctx->mutable_stats(0); auto &heap = ctx->mutable_result_heap(); heap.clear(); size_t total_scan_count = 0; @@ -242,15 +245,95 @@ int IVFSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, IndexError::What(ret)); total_scan_count += scan_count; } - heap.sort(); // sort the results + heap.sort(); + if (!filter.is_valid()) { + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(0, &heap); + ctx->topk_to_result(0); + return 0; + } + + // Cluster-first batch path: group queries by cluster, load each cluster once + ctx->reset_batch_heaps(count); + + // Build cluster_id -> list of query indices + struct QueryRef { + uint32_t q_idx; + const void *query_ptr; + }; + std::unordered_map> cluster_queries; + std::vector per_query_scan(count, 0); + + const char *qbase = static_cast(query); + for (size_t q = 0; q < count; ++q) { + auto ¢roids = centroid_index_ctx->result(q); + const void *qptr = qbase + q * iv_qmeta.element_size(); + for (size_t i = 0; i < centroids.size(); ++i) { + auto cid = centroids[i].key(); + cluster_queries[cid].push_back({static_cast(q), qptr}); + } + } + + // Iterate by cluster: block-level multi-query parallel search + // Sort cluster IDs for stable LRU behavior and sequential I/O access. + std::vector sorted_cids; + sorted_cids.reserve(cluster_queries.size()); + for (auto &kv : cluster_queries) sorted_cids.push_back(kv.first); + std::sort(sorted_cids.begin(), sorted_cids.end()); + + for (size_t cid : sorted_cids) { + auto &qrefs = cluster_queries[cid]; + + // Build BatchQueryItems for this cluster (skip queries exceeding scan + // limit) + std::vector batch_items; + std::vector batch_q_indices; + batch_items.reserve(qrefs.size()); + batch_q_indices.reserve(qrefs.size()); + + for (auto &ref : qrefs) { + if (per_query_scan[ref.q_idx] >= ctx->max_scan_count()) { + continue; + } + IVFEntity::BatchQueryItem item; + item.query = ref.query_ptr; + item.heap = &ctx->mutable_batch_heap(ref.q_idx); + item.stats = &ctx->mutable_stats(ref.q_idx); + batch_items.push_back(item); + batch_q_indices.push_back(ref.q_idx); + } + + if (batch_items.empty()) continue; + + // Single call scans the cluster once for all queries + uint32_t scan_count = 0; + if (!filter.is_valid()) { + ret = entity->search_batch(cid, batch_items.data(), batch_items.size(), + &scan_count); + } else { + ret = entity->search_batch(cid, filter, batch_items.data(), + batch_items.size(), &scan_count); + } + ivf_check_with_msg(ret, "Failed to search_batch in entity for %s", + IndexError::What(ret)); + + // Update per-query scan counts + for (auto q_idx : batch_q_indices) { + per_query_scan[q_idx] += scan_count; + } + } + + // Finalize results for each query + for (size_t q = 0; q < count; ++q) { + auto &heap = ctx->mutable_batch_heap(q); if (!filter.is_valid()) { - // mapping the local id to key if query without filter ret = entity->retrieve_keys(&heap); ivf_check_error_code(ret); } entity->normalize(q, &heap); - ctx->topk_to_result(q); - query = static_cast(query) + iv_qmeta.element_size(); + ctx->batch_topk_to_result(static_cast(q)); } return 0; diff --git a/src/core/algorithm/ivf/ivf_searcher_context.h b/src/core/algorithm/ivf/ivf_searcher_context.h index 594e04533..6e85acf37 100644 --- a/src/core/algorithm/ivf/ivf_searcher_context.h +++ b/src/core/algorithm/ivf/ivf_searcher_context.h @@ -160,6 +160,51 @@ class IVFSearcherContext : public IndexSearcher::Context { result_heap_.set_threshold(this->threshold()); } + //! Reset batch heaps for cluster-first batch search + void reset_batch_heaps(size_t qnum) { + reset_results(qnum); + batch_heaps_.resize(qnum); + for (size_t i = 0; i < qnum; ++i) { + batch_heaps_[i].clear(); + batch_heaps_[i].limit(topk_); + batch_heaps_[i].set_threshold(this->threshold()); + } + } + + //! Get mutable batch heap for a specific query + IndexDocumentHeap &mutable_batch_heap(size_t q) { + ailego_assert_with(q < batch_heaps_.size(), "invalid q"); + return batch_heaps_[q]; + } + + //! Convert batch heap to final result for a query + void batch_topk_to_result(uint32_t q) { + auto &heap = batch_heaps_[q]; + if (ailego_unlikely(heap.size() == 0)) { + return; + } + + ailego_assert_with(q < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(heap.size())); + heap.sort(); + results_[q].clear(); + for (int i = 0; i < size; ++i) { + auto score = heap[i].score(); + if (score > this->threshold()) { + break; + } + + key_t key = heap[i].key(); + if (fetch_vector_) { + IndexStorage::MemoryBlock block; + entity_->get_vector_by_key(key, block); + results_[q].emplace_back(key, score, key, block); + } else { + results_[q].emplace_back(key, score); + } + } + } + //! Update context, the context may be shared by different searcher int update_context(IVFEntity::Pointer &new_entity, IndexSearcher::Context::Pointer ¢roid_ctx, @@ -224,6 +269,7 @@ class IVFSearcherContext : public IndexSearcher::Context { IVFEntity::Pointer entity_{}; IndexSearcher::Context::Pointer centroid_searcher_ctx_{}; IndexDocumentHeap result_heap_; + std::vector batch_heaps_{}; std::vector results_{}; std::vector stats_vec_{}; diff --git a/src/core/algorithm/ivf/ivf_streamer.cc b/src/core/algorithm/ivf/ivf_streamer.cc index 2bfab59de..30ec15768 100644 --- a/src/core/algorithm/ivf/ivf_streamer.cc +++ b/src/core/algorithm/ivf/ivf_streamer.cc @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "ivf_streamer.h" +#include +#include #include #include #include "ivf_centroid_index.h" @@ -224,9 +226,10 @@ int IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, ret = entity->transform(query, qmeta, count, &query, &iv_qmeta); ivf_check_with_msg(ret, "Failed to transform querys"); - for (size_t q = 0; q < count; ++q) { - auto ¢roids = centroid_index_ctx->result(q); - auto &context_stats = ctx->mutable_stats(q); + if (count == 1) { + // Single query: original per-query path + auto ¢roids = centroid_index_ctx->result(0); + auto &context_stats = ctx->mutable_stats(0); auto &heap = ctx->mutable_result_heap(); heap.clear(); size_t total_scan_count = 0; @@ -245,15 +248,91 @@ int IVFStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, IndexError::What(ret)); total_scan_count += scan_count; } - heap.sort(); // sort the results + heap.sort(); if (!filter.is_valid()) { - // mapping the local id to key if query without filter ret = entity->retrieve_keys(&heap); ivf_check_error_code(ret); } - entity->normalize(q, &heap); - ctx->topk_to_result(q); + entity->normalize(0, &heap); + ctx->topk_to_result(0); query = static_cast(query) + iv_qmeta.element_size(); + return 0; + } + + // Cluster-first batch path: group queries by cluster, load each cluster once + ctx->reset_batch_heaps(count); + + struct QueryRef { + uint32_t q_idx; + const void *query_ptr; + }; + std::unordered_map> cluster_queries; + std::vector per_query_scan(count, 0); + + const char *qbase = static_cast(query); + for (size_t q = 0; q < count; ++q) { + auto ¢roids = centroid_index_ctx->result(q); + const void *qptr = qbase + q * iv_qmeta.element_size(); + for (size_t i = 0; i < centroids.size(); ++i) { + auto cid = centroids[i].key(); + cluster_queries[cid].push_back({static_cast(q), qptr}); + } + } + + // Sort cluster IDs for stable LRU behavior and sequential I/O access + std::vector sorted_cids; + sorted_cids.reserve(cluster_queries.size()); + for (auto &kv : cluster_queries) sorted_cids.push_back(kv.first); + std::sort(sorted_cids.begin(), sorted_cids.end()); + + // Iterate by cluster in sorted order: block-level multi-query parallel search + for (size_t cid : sorted_cids) { + auto &qrefs = cluster_queries[cid]; + + std::vector batch_items; + std::vector batch_q_indices; + batch_items.reserve(qrefs.size()); + batch_q_indices.reserve(qrefs.size()); + + for (auto &ref : qrefs) { + if (per_query_scan[ref.q_idx] >= ctx->max_scan_count()) { + continue; + } + IVFEntity::BatchQueryItem item; + item.query = ref.query_ptr; + item.heap = &ctx->mutable_batch_heap(ref.q_idx); + item.stats = &ctx->mutable_stats(ref.q_idx); + batch_items.push_back(item); + batch_q_indices.push_back(ref.q_idx); + } + + if (batch_items.empty()) continue; + + uint32_t scan_count = 0; + if (!filter.is_valid()) { + ret = entity->search_batch(cid, batch_items.data(), batch_items.size(), + &scan_count); + } else { + ret = entity->search_batch(cid, filter, batch_items.data(), + batch_items.size(), &scan_count); + } + ivf_check_with_msg(ret, "Failed to search_batch in entity for %s", + IndexError::What(ret)); + + for (auto q_idx : batch_q_indices) { + per_query_scan[q_idx] += scan_count; + } + } + + // Finalize results for each query + for (size_t q = 0; q < count; ++q) { + auto &heap = ctx->mutable_batch_heap(q); + if (!filter.is_valid()) { + ret = entity->retrieve_keys(&heap); + ivf_check_error_code(ret); + } + entity->normalize(q, &heap); + ctx->batch_topk_to_result(static_cast(q)); } return 0; diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 755f673f9..dbf0ca0cf 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -585,6 +585,83 @@ int Index::Search(const VectorData &vector_data, return ret; } +int Index::BatchSearch(const void *queries, uint32_t count, + const BaseIndexQueryParam::Pointer &search_param, + std::vector *results) { + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + if (!queries || count == 0 || !results) { + LOG_ERROR("Invalid batch search arguments"); + return core::IndexError_InvalidArgument; + } + if (is_sparse_) { + LOG_ERROR("BatchSearch is not supported for sparse index"); + return core::IndexError_Unsupported; + } + if (!is_trained_ && this->Train() != 0) { + LOG_ERROR("Failed to train index"); + return core::IndexError_Runtime; + } + + auto &context = acquire_context(); + if (!context) { + LOG_ERROR("Failed to acquire context"); + return core::IndexError_Runtime; + } + + // Prepare context (set topk, filter, etc.) + DenseVector dummy_dense; + dummy_dense.data = queries; + VectorData dummy_vector_data; + dummy_vector_data.vector = dummy_dense; + if (_prepare_for_search(dummy_vector_data, search_param, context) != 0) { + LOG_ERROR("Failed to prepare for batch search"); + context->reset(); + return core::IndexError_Runtime; + } + + // Apply reformer if needed + const void *search_queries = queries; + core::IndexQueryMeta search_meta = input_vector_meta_; + std::string reformer_buffer; + if (reformer_ != nullptr) { + core::IndexQueryMeta new_meta; + if (reformer_->transform(queries, input_vector_meta_, count, + &reformer_buffer, &new_meta) != 0) { + LOG_ERROR("Failed to transform batch queries"); + context->reset(); + return core::IndexError_Runtime; + } + search_queries = reformer_buffer.data(); + search_meta = new_meta; + } + + // Call streamer with batch count + int ret = streamer_->search_impl(search_queries, search_meta, count, context); + if (ret != 0) { + LOG_ERROR("Failed to batch search, ret=%d", ret); + context->reset(); + return ret; + } + + // Extract per-query results + results->resize(count); + for (uint32_t q = 0; q < count; ++q) { + auto &doc_list = context->result(q); + (*results)[q].doc_list_ = doc_list; + if (metric_ && metric_->support_normalize()) { + for (auto &doc : (*results)[q].doc_list_) { + metric_->normalize(doc.mutable_score()); + } + } + } + + context->reset(); + return 0; +} + int Index::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { diff --git a/src/core/interface/indexes/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc index 9df540f35..258faec50 100644 --- a/src/core/interface/indexes/ivf_index.cc +++ b/src/core/interface/indexes/ivf_index.cc @@ -83,20 +83,19 @@ int IVFIndex::Open(const std::string &file_path, break; } case StorageOptions::StorageType::kBufferPool: { - // NOTE: IVF index is dumped via FileDumper (plain binary file), which is - // not compatible with BufferStorage's IndexFormat layout (header/footer - // chain). Until IVF gains a BufferStorage-aware dump path, fall back to - // MMapFileReadStorage so the freshly-dumped file can be reopened. - storage_ = core::IndexFactory::CreateStorage("MMapFileReadStorage"); + // IVF index is dumped via FileDumper (FileDumper container layout). + // BufferReadStorage parses that layout through IndexUnpacker (same as + // MMapFileReadStorage) but serves reads through a VecBufferPool, so the + // freshly-dumped file can be reopened with buffer-pool memory control. + storage_ = core::IndexFactory::CreateStorage("BufferReadStorage"); if (storage_ == nullptr) { - LOG_ERROR( - "Failed to create MMapFileReadStorage (IVF buffer-pool fallback)"); + LOG_ERROR("Failed to create BufferReadStorage (IVF buffer-pool)"); return core::IndexError_Runtime; } int ret = storage_->init(storage_params); if (ret != 0) { LOG_ERROR( - "Failed to init MMapFileReadStorage (IVF buffer-pool fallback), " + "Failed to init BufferReadStorage (IVF buffer-pool), " "path: %s, err: %s", file_path_.c_str(), core::IndexError::What(ret)); return ret; diff --git a/src/core/quantizer/cosine_reformer.cc b/src/core/quantizer/cosine_reformer.cc index ae85a242c..87d2a00ec 100644 --- a/src/core/quantizer/cosine_reformer.cc +++ b/src/core/quantizer/cosine_reformer.cc @@ -164,18 +164,34 @@ class CosineReformer : public IndexReformer { return 0; } - //! Transform queries - int transform(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, - uint32_t /*count*/, std::string * /*out*/, - IndexQueryMeta * /*ometa*/) const override { - return IndexError_Unsupported; + //! Transform queries (batch) + int transform(const void *query, const IndexQueryMeta &qmeta, uint32_t count, + std::string *out, IndexQueryMeta *ometa) const override { + if (count == 0) return 0; + + // Compute output meta once + *ometa = qmeta; + ometa->set_meta(dst_type_, qmeta.dimension() + ExtraDimension(dst_type_)); + size_t out_elem_size = ometa->element_size(); + out->resize(static_cast(count) * out_elem_size); + + const char *qbase = reinterpret_cast(query); + for (uint32_t i = 0; i < count; ++i) { + const void *q = qbase + i * qmeta.element_size(); + std::string single_out; + IndexQueryMeta single_meta; + int ret = + this->CosineReformer::transform(q, qmeta, &single_out, &single_meta); + if (ret != 0) return ret; + ::memcpy(&(*out)[i * out_elem_size], single_out.data(), out_elem_size); + } + return 0; } - //! Convert records - int convert(const void * /*records*/, const IndexQueryMeta & /*rmeta*/, - uint32_t /*count*/, std::string * /*out*/, - IndexQueryMeta * /*ometa*/) const override { - return IndexError_Unsupported; + //! Convert records (batch) + int convert(const void *records, const IndexQueryMeta &rmeta, uint32_t count, + std::string *out, IndexQueryMeta *ometa) const override { + return this->transform(records, rmeta, count, out, ometa); } //! Normalize results diff --git a/src/core/utility/block_heap.h b/src/core/utility/block_heap.h index 09d7ad72c..04f24795c 100644 --- a/src/core/utility/block_heap.h +++ b/src/core/utility/block_heap.h @@ -70,6 +70,12 @@ struct BlockHeap { return cur_ < data_.size(); } + // Peek the next candidate id without popping. + // Caller must ensure has_next() is true. + uint32_t peek() const { + return get_id(data_[cur_].first); + } + // Pop the closest unpopped candidate id (without the check bit). // Caller must ensure has_next() is true. uint32_t pop(); diff --git a/src/core/utility/buffer_read_storage.cc b/src/core/utility/buffer_read_storage.cc new file mode 100644 index 000000000..f110e05d8 --- /dev/null +++ b/src/core/utility/buffer_read_storage.cc @@ -0,0 +1,427 @@ +// 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. +// +// BufferReadStorage is a read-only IndexStorage that mirrors the structure of +// MMapFileReadStorage (it parses the FileDumper container layout through +// IndexUnpacker and exposes segment-based access), but instead of mmap-ing the +// file it reads through a VecBufferPool. This lets IVF / DiskANN(Vamana) +// indexes -- which are dumped via FileDumper -- benefit from the buffer-pool's +// paged cache + LRU eviction + memory-budget control, while keeping the same +// Segment interface that those indexes already consume. +#include +#include +#include +#include +#include +#include +#include +#include "utility_params.h" + +namespace zvec { +namespace core { + +/*! Buffer Read Storage (backed by VecBufferPool) + */ +class BufferReadStorage : public IndexStorage { + public: + /*! Buffer Read Storage Segment + * + * Each segment keeps the owning VecBufferPool / VecBufferPoolHandle alive + * (shared_ptr) so that pages it reads remain valid for the segment's + * lifetime. Reads go through the pool's paged cache: + * - fetch() -> read_range into the caller's buffer + * - read(const void**) -> read_range into a per-segment buffer (stable + * pointer, never pins a page) + * - read(MemoryBlock&) -> single page: zero-copy pin tied to the + * MemoryBlock lifecycle; cross page: owned copy + * - read(SegmentData*) -> read_range into the per-segment buffer + */ + class Segment : public IndexStorage::Segment, + public std::enable_shared_from_this { + public: + //! Index Storage Pointer + typedef std::shared_ptr Pointer; + + //! Constructor + Segment(const std::shared_ptr &pool, + const std::shared_ptr &handle, + 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()), + pool_(pool), + handle_(handle) {} + + //! Constructor (clone) + Segment(const Segment &rhs) + : std::enable_shared_from_this(), + 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_), + pool_(rhs.pool_), + handle_(rhs.handle_) {} + + //! Destructor + ~Segment(void) override {} + + //! Retrieve size of data + size_t data_size(void) const override { + return data_size_; + } + + //! 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 { + if (ailego_unlikely(offset + len > region_size_)) { + if (offset > region_size_) { + offset = region_size_; + } + len = region_size_ - offset; + } + if (len == 0) { + return 0; + } + if (!handle_->read_range(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 pointer via per-segment buffer) + size_t read(size_t offset, const void **data, size_t len) override { + if (ailego_unlikely(offset + len > region_size_)) { + if (offset > region_size_) { + offset = region_size_; + } + len = region_size_ - offset; + } + if (len == 0) { + *data = buffer_.data(); + return 0; + } + buffer_.resize(len); + if (!handle_->read_range(data_offset_ + offset, len, + reinterpret_cast(buffer_.data()))) { + LOG_ERROR( + "BufferReadStorage::Segment::read: read_range failed, " + "abs_offset=%zu, len=%zu", + data_offset_ + offset, len); + *data = nullptr; + return 0; + } + *data = buffer_.data(); + return len; + } + + //! Read data from segment into a MemoryBlock + size_t read(size_t offset, MemoryBlock &data, size_t len) override { + if (ailego_unlikely(offset + len > region_size_)) { + if (offset > region_size_) { + offset = region_size_; + } + len = region_size_ - offset; + } + size_t abs_offset = data_offset_ + 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) { + // Single-page: zero-copy pin whose release is tied to the + // MemoryBlock lifecycle (release_one on destruction). + size_t page_id = 0; + char *raw = handle_->get_single_page(abs_offset, len, page_id); + if (!raw) { + LOG_ERROR( + "BufferReadStorage::Segment::read(MemoryBlock&): single-page " + "acquire failed, abs_offset=%zu, len=%zu", + abs_offset, len); + return 0; + } + data.reset(handle_.get(), page_id, raw); + return len; + } + // Cross-page: copy into a freshly-allocated 4K-aligned buffer that the + // MemoryBlock owns (freed via ailego_free on destruction). + 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( + "BufferReadStorage::Segment::read(MemoryBlock&): cross-page alloc " + "failed, abs_offset=%zu, len=%zu", + abs_offset, len); + return 0; + } + if (!handle_->read_range(abs_offset, len, tmp)) { + 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; + } + + //! Read scattered data from segment (stable pointers via per-segment buf) + bool read(SegmentData *iovec, size_t count) override { + size_t total = 0u; + for (auto *it = iovec, *end = iovec + count; it != end; ++it) { + ailego_false_if_false(it->offset + it->length <= region_size_); + total += it->length; + } + ailego_false_if_false(total != 0); + + buffer_.resize(total); + uint8_t *buf = buffer_.data(); + for (auto *it = iovec, *end = iovec + count; it != end; ++it) { + ailego_false_if_false( + handle_->read_range(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 { + return; + } + + //! Clone the segment + IndexStorage::Segment::Pointer clone(void) override { + return std::make_shared(*this); + } + + void prefetch(size_t offset, size_t len) override { + if (offset + len > region_size_) { + len = (offset > region_size_) ? 0 : region_size_ - offset; + } + if (len == 0) return; + handle_->prefetch_range(data_offset_ + offset, len); + } + + //! Free bytes in the shared buffer pool. Used by the caller to decide + //! whether a whole cluster fits before issuing prefetch. + size_t prefetch_budget(void) const override { + return ailego::MemoryLimitPool::get_instance().available(); + } + + //! No stable base pointer: data lives in an evictable paged cache. + const uint8_t *base_data(void) const override { + return nullptr; + } + + private: + 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::vector buffer_{}; + std::shared_ptr pool_{nullptr}; + std::shared_ptr handle_{nullptr}; + }; + + //! Destructor + ~BufferReadStorage(void) override {} + + //! 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_ENABLE_DIRECT_IO, &enable_direct_io_); + return 0; + } + + int flush(void) override { + return 0; + } + + int append(const std::string &, size_t) override { + return IndexError_NotImplemented; + } + + void refresh(uint64_t) override { + return; + } + + 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 { + // Read-only buffer pool over the freshly-dumped FileDumper container. + buffer_pool_ = std::make_shared( + path, /*writable=*/false, /*enable_direct_io=*/enable_direct_io_); + if (!buffer_pool_) { + LOG_ERROR("Failed to create VecBufferPool, path: %s", path.c_str()); + return IndexError_NoMemory; + } + handle_ = std::make_shared( + buffer_pool_->get_handle()); + + size_t file_size = buffer_pool_->file_size(); + index_offset_ = (header_offset_ >= 0 ? 0 : file_size) + header_offset_; + size_t end_offset = (footer_offset_ > 0 ? 0 : file_size) + footer_offset_; + size_t size = end_offset > index_offset_ ? end_offset - index_offset_ : 0; + + // read_data for IndexUnpacker: provide a stable pointer by copying the + // requested range into a reused scratch buffer via get_meta (direct + // pread, valid before buffer_pool_->init()). + auto read_data = [this, end_offset](size_t offset, const void **data, + size_t len) -> size_t { + size_t off = offset + index_offset_; + if (off + len > end_offset) { + if (off > end_offset) { + off = end_offset; + } + len = end_offset - off; + } + scratch_.resize(len); + *data = scratch_.data(); + if (len == 0) { + return 0; + } + if (handle_->get_meta(off, len, + reinterpret_cast(scratch_.data())) != 0) { + return 0; + } + return len; + }; + + IndexUnpacker unpacker; + if (!unpacker.unpack(read_data, size, checksum_validation_)) { + LOG_ERROR("Failed to unpack file: %s", path.c_str()); + return IndexError_UnpackIndex; + } + segments_ = std::move(*unpacker.mutable_segments()); + magic_ = unpacker.magic(); + + // Allocate the page table now that the layout is known. + int ret = buffer_pool_->init(); + if (ret != 0) { + LOG_ERROR("Failed to init VecBufferPool, path: %s", path.c_str()); + return IndexError_Runtime; + } + buffer_pool_->warmup(); + return 0; + } + + int close(void) override { + segments_.clear(); + handle_ = nullptr; + buffer_pool_ = nullptr; + return 0; + } + + //! Retrieve a segment by id + IndexStorage::Segment::Pointer get(const std::string &id, int) override { + if (!buffer_pool_ || !handle_) { + return IndexStorage::Segment::Pointer(); + } + auto it = segments_.find(id); + if (it == segments_.end()) { + return IndexStorage::Segment::Pointer(); + } + return std::make_shared( + buffer_pool_, handle_, index_offset_, it->second); + } + + std::map get_all( + void) const override { + std::map result; + if (buffer_pool_ && handle_) { + for (const auto &it : segments_) { + result.emplace(it.first, + std::make_shared( + buffer_pool_, handle_, 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; + } + + private: + bool checksum_validation_{false}; + bool enable_direct_io_{true}; + int64_t header_offset_{0}; + int64_t footer_offset_{0}; + size_t index_offset_{0}; + uint32_t magic_{0}; + std::vector scratch_{}; + std::map segments_{}; + std::shared_ptr buffer_pool_{nullptr}; + std::shared_ptr handle_{nullptr}; +}; + +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..ed4b43d38 100644 --- a/src/core/utility/buffer_storage.cc +++ b/src/core/utility/buffer_storage.cc @@ -449,6 +449,12 @@ class BufferStorage : public IndexStorage { return shared_from_this(); } + size_t abs_data_offset(void) const override { + return segment_info_->segment_header_start_offset + + segment_info_->segment_header->content_offset + + segment_info_->segment.meta()->data_index; + } + protected: friend BufferStorage; // Pointer into BufferStorage::segments_ (unordered_map mapped value). @@ -481,6 +487,8 @@ class BufferStorage : public IndexStorage { if (val != 0) { segment_meta_capacity_ = val; } + // O_DIRECT is always on for memory controllability. + params.get(BUFFER_STORAGE_ENABLE_DIRECT_IO, &enable_direct_io_); return 0; } @@ -506,10 +514,11 @@ 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). + // O_DIRECT is used for both build and search to keep memory + // controllable (no uncontrolled page-cache growth). buffer_pool_ = std::make_shared( - path, /*writable=*/create_if_missing); + path, /*writable=*/create_if_missing, + /*enable_direct_io=*/enable_direct_io_); buffer_pool_handle_ = std::make_shared( buffer_pool_->get_handle()); int ret = ParseToMapping(); @@ -522,6 +531,7 @@ class BufferStorage : public IndexStorage { this->close_index(); return ret; } + buffer_pool_->warmup(); LOG_INFO( "BufferStorage opened: file=%s, writable=%d, max_segment_size=%" PRIu64 ", segment_count=%zu", @@ -804,6 +814,10 @@ class BufferStorage : public IndexStorage { return chain_headers_.front()->magic; } + ailego::VecBufferPool *vec_buffer_pool(void) const override { + return buffer_pool_.get(); + } + protected: //! Initialize index version segment (writes content into an IndexMapping). //! Only intended to be called from init_index() while `mapping` is still @@ -1530,6 +1544,10 @@ class BufferStorage : public IndexStorage { // init_index(). uint32_t segment_meta_capacity_{4096u}; + // When true, the page-data fd is opened with O_DIRECT (metadata fd stays + // buffered). Defaults to true for memory controllability. + bool enable_direct_io_{true}; + // Per-header-chain file offsets used by flush_index() and append_segment(). struct MetaChain { uint64_t header_start_offset; diff --git a/src/core/utility/linear_pool.h b/src/core/utility/linear_pool.h index 20c3bc8c4..d67314420 100644 --- a/src/core/utility/linear_pool.h +++ b/src/core/utility/linear_pool.h @@ -184,6 +184,11 @@ struct LinearPool { bool has_next() const { return cur_ < size_ && cur_ < ef_; } + // Peek the next candidate id without popping. + // Caller must ensure has_next() is true. + int peek() const { + return get_id(data_[cur_].id); + } int id(int i) const { return get_id(data_[i].id); } diff --git a/src/core/utility/utility_params.h b/src/core/utility/utility_params.h index c57e6e980..59d21788f 100644 --- a/src/core/utility/utility_params.h +++ b/src/core/utility/utility_params.h @@ -60,6 +60,16 @@ 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_ENABLE_DIRECT_IO = + "proxima.buffer.read_storage.enable_direct_io"; + //! MMapFileStorage static const std::string MMAPFILE_STORAGE_MEMORY_LOCKED = "proxima.mmap_file.storage.memory_locked"; @@ -72,6 +82,10 @@ static const std::string MMAPFILE_STORAGE_FORCE_FLUSH = static const std::string MMAPFILE_STORAGE_SEGMENT_META_CAPACITY = "proxima.mmap_file.storage.segment_meta_capacity"; +//! BufferStorage +static const std::string BUFFER_STORAGE_ENABLE_DIRECT_IO = + "proxima.buffer.storage.enable_direct_io"; + //! MipsConverter static const std::string MIPS_CONVERTER_M_VALUE = "proxima.mips.converter.m_value"; diff --git a/src/include/zvec/ailego/buffer/block_eviction_queue.h b/src/include/zvec/ailego/buffer/block_eviction_queue.h index fa5aff2a5..ee78a72d3 100644 --- a/src/include/zvec/ailego/buffer/block_eviction_queue.h +++ b/src/include/zvec/ailego/buffer/block_eviction_queue.h @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include #include #include @@ -30,8 +32,10 @@ #include #include #include +#include #include #include +#include #include #include "concurrentqueue.h" @@ -52,7 +56,11 @@ class 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; + //! Attempt to evict the block identified by owner_key. + //! Returns true if the block was actually evicted (its memory was released), + //! false if it was spared (CLOCK second chance / still pinned) so callers + //! can tell real reclamation from a no-op and make progress accordingly. + virtual bool evict_block(eviction_key_t owner_key) = 0; }; class BlockEvictionQueue { @@ -81,13 +89,6 @@ class BlockEvictionQueue { 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(); - } - void set_valid(EvictableBlockOwner *owner) { std::unique_lock lock(valid_owners_mutex_); valid_owners_.insert(owner); @@ -98,13 +99,12 @@ 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); + private: BlockEvictionQueue() { init(); @@ -141,12 +141,116 @@ class MemoryLimitPool { bool is_full(); + //! Currently available (free) bytes in the pool. Lock-free read-only + //! estimate: used_size_ is atomic and pool_size_ is fixed after init(). + //! Returns 0 when the pool is at or over capacity. Callers use this to gate + //! whole-cluster prefetch under memory pressure. + size_t available() const { + size_t used = used_size_.load(); + return (used >= pool_size_) ? 0 : (pool_size_ - used); + } + + 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); + } + + //! Total capacity in bytes (fixed after init()). + size_t capacity() const { + return pool_size_; + } + + //! Snapshot of pool-level counters for monitoring / export. + struct PoolStats { + size_t pool_size{0}; + size_t used{0}; + size_t free_buffers{0}; // buffers cached across all shards + uint64_t alloc_from_freelist{0}; // acquisitions served from a shard + uint64_t alloc_from_slab{0}; // acquisitions that carved a new buffer + uint64_t bg_evict_rounds{0}; // background reclaim passes + uint64_t bg_evicted_buffers{0}; // buffers reclaimed by background thread + uint64_t high_watermark_hits{0}; // foreground acquire hit the capacity cap + }; + PoolStats stats() const; + void log_stats() const; + private: MemoryLimitPool() = default; + ~MemoryLimitPool(); + + void drain_free_list(); + char *carve_from_slab_locked(size_t buffer_size); + void free_all_slabs_locked(); + size_t pick_shard(); + + // Background evictor: proactively reclaims buffers down to the low + // watermark so foreground acquire_buffer() rarely pays eviction/flush cost + // inline, smoothing tail latency under memory pressure. + void start_background_evictor(); + void stop_background_evictor(); + void background_evict_loop(); + // Proactive reclaim keeps a small free margin so the foreground path finds + // ready buffers without paying inline eviction cost. These used to be 90% / + // 75% of pool_size, but a percentage margin makes the background evictor + // discard live pages whenever used exceeds 75% -- even when the entire + // working set would fit in the pool. For read-mostly page caches that is + // self-inflicted thrashing: a pool larger than the index still perpetually + // evicts and re-reads its hottest pages (e.g. a 1.4 GB pool caps at ~1.05 GB + // and thrashes a 1.17 GB index). Reserve a small absolute margin instead so + // the pool fills to near capacity before anything is evicted. + size_t reserve_margin() const { + size_t m = pool_size_ / 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 >= pool_size_) m = pool_size_ / 8; // tiny pools: fall back + return m; + } + size_t high_watermark() const { + return pool_size_ - reserve_margin(); + } + size_t low_watermark() const { + return pool_size_ - reserve_margin() * 2; + } private: + // Sharded free-list: spreads the hot alloc/free path across many locks so + // foreground threads and the background evictor rarely contend. Each shard + // is cache-line aligned to avoid false sharing. + static constexpr size_t kNumFreeShards = 64; + struct alignas(64) FreeShard { + std::mutex mutex; + char *head{nullptr}; + std::atomic count{0}; + }; + size_t pool_size_{0}; std::atomic used_size_{0}; + + FreeShard free_shards_[kNumFreeShards]; + std::atomic shard_seq_{0}; + + // Slab allocator (cold path) guarded by its own mutex, decoupled from the + // free-list shards. + std::mutex slab_mutex_; + std::vector slabs_; + char *slab_cursor_{nullptr}; + size_t slab_remaining_{0}; + + // Observability counters (relaxed atomics; statistics only). + std::atomic alloc_from_freelist_{0}; + std::atomic alloc_from_slab_{0}; + std::atomic bg_evict_rounds_{0}; + std::atomic bg_evicted_buffers_{0}; + std::atomic high_watermark_hits_{0}; + + 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..b6bed6f18 100644 --- a/src/include/zvec/ailego/buffer/external_cache.h +++ b/src/include/zvec/ailego/buffer/external_cache.h @@ -133,16 +133,16 @@ class ExternalCache : public EvictableBlockOwner { return iter->second.generation.load(std::memory_order_relaxed) != version; } - 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; + return false; } Entry &entry = iter->second; @@ -152,7 +152,9 @@ class ExternalCache : public EvictableBlockOwner { MemoryLimitPool::get_instance().release_external(entry.size); entry.size = 0; loader_.clear(entry.payload); + return true; } + return false; } private: diff --git a/src/include/zvec/ailego/buffer/vector_page_table.h b/src/include/zvec/ailego/buffer/vector_page_table.h index 02d19bbc7..876771b17 100644 --- a/src/include/zvec/ailego/buffer/vector_page_table.h +++ b/src/include/zvec/ailego/buffer/vector_page_table.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,9 @@ #include #include #include +#if defined(__linux) || defined(__linux__) +#include +#endif #include "block_eviction_queue.h" #include "concurrentqueue.h" @@ -50,6 +54,13 @@ class VectorPageTable : public EvictableBlockOwner { std::atomic ref_count; std::atomic in_evict_queue; std::atomic is_dirty; + // CLOCK second-chance bit: set on every cache hit, cleared when the + // evictor spares the page. Turns the FIFO eviction queue into an + // access-aware (approximate LRU) policy without a global LRU list. + std::atomic referenced; + uint8_t evict_priority{0}; + bool ever_loaded{ + false}; // true once the page has been loaded at least once char *buffer; size_t file_offset; }; @@ -94,7 +105,14 @@ class VectorPageTable : public EvictableBlockOwner { 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 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 = priority; + e.in_evict_queue.store(false, std::memory_order_relaxed); + } char *set_block_acquired(block_id_t block_id, char *buffer, size_t file_offset); @@ -114,6 +132,19 @@ class VectorPageTable : public EvictableBlockOwner { return entry_at(block_id).is_dirty.load(std::memory_order_relaxed); } + //! Get the raw buffer pointer for a loaded page (nullptr if not loaded). + //! Used by batched flush to memcpy page contents into a coalescing buffer. + char *get_block_buffer(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + return entry_at(block_id).buffer; + } + + //! 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)); + entry_at(block_id).is_dirty.store(false, std::memory_order_relaxed); + } + //! Flush a single dirty block without evicting it. Caller guarantees the //! block is currently loaded (buffer != nullptr). int flush_block(block_id_t block_id) { @@ -140,6 +171,22 @@ class VectorPageTable : public EvictableBlockOwner { return entry_num_.load(std::memory_order_acquire); } + //! Cache observability counters (monotonic, relaxed atomics). + struct Stats { + uint64_t hit{0}; // acquire_block served from memory + 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 + }; + Stats stats() const { + Stats s; + s.hit = hit_count_.load(std::memory_order_relaxed); + s.evict = evict_count_.load(std::memory_order_relaxed); + s.second_chance = second_chance_count_.load(std::memory_order_relaxed); + s.dirty_flush = dirty_flush_count_.load(std::memory_order_relaxed); + return s; + } + 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; @@ -151,6 +198,21 @@ class VectorPageTable : public EvictableBlockOwner { return !e.in_evict_queue.load(std::memory_order_relaxed); } + //! Check if a page is loaded (has a non-null buffer). + //! Used by try_acquire_buffer to avoid ref_count leaks on unloaded pages. + 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 != nullptr; + } + + //! Check if a page has ever been loaded (so pread is needed on reload + //! after eviction). A page that was never loaded can be zero-filled + //! if it lies beyond the initial file size (created by extend_file). + bool is_ever_loaded(block_id_t block_id) const { + assert(block_id < entry_num_.load(std::memory_order_acquire)); + return entry_at(block_id).ever_loaded; + } + private: // Segmented page table: entries are split across fixed-size segments so // that extend() can grow the table without moving existing entries. @@ -192,6 +254,13 @@ class VectorPageTable : public EvictableBlockOwner { } FlushCallback flush_callback_{}; + + // Observability counters. Relaxed ordering: these are statistics, never + // used to make correctness decisions. + std::atomic hit_count_{0}; + std::atomic evict_count_{0}; + std::atomic second_chance_count_{0}; + std::atomic dirty_flush_count_{0}; }; class VecBufferPoolHandle; @@ -202,8 +271,12 @@ class VecBufferPool { static constexpr size_t kMutexBucketCount = 64UL * 1024UL; - VecBufferPool(const std::string &filename, bool writable = false); + VecBufferPool(const std::string &filename, bool writable = false, + bool enable_direct_io = false); ~VecBufferPool() { + // Emit a one-line cache summary (hit rate / evictions) before teardown + // so operators can reason about buffer-pool efficiency per file. + log_stats(); // 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(); @@ -211,15 +284,49 @@ class VecBufferPool { assert(page_table_.is_released(i)); page_table_.evict_block(i); } +#if defined(__linux) || defined(__linux__) + if (aio_enabled_ && aio_ctx_) { + LibAioLoader::Instance().io_destroy(aio_ctx_); + } +#endif #if defined(_MSC_VER) _close(fd_); + _close(meta_fd_); #else close(fd_); + close(meta_fd_); #endif } 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}; + double hit_rate() const { + uint64_t total = hit + miss; + return total ? static_cast(hit) / static_cast(total) + : 0.0; + } + }; + 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.miss = miss_count_.load(std::memory_order_relaxed); + return s; + } + + //! 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); @@ -253,11 +360,69 @@ class VecBufferPool { 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); + + void prefetch_pages_aio(block_id_t first_page, size_t page_count); + + //! Submit AIO for scattered pages without waiting for completion. + //! Results are stored in thread-local state; call harvest_aio() later. + void submit_aio_async(const block_id_t *page_ids, size_t count); + + //! Harvest previously submitted async AIO results (non-blocking). + //! Installs completed pages into page_table. Safe to call when no pending. + void harvest_aio(); + + //! Blocking counterpart of harvest_aio(): waits (io_getevents with a NULL + //! timeout) until every in-flight request submitted via submit_aio_async() + //! has completed and been installed into page_table. Lets a search hop + //! issue one concurrent read burst and block once for the whole batch, + //! instead of rationing per-neighbour synchronous preads. Safe to call + //! when nothing is pending. + void wait_aio(); + + bool aio_enabled() const { +#if defined(__linux) || defined(__linux__) + return aio_enabled_; +#else + return false; +#endif + } + + //! Try to acquire a page buffer WITHOUT triggering disk I/O. + //! Returns the buffer pointer if the page is already in memory, + //! or nullptr if the page would need a pread (cache miss). + //! Avoids touching ref_count for unloaded pages to prevent leaks. + char *try_acquire_buffer(block_id_t page_id) { + assert(page_id < page_table_.entry_num()); + // Quick check: if buffer is null, page not loaded - skip without + // incrementing ref_count (acquire_block would leak ref_count on + // pages with ref_count>=0 but buffer==nullptr). + if (!page_table_.is_loaded(page_id)) return nullptr; + return page_table_.acquire_block(page_id); + } + private: - int fd_; + friend class VecBufferPoolHandle; + int fd_; // page-data channel: may carry O_DIRECT + 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}; + // Cache miss counter: incremented once per page fetched from disk on the + // cold acquire path (pread / zero-fill). Combined with page_table_ hits it + // yields the pool hit rate. + std::atomic miss_count_{0}; +#if defined(__linux) || defined(__linux__) + io_context_t aio_ctx_{nullptr}; + bool aio_enabled_{false}; +#endif public: VectorPageTable page_table_; @@ -279,6 +444,8 @@ class VecBufferPoolHandle { bool read_range(size_t file_offset, size_t len, char *out); + void prefetch_range(size_t file_offset, size_t len); + int get_meta(size_t offset, size_t length, char *buffer); int write_range(size_t file_offset, size_t len, const char *src); diff --git a/src/include/zvec/ailego/io/libaio_def.h b/src/include/zvec/ailego/io/libaio_def.h new file mode 100644 index 000000000..fa6bbb34b --- /dev/null +++ b/src/include/zvec/ailego/io/libaio_def.h @@ -0,0 +1,155 @@ +#pragma once + +#include +#include + +#if defined(__linux) || defined(__linux__) + +// If system is already included, skip our definitions. +// Conversely, define __LIBAIO_H so that a later #include becomes a +// no-op (guarded by #ifndef __LIBAIO_H in the system header). +#ifdef __LIBAIO_H +// System header already provided all AIO types; nothing to do. +#else +#define __LIBAIO_H + +struct sockaddr; +struct iovec; + +typedef struct io_context *io_context_t; + +typedef enum io_iocb_cmd { + IO_CMD_PREAD = 0, + IO_CMD_PWRITE = 1, + IO_CMD_FSYNC = 2, + IO_CMD_FDSYNC = 3, + IO_CMD_POLL = 5, + IO_CMD_NOOP = 6, + IO_CMD_PREADV = 7, + IO_CMD_PWRITEV = 8, +} io_iocb_cmd_t; + +#if defined(__i386__) || (defined(__x86_64__) && defined(__ILP32__)) || \ + (defined(__arm__) && !defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__LITTLE_ENDIAN__)) || defined(__bfin__) || \ + (defined(__MIPSEL__) && !defined(__mips64)) || defined(__cris__) || \ + defined(__loongarch32) || (defined(__riscv) && __riscv_xlen == 32) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDptr(x, y) \ + x; \ + unsigned y +#define AIO_PADDEDul(x, y) \ + unsigned long x; \ + unsigned y + +#elif defined(__ia64__) || defined(__x86_64__) || defined(__alpha__) || \ + (defined(__mips64) && defined(__MIPSEL__)) || \ + (defined(__aarch64__) && defined(__AARCH64EL__)) || \ + defined(__loongarch64) || (defined(__riscv) && __riscv_xlen == 64) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) x, y +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +#elif defined(__powerpc64__) || defined(__s390x__) || \ + (defined(__hppa__) && defined(__arch64__)) || \ + (defined(__sparc__) && defined(__arch64__)) || \ + (defined(__mips64) && defined(__MIPSEB__)) || \ + (defined(__aarch64__) && defined(__AARCH64EB__)) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 8) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) x +#define AIO_PADDEDul(x, y) unsigned long x + +#elif defined(__PPC__) || defined(__s390__) || \ + (defined(__arm__) && defined(__ARMEB__)) || \ + (defined(__sh__) && defined(__BIG_ENDIAN__)) || defined(__sparc__) || \ + defined(__MIPSEB__) || defined(__m68k__) || defined(__hppa__) || \ + defined(__frv__) || defined(__avr32__) || \ + (defined(__GNUC__) && defined(__BYTE_ORDER__) && \ + __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 4) +#define AIO_PADDED(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDptr(x, y) \ + unsigned y; \ + x +#define AIO_PADDEDul(x, y) \ + unsigned y; \ + unsigned long x + +#else +#error endianness? +#endif + +struct io_iocb_poll { + AIO_PADDED(int events, __pad1); +}; + +struct io_iocb_sockaddr { + AIO_PADDEDptr(struct sockaddr *addr, __pad1); + AIO_PADDEDul(len, __pad2); +}; + +struct io_iocb_common { + AIO_PADDEDptr(void *buf, __pad1); + AIO_PADDEDul(nbytes, __pad2); + long long offset; + long long __pad3; + unsigned flags; + unsigned resfd; +}; + +struct io_iocb_vector { + AIO_PADDEDptr(const struct iovec *vec, __pad1); + AIO_PADDEDul(nr, __pad2); + long long offset; +}; + +struct iocb { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDED(unsigned key, aio_rw_flags); + short aio_lio_opcode; + short aio_reqprio; + int aio_fildes; + union { + struct io_iocb_common c; + struct io_iocb_vector v; + struct io_iocb_poll poll; + struct io_iocb_sockaddr saddr; + } u; +}; + +struct io_event { + AIO_PADDEDptr(void *data, __pad1); + AIO_PADDEDptr(struct iocb *obj, __pad2); + AIO_PADDEDul(res, __pad3); + AIO_PADDEDul(res2, __pad4); +}; + +#undef AIO_PADDED +#undef AIO_PADDEDptr +#undef AIO_PADDEDul + +static inline void io_prep_pread(struct iocb *iocb, int fd, void *buf, + size_t count, long long offset) { + memset(iocb, 0, sizeof(*iocb)); + iocb->aio_fildes = fd; + iocb->aio_lio_opcode = IO_CMD_PREAD; + iocb->aio_reqprio = 0; + iocb->u.c.buf = buf; + iocb->u.c.nbytes = count; + iocb->u.c.offset = offset; +} + +#endif // !__LIBAIO_H (our definitions block) + +#endif // __linux diff --git a/src/include/zvec/ailego/io/libaio_loader.h b/src/include/zvec/ailego/io/libaio_loader.h new file mode 100644 index 000000000..fa82b9e5d --- /dev/null +++ b/src/include/zvec/ailego/io/libaio_loader.h @@ -0,0 +1,94 @@ +#pragma once + +#if defined(__linux) || defined(__linux__) + +#include +#include +#include +#include "libaio_def.h" + +typedef int (*aio_setup_fn)(int maxevents, io_context_t *ctxp); +typedef int (*aio_destroy_fn)(io_context_t ctx); +typedef int (*aio_submit_fn)(io_context_t ctx, long nr, struct iocb *ios[]); +typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr, + struct io_event *events, + struct timespec *timeout); + +class LibAioLoader { + public: + static LibAioLoader &Instance() { + static LibAioLoader instance; + return instance; + } + + bool Load() { + if (available_.load(std::memory_order_acquire)) { + return true; + } + std::call_once(once_, [this] { this->TryLoad(); }); + return available_.load(std::memory_order_relaxed); + } + + bool IsAvailable() const { + return available_.load(std::memory_order_acquire); + } + + aio_setup_fn io_setup; + aio_destroy_fn io_destroy; + aio_submit_fn io_submit; + aio_getevents_fn io_getevents; + + private: + LibAioLoader() + : io_setup(nullptr), + io_destroy(nullptr), + io_submit(nullptr), + io_getevents(nullptr) {} + + ~LibAioLoader() = default; + + LibAioLoader(const LibAioLoader &) = delete; + LibAioLoader &operator=(const LibAioLoader &) = delete; + + void TryLoad() { + static constexpr const char *kSonames[] = { + "libaio.so.1", + "libaio.so.1t64", + }; + + for (const char *soname : kSonames) { + void *h = dlopen(soname, RTLD_LAZY); + if (h == nullptr) { + continue; + } + + io_setup = reinterpret_cast(dlsym(h, "io_setup")); + io_destroy = reinterpret_cast(dlsym(h, "io_destroy")); + io_submit = reinterpret_cast(dlsym(h, "io_submit")); + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents")); + if (io_getevents == nullptr) { + io_getevents = + reinterpret_cast(dlsym(h, "io_getevents_time64")); + } + + if (io_setup && io_destroy && io_submit && io_getevents) { + handle_ = h; + available_.store(true, std::memory_order_release); + return; + } + + dlclose(h); + io_setup = nullptr; + io_destroy = nullptr; + io_submit = nullptr; + io_getevents = nullptr; + } + } + + std::once_flag once_; + std::atomic available_{false}; + void *handle_{nullptr}; +}; + +#endif 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..083b3d9a5 100644 --- a/src/include/zvec/core/framework/index_storage.h +++ b/src/include/zvec/core/framework/index_storage.h @@ -332,6 +332,25 @@ class IndexStorage : public IndexModule { 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) { + (void)offset; + (void)len; + } + + //! Bytes the storage backend can currently accept for prefetch without + //! evicting still-useful data. Pooled/evictable backends (e.g. + //! BufferReadStorage) return the free space of their buffer pool; backends + //! without a bounded pool (e.g. mmap) return SIZE_MAX to signal "always + //! fits". Callers use it to gate whole-cluster prefetch under memory + //! pressure. + virtual size_t prefetch_budget(void) const { + return static_cast(-1); + } }; //! Destructor @@ -399,6 +418,10 @@ class IndexStorage : public IndexModule { virtual std::string file_path(void) const { return ""; } + + virtual ailego::VecBufferPool *vec_buffer_pool(void) const { + return nullptr; + } }; } // namespace core diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 118d06d44..b4704c5aa 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -138,6 +138,11 @@ class Index { const BaseIndexQueryParam::Pointer &search_param, SearchResult *result); + //! Batch search: N queries in one call, cluster-first for IVF + virtual int BatchSearch(const void *queries, uint32_t count, + const BaseIndexQueryParam::Pointer &search_param, + std::vector *results); + virtual int AddWithSource(const VectorData &vector, uint32_t doc_id, const core::VectorSource &src); virtual int SearchWithSource(const VectorData &query, 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..479e5a3be --- /dev/null +++ b/tests/ailego/buffer/vector_page_table_test.cc @@ -0,0 +1,248 @@ +// 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. Observability counters (hit / miss / evict / second_chance / stats) + +#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) { + MemoryLimitPool::get_instance().init(capacity_pages * kVectorPageSize); + } + 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_; +}; + +} // namespace + +// --------------------------------------------------------------------------- +// 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; + InitPool(/*capacity_pages=*/16); // 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 +} + +// --------------------------------------------------------------------------- +// 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; + InitPool(/*capacity_pages=*/32); + 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); + }; + + // Interleave a hot set (0..7) with a rolling cold scan to create pressure. + for (int round = 0; round < 20; ++round) { + for (size_t h = 0; h < 8; ++h) read_page(h); // keep hot set referenced + for (size_t c = 8; c < num_pages; ++c) read_page(c); // cold churn + } + + 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; + InitPool(cap_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 = mp.capacity() / 4 * 3; // 75% + // Poll up to ~2s for the background thread to reclaim down to the low mark. + for (int i = 0; i < 200 && mp.used() > low + kVectorPageSize; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_LE(mp.used(), low + kVectorPageSize); + EXPECT_GT(mp.stats().bg_evicted_buffers, 0u); +} + +// --------------------------------------------------------------------------- +// 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; + InitPool(/*capacity_pages=*/48); + 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); + + // 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); +}