diff --git a/src/ailego/buffer/vector_page_table.cc b/src/ailego/buffer/vector_page_table.cc index 2f6314ed6..6364aa346 100644 --- a/src/ailego/buffer/vector_page_table.cc +++ b/src/ailego/buffer/vector_page_table.cc @@ -993,6 +993,18 @@ VecBufferPool::VecBufferPool(const std::string &filename, bool writable) { #endif throw std::runtime_error("Failed to stat file: " + filename); } +#if !defined(_MSC_VER) + // Atomic replacement between the two opens must not mix metadata from one + // file with pages from another. Windows CRT opens prevent deletion while + // either descriptor is live; POSIX needs an explicit identity check. + struct stat meta_st; + if (fstat(meta_fd_, &meta_st) < 0 || st.st_dev != meta_st.st_dev || + st.st_ino != meta_st.st_ino) { + ::close(fd_); + ::close(meta_fd_); + throw std::runtime_error("Backing file changed while opening: " + filename); + } +#endif file_size_ = st.st_size; initial_file_size_ = file_size_; #if defined(__linux__) && !defined(__ANDROID__) diff --git a/src/ailego/io/libaio_loader.h b/src/ailego/io/libaio_loader.h index 8bdd415f8..60874d137 100644 --- a/src/ailego/io/libaio_loader.h +++ b/src/ailego/io/libaio_loader.h @@ -54,8 +54,12 @@ typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr, class LibAioLoader { public: static LibAioLoader &Instance() { - static LibAioLoader instance; - return instance; + // Global thread pools may join workers during static destruction, after + // this loader would be destroyed. Their thread-local AIO contexts still + // need io_destroy(), so retain the loader and library until process exit. + // Individual AIO contexts must still be destroyed by their owners. + static LibAioLoader *const instance = new LibAioLoader(); + return *instance; } // Load (or confirm already loaded) libaio. Returns true on success. diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index 5bc6f857f..5b9c0f1ec 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -85,7 +85,8 @@ int DiskAnnContext::resize_fetch_sector_buffer( } int DiskAnnContext::init(ContextType type, uint32_t /*graph_degree*/, - uint32_t pq_chunk_num, uint32_t element_size) { + uint32_t pq_chunk_num, uint32_t element_size, + bool setup_io_context) { if (!entity_ || element_size == 0) { LOG_ERROR("Invalid DiskAnn context parameters"); return IndexError_InvalidArgument; @@ -137,10 +138,12 @@ int DiskAnnContext::init(ContextType type, uint32_t /*graph_degree*/, return IndexError_NoMemory; } - ret = setup_io_ctx(io_ctx_); - if (ret != 0) { - LOG_ERROR("setup io ctx error, ret=%d", ret); - return ret; + if (setup_io_context) { + ret = setup_io_ctx(io_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error, ret=%d", ret); + return ret; + } } break; @@ -150,10 +153,12 @@ int DiskAnnContext::init(ContextType type, uint32_t /*graph_degree*/, return ret; } - ret = setup_io_ctx(io_ctx_); - if (ret != 0) { - LOG_ERROR("setup fetch io ctx error, ret=%d", ret); - return ret; + if (setup_io_context) { + ret = setup_io_ctx(io_ctx_); + if (ret != 0) { + LOG_ERROR("setup fetch io ctx error, ret=%d", ret); + return ret; + } } break; diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index 2d2107da8..89f21d225 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -67,7 +67,7 @@ class DiskAnnContext : public IndexContext, public: //! Init int init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, - uint32_t element_size); + uint32_t element_size, bool setup_io_context = true); //! Update context, the context may be shared by different searcher/streamer int update_context(ContextType type, const IndexMeta &meta, diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 7b881b063..b65fa4e59 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -26,9 +26,11 @@ #include #include #include +#include #include #include #if defined(_WIN32) || defined(_WIN64) +#include #include #endif #if defined(__APPLE__) || defined(__MACH__) @@ -762,6 +764,331 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return ret; } +#endif // POSIX reader declarations + +BufferPoolAlignedFileReader::BufferPoolAlignedFileReader( + std::shared_ptr pool) + : pool_(std::move(pool)) {} + +BufferPoolAlignedFileReader::~BufferPoolAlignedFileReader() = default; + +void BufferPoolAlignedFileReader::open(const std::string &fname) { + (void)open_from_pool(fname); +} + +int BufferPoolAlignedFileReader::open_from_pool(const std::string &fname) { + opened_ = false; + bypass_reader_.close(); + if (!pool_) { + return IndexError_InvalidArgument; + } +#if defined(_WIN32) || defined(_WIN64) + const auto handle = + reinterpret_cast(_get_osfhandle(pool_->file_descriptor())); + const int ret = bypass_reader_.open_from_handle(fname, handle); +#else + const int ret = + bypass_reader_.open_from_handle(fname, pool_->file_descriptor()); +#endif + opened_ = ret == 0; + return ret; +} + +void BufferPoolAlignedFileReader::close() { + opened_ = false; + bypass_reader_.close(); + pool_.reset(); +} + +int BufferPoolAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx, bool /*async*/) { + if (!pool_ || !opened_) { + LOG_ERROR("BufferPoolAlignedFileReader: backing file is not open"); + return IndexError_Runtime; + } + if (read_reqs.empty()) { + return 0; + } + + try { + struct UniquePage { + ailego::block_id_t page_id; + char *cached_page{nullptr}; + bool bypass_candidate{false}; + }; + struct PageOccurrence { + size_t unique_index; + char *destination; + size_t offset_in_page; + size_t length; + size_t canonical_index; + }; + + size_t total_pages = 0; + for (const AlignedRead &req : read_reqs) { + if (req.buf == nullptr || req.len == 0 || + req.offset > std::numeric_limits::max() || + req.len > std::numeric_limits::max()) { + return IndexError_InvalidArgument; + } + const size_t offset = static_cast(req.offset); + const size_t length = static_cast(req.len); + if (ailego::kVectorPageSize < DiskAnnUtil::kSectorSize || + ailego::kVectorPageSize % DiskAnnUtil::kSectorSize != 0 || + offset % DiskAnnUtil::kSectorSize != 0 || + length % DiskAnnUtil::kSectorSize != 0 || + offset > pool_->file_size() || length > pool_->file_size() - offset) { + return IndexError_InvalidArgument; + } + const size_t first_page = offset / ailego::kVectorPageSize; + const size_t last_page = (offset + length - 1) / ailego::kVectorPageSize; + const size_t pages = last_page - first_page + 1; + if (pages > std::numeric_limits::max() - total_pages) { + return IndexError_InvalidLength; + } + total_pages += pages; + } + + // A pool may own the file without a page table when even metadata plus + // one page exceeds its budget. Keep batched direct I/O in this case and + // never enter page-cache operations that require initialized entries. + if (!pool_->cache_enabled()) { + if (ctx == nullptr && setup_io_ctx(ctx) != 0) { + return IndexError_Runtime; + } + const int ret = bypass_reader_.read(read_reqs, ctx); + if (ret == 0) { + for (const AlignedRead &req : read_reqs) { + pool_->record_bypass_read(static_cast(req.len)); + } + } + return ret; + } + + std::vector unique_pages; + std::vector occurrences; + unique_pages.reserve(total_pages); + occurrences.reserve(total_pages); + for (const AlignedRead &req : read_reqs) { + size_t source_offset = static_cast(req.offset); + size_t remaining = static_cast(req.len); + char *destination = static_cast(req.buf); + while (remaining != 0) { + const auto page_id = static_cast( + source_offset / ailego::kVectorPageSize); + const size_t offset_in_page = source_offset % ailego::kVectorPageSize; + const size_t copy_length = + std::min(remaining, ailego::kVectorPageSize - offset_in_page); + size_t unique_index = 0; + while (unique_index < unique_pages.size() && + unique_pages[unique_index].page_id != page_id) { + ++unique_index; + } + if (unique_index == unique_pages.size()) { + unique_pages.push_back(UniquePage{page_id, nullptr, false}); + } + size_t canonical_index = occurrences.size(); + for (size_t i = 0; i < occurrences.size(); ++i) { + const PageOccurrence &prior = occurrences[i]; + if (prior.unique_index == unique_index && + prior.offset_in_page == offset_in_page && + prior.length == copy_length) { + canonical_index = prior.canonical_index; + break; + } + } + occurrences.push_back(PageOccurrence{unique_index, destination, + offset_in_page, copy_length, + canonical_index}); + source_offset += copy_length; + destination += copy_length; + remaining -= copy_length; + } + } + + std::vector admitted_ids; + std::vector admitted_indices; + std::vector admitted_pages(unique_pages.size(), nullptr); + std::vector bypass_requests; + admitted_ids.reserve(unique_pages.size()); + admitted_indices.reserve(unique_pages.size()); + bypass_requests.reserve(unique_pages.size()); + + auto release_cached_pages = [&]() { + for (UniquePage &page : unique_pages) { + if (page.cached_page != nullptr) { + pool_->release_pages(&page.page_id, 1); + page.cached_page = nullptr; + } + } + }; + struct CachedPageGuard { + decltype(release_cached_pages) &release; + ~CachedPageGuard() { + release(); + } + } cached_page_guard{release_cached_pages}; + + for (size_t i = 0; i < unique_pages.size(); ++i) { + UniquePage &page = unique_pages[i]; + page.cached_page = pool_->try_acquire_buffer(page.page_id); + if (page.cached_page != nullptr) { + continue; + } + if (pool_->should_admit_page(page.page_id)) { + admitted_ids.push_back(page.page_id); + admitted_indices.push_back(i); + } else { + page.bypass_candidate = true; + } + } + + if (!admitted_ids.empty()) { + if (pool_->acquire_pages(admitted_ids.data(), admitted_ids.size(), + admitted_pages.data())) { + for (size_t i = 0; i < admitted_ids.size(); ++i) { + unique_pages[admitted_indices[i]].cached_page = admitted_pages[i]; + } + } else { + for (const size_t index : admitted_indices) { + unique_pages[index].bypass_candidate = true; + } + } + } + + size_t bypass_rechecks = 0; + size_t bypass_cache_joins = 0; + admitted_ids.clear(); + admitted_indices.clear(); + for (size_t i = 0; i < unique_pages.size(); ++i) { + UniquePage &page = unique_pages[i]; + if (!page.bypass_candidate) { + continue; + } + ++bypass_rechecks; + page.cached_page = pool_->try_acquire_buffer(page.page_id); + if (page.cached_page != nullptr) { + page.bypass_candidate = false; + ++bypass_cache_joins; + } else if (pool_->should_join_cache_path(page.page_id)) { + admitted_ids.push_back(page.page_id); + admitted_indices.push_back(i); + } + } + if (!admitted_ids.empty() && + pool_->acquire_pages(admitted_ids.data(), admitted_ids.size(), + admitted_pages.data())) { + for (size_t i = 0; i < admitted_ids.size(); ++i) { + UniquePage &page = unique_pages[admitted_indices[i]]; + page.cached_page = admitted_pages[i]; + page.bypass_candidate = false; + } + bypass_cache_joins += admitted_ids.size(); + } + pool_->record_bypass_recheck(bypass_rechecks, bypass_cache_joins); + + uint64_t run_offset = 0; + uint64_t run_length = 0; + char *run_destination = nullptr; + size_t bypass_bytes = 0; + auto flush_bypass_run = [&]() { + if (run_length != 0) { + bypass_requests.emplace_back(run_offset, run_length, run_destination); + run_length = 0; + } + }; + for (size_t i = 0; i < occurrences.size(); ++i) { + const PageOccurrence &occurrence = occurrences[i]; + const UniquePage &page = unique_pages[occurrence.unique_index]; + if (!page.bypass_candidate || occurrence.canonical_index != i) { + continue; + } + const uint64_t slice_offset = + static_cast(page.page_id) * ailego::kVectorPageSize + + occurrence.offset_in_page; + if (run_length != 0 && run_offset + run_length == slice_offset && + run_destination + run_length == occurrence.destination) { + run_length += occurrence.length; + } else { + flush_bypass_run(); + run_offset = slice_offset; + run_length = occurrence.length; + run_destination = occurrence.destination; + } + bypass_bytes += occurrence.length; + } + flush_bypass_run(); + + if (!bypass_requests.empty()) { + if (ctx == nullptr && setup_io_ctx(ctx) != 0) { + return IndexError_Runtime; + } + const int read_ret = bypass_reader_.read(bypass_requests, ctx); + if (read_ret != 0) { + return read_ret; + } + pool_->record_bypass_read(bypass_bytes, bypass_requests.size()); + } + + for (size_t i = 0; i < occurrences.size(); ++i) { + const PageOccurrence &occurrence = occurrences[i]; + const UniquePage &page = unique_pages[occurrence.unique_index]; + if (page.cached_page != nullptr) { + std::memcpy(occurrence.destination, + page.cached_page + occurrence.offset_in_page, + occurrence.length); + } else if (occurrence.canonical_index != i) { + std::memcpy(occurrence.destination, + occurrences[occurrence.canonical_index].destination, + occurrence.length); + } + } + return 0; + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } +} + +int BufferPoolAlignedFileReader::submit(PendingBatch &batch, + std::vector &read_reqs, + IOContext &ctx) { + batch.n_submitted = 0; + batch.n_reaped = 0; + batch.used_pread = false; +#if defined(__linux__) && !defined(__ANDROID__) + batch.cbs.clear(); + batch.cb_ptrs.clear(); +#elif defined(_WIN32) || defined(_WIN64) + batch.expected_lengths.clear(); + batch.completed.clear(); + batch.generation = 0; +#endif + + const int ret = read(read_reqs, ctx); + if (ret != 0) { + return ret; + } + batch.used_pread = true; + batch.n_submitted = static_cast(read_reqs.size()); + return 0; +} + +int BufferPoolAlignedFileReader::get_completed( + PendingBatch &batch, IOContext & /*ctx*/, int /*min_completed*/, + std::vector &completed_indices) { + completed_indices.clear(); + for (uint32_t i = batch.n_reaped; i < batch.n_submitted; ++i) { + completed_indices.push_back(i); + } + batch.n_reaped = batch.n_submitted; + return static_cast(completed_indices.size()); +} + +void BufferPoolAlignedFileReader::release_io_ctx(IOContext &ctx) { + bypass_reader_.release_io_ctx(ctx); +} + +#if !defined(_WIN32) && !defined(_WIN64) #if defined(__linux__) && !defined(__ANDROID__) int LinuxAlignedFileReader::submit(PendingBatch &batch, std::vector &read_reqs, diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 22b812a23..032ae595e 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -43,6 +43,7 @@ #if !defined(_WIN32) && !defined(_WIN64) #include #endif +#include #include #include #include @@ -50,6 +51,9 @@ #include "diskann_util.h" namespace zvec { +namespace ailego { +class VecBufferPool; +} namespace core { // IoBackend holds the selected backend for each thread. On Linux it also owns @@ -148,6 +152,10 @@ class AlignedFileReader { // POSIX backends keep their process-local queue resources for reuse; Windows // overrides this to close private file and completion-port handles. virtual void release_io_ctx(IOContext &ctx) = 0; + + virtual bool requires_io_context() const { + return true; + } }; // POSIX reader implementation. Linux selects io_uring, libaio, or pread; @@ -223,5 +231,34 @@ using PlatformAlignedFileReader = WindowsAlignedFileReader; using PlatformAlignedFileReader = LinuxAlignedFileReader; #endif +class BufferPoolAlignedFileReader : public AlignedFileReader { + public: + explicit BufferPoolAlignedFileReader( + std::shared_ptr pool); + ~BufferPoolAlignedFileReader() override; + + void open(const std::string &fname) override; + //! Capture the pool's file, not a potentially replaced pathname. Also works + //! when the pool cannot reserve page-table memory and uses bypass-only I/O. + int open_from_pool(const std::string &fname); + void close() override; + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false) override; + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) override; + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) override; + void release_io_ctx(IOContext &ctx) override; + + bool requires_io_context() const override { + return false; + } + + private: + std::shared_ptr pool_; + PlatformAlignedFileReader bypass_reader_; + bool opened_{false}; +}; + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 46d0aa13b..955938e22 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include namespace zvec { @@ -116,13 +117,14 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity, return IndexError_InvalidFormat; } + auto pool = storage->vec_buffer_pool(); auto cached_file = storage->file(); #if defined(_WIN32) || defined(_WIN64) // Windows DiskAnn must be able to close the single buffered handle before // opening its unbuffered IOCP handles. FileReadStorage's // alone_file_handle mode gives every Segment an independent handle, which // cannot be closed through IndexStorage and may be retained by the caller. - if (!cached_file) { + if (!pool && !cached_file) { LOG_ERROR( "DiskAnn on Windows requires FileReadStorage with " "proxima.file.read_storage.alone_file_handle disabled"); @@ -158,79 +160,102 @@ int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity, const auto file_path = storage->file_path(); int ret = 0; - reader_.reset(new PlatformAlignedFileReader()); + if (pool) { + if (ailego::kVectorPageSize < DiskAnnUtil::kSectorSize || + ailego::kVectorPageSize % DiskAnnUtil::kSectorSize != 0) { + LOG_ERROR( + "DiskAnn BufferPool page size is incompatible with the sector " + "size: page_size=%zu sector_size=%zu", + ailego::kVectorPageSize, + static_cast(DiskAnnUtil::kSectorSize)); + return IndexError_Unsupported; + } + storage_ = storage; + auto buffer_reader = + std::make_shared(std::move(pool)); + ret = buffer_reader->open_from_pool(file_path); + if (ret != 0) { + LOG_ERROR("Failed to capture DiskAnn buffer file, ret=%d", ret); + return ret; + } + reader_ = std::move(buffer_reader); + } else { + reader_.reset(new PlatformAlignedFileReader()); #if defined(_WIN32) || defined(_WIN64) - // Drop every Segment reference created by entity.load() before checking the - // File control block. Without an external alias, only cached_file and the - // FileReadStorage itself remain as owners. - entity.release_storage(); - vector_segment.reset(); - if (cached_file.use_count() != 2) { - LOG_ERROR( - "DiskAnn on Windows cannot load while the caller retains the " - "FileReadStorage file or one of its segments"); - return IndexError_InvalidArgument; - } + // Drop every Segment reference created by entity.load() before checking the + // File control block. Without an external alias, only cached_file and the + // FileReadStorage itself remain as owners. + entity.release_storage(); + vector_segment.reset(); + if (cached_file.use_count() != 2) { + LOG_ERROR( + "DiskAnn on Windows cannot load while the caller retains the " + "FileReadStorage file or one of its segments"); + return IndexError_InvalidArgument; + } - // Capture the exact file object that supplied the in-memory metadata before - // releasing FileReadStorage. Reopening file_path after cleanup could bind - // graph reads to a replacement file while PQ/keys still belong to the old - // one. - ret = static_cast(reader_.get()) - ->open_from_handle(file_path, cached_file->native_handle()); -#else - if (cached_file) { - // POSIX atomic replacement leaves an open descriptor bound to the old - // inode. Capture an independent descriptor before cleanup so graph reads - // use the same file object that supplied the in-memory metadata. - ret = static_cast(reader_.get()) + // Capture the exact file object that supplied the in-memory metadata before + // releasing FileReadStorage. Reopening file_path after cleanup could bind + // graph reads to a replacement file while PQ/keys still belong to the old + // one. + ret = static_cast(reader_.get()) ->open_from_handle(file_path, cached_file->native_handle()); - } else { - // Preserve support for FileReadStorage's alone_file_handle mode. Its - // Segment abstraction does not expose a descriptor, so retain the - // origin/main ordering and bind the path before releasing the storage. - reader_->open(file_path); - } +#else + if (cached_file) { + // POSIX atomic replacement leaves an open descriptor bound to the old + // inode. Capture an independent descriptor before cleanup so graph reads + // use the same file object that supplied the in-memory metadata. + ret = static_cast(reader_.get()) + ->open_from_handle(file_path, cached_file->native_handle()); + } else { + // Preserve support for FileReadStorage's alone_file_handle mode. Its + // Segment abstraction does not expose a descriptor, so retain the + // origin/main ordering and bind the path before releasing the storage. + reader_->open(file_path); + } #endif - if (ret != 0) { - LOG_ERROR("Failed to capture DiskAnn index file, ret=%d", ret); - return ret; - } + if (ret != 0) { + LOG_ERROR("Failed to capture DiskAnn index file, ret=%d", ret); + return ret; + } - ret = storage->cleanup(); + ret = storage->cleanup(); #if !defined(_WIN32) && !defined(_WIN64) - entity.release_storage(); - vector_segment.reset(); + entity.release_storage(); + vector_segment.reset(); #endif - storage.reset(); - if (ret != 0) { - reader_->close(); - LOG_ERROR("Failed to release DiskAnn index storage, ret=%d", ret); - return ret; - } + storage.reset(); + if (ret != 0) { + reader_->close(); + LOG_ERROR("Failed to release DiskAnn index storage, ret=%d", ret); + return ret; + } #if defined(_WIN32) || defined(_WIN64) - // Windows cannot keep an ordinary buffered alias to this file object beside - // DiskAnn's unbuffered handles without a severe random-read regression. The - // preflight check above avoids consuming the storage on an ordinary - // ownership error. Check again after cleanup so an unexpected remaining - // owner cannot make the successful load retain a buffered handle. - if (cached_file.use_count() != 1) { - reader_->close(); - LOG_ERROR( - "DiskAnn on Windows cannot load while the caller retains the " - "FileReadStorage file or one of its segments"); - return IndexError_InvalidArgument; - } + // Windows cannot keep an ordinary buffered alias to this file object beside + // DiskAnn's unbuffered handles without a severe random-read regression. The + // preflight check above avoids consuming the storage on an ordinary + // ownership error. Check again after cleanup so an unexpected remaining + // owner cannot make the successful load retain a buffered handle. + if (cached_file.use_count() != 1) { + reader_->close(); + LOG_ERROR( + "DiskAnn on Windows cannot load while the caller retains the " + "FileReadStorage file or one of its segments"); + return IndexError_InvalidArgument; + } #endif - // Releasing the last internal reference closes the buffered source handle. - // POSIX caller-owned aliases remain valid; Windows has rejected them above. - cached_file.reset(); + // Releasing the last internal reference closes the buffered source handle. + // POSIX caller-owned aliases remain valid; Windows has rejected them above. + cached_file.reset(); + } - ret = setup_io_ctx(init_ctx_); - if (ret != 0) { - LOG_ERROR("setup io ctx error"); - return ret; + if (reader_->requires_io_context()) { + ret = setup_io_ctx(init_ctx_); + if (ret != 0) { + LOG_ERROR("setup io ctx error"); + return ret; + } } disk_bytes_per_point_ = meta_.element_size(); diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index 7c325169c..e0a8d75c0 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -58,6 +58,10 @@ class DiskAnnIndexer { int get_vector(diskann_id_t id, IndexContext::Pointer &context, std::string &vector); + bool requires_io_context() const { + return reader_ && reader_->requires_io_context(); + } + diskann_key_t get_key(diskann_id_t id) const; diskann_id_t get_id(diskann_key_t key) const; @@ -95,6 +99,7 @@ class DiskAnnIndexer { int cached_beam_search_impl(DiskAnnContext *ctx); DiskAnnEntity::Pointer entity_{}; + IndexStorage::Pointer storage_{}; IndexMeta meta_; diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index f1c48a859..77b1e4a66 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -459,7 +459,8 @@ IndexSearcher::Context::Pointer DiskAnnSearcher::create_context() const { } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), meta_.element_size(), + diskann_indexer_->requires_io_context())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 6b43ec959..ac0123b0a 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -534,7 +534,8 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } if (ailego_unlikely(ctx->init( DiskAnnContext::kSearcherContext, search_ctx_entity->max_degree(), - search_ctx_entity->pq_chunk_num(), meta_.element_size())) != 0) { + search_ctx_entity->pq_chunk_num(), meta_.element_size(), + diskann_indexer_->requires_io_context())) != 0) { LOG_ERROR("Init DiskAnn Context failed"); delete ctx; diff --git a/src/core/algorithm/ivf/ivf_builder.cc b/src/core/algorithm/ivf/ivf_builder.cc index 90c0ba995..285d37048 100644 --- a/src/core/algorithm/ivf/ivf_builder.cc +++ b/src/core/algorithm/ivf/ivf_builder.cc @@ -402,9 +402,26 @@ int IVFBuilder::build(IndexThreads::Pointer threads, if (holder->count() > 0) { holder_->reserve(holder->count()); } - for (auto iter = holder->create_iterator(); iter && iter->is_valid(); - iter->next()) { - holder_->emplace(iter->key(), iter->data()); + auto iter = holder->create_iterator(); + if (!iter) { + return IndexError_NoMemory; + } + for (; iter->is_valid(); iter->next()) { + const uint64_t key = iter->key(); + if (iter->status() != 0) { + return iter->status(); + } + const void *data = iter->data(); + if (iter->status() != 0) { + return iter->status(); + } + if (!data) { + return IndexError_ReadData; + } + holder_->emplace(key, data); + } + if (iter->status() != 0) { + return iter->status(); } converted_holder = holder_; } @@ -683,6 +700,9 @@ int IVFBuilder::build_label_index(IndexThreads *threads, return IndexError_Mismatch; } const void *data = iter->data(); + if (iter->status() != 0) { + return iter->status(); + } if (!data) { return IndexError_Runtime; } @@ -706,6 +726,9 @@ int IVFBuilder::build_label_index(IndexThreads *threads, LOG_INFO("Current built count:%zu", id); } } + if (iter && iter->status() != 0) { + return iter->status(); + } if (id != holder->count()) { return IndexError_Mismatch; } @@ -777,11 +800,23 @@ int IVFBuilder::dump_index(const IndexDumper::Pointer &dumper) { auto iter = quantizer->result()->create_iterator(); for (; iter->is_valid(); iter->next()) { uint32_t id = iter->key(); + if (iter->status() != 0) { + return iter->status(); + } + const void *data = iter->data(); + if (iter->status() != 0) { + return iter->status(); + } + if (!data) { + return IndexError_ReadData; + } record_dumped_id(id); - ret = - ivf_dumper->dump_inverted_vector(i, holder_->key(id), iter->data()); + ret = ivf_dumper->dump_inverted_vector(i, holder_->key(id), data); ivf_check_error_code(ret); } + if (iter->status() != 0) { + return iter->status(); + } } } diff --git a/src/core/algorithm/ivf/ivf_entity.cc b/src/core/algorithm/ivf/ivf_entity.cc index 97771d81a..75b2a8f13 100644 --- a/src/core/algorithm/ivf/ivf_entity.cc +++ b/src/core/algorithm/ivf/ivf_entity.cc @@ -12,11 +12,127 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "ivf_entity.h" +#include #include +#include #include "ivf_utility.h" namespace zvec { namespace core { +namespace { + +//! Cursor over a scatter read. IVF blocks and vectors are logical byte +//! ranges, so neither needs to be aligned with an underlying storage page. +class ScatterCursor { + public: + explicit ScatterCursor(const IndexStorage::Segment::ScatterBlock &block) + : spans_(block.spans) {} + + const uint8_t *data() { + normalize(); + return span_index_ < spans_.size() ? spans_[span_index_].data + span_offset_ + : nullptr; + } + + size_t contiguous_size() { + normalize(); + return span_index_ < spans_.size() ? spans_[span_index_].size - span_offset_ + : 0; + } + + bool skip(size_t size) { + while (size != 0) { + normalize(); + if (span_index_ == spans_.size()) { + return false; + } + const size_t available = spans_[span_index_].size - span_offset_; + const size_t consumed = std::min(size, available); + span_offset_ += consumed; + size -= consumed; + } + return true; + } + + bool copy(void *output, size_t size) { + auto *dst = static_cast(output); + while (size != 0) { + normalize(); + if (span_index_ == spans_.size()) { + return false; + } + const auto &span = spans_[span_index_]; + const size_t available = span.size - span_offset_; + const size_t copied = std::min(size, available); + std::memcpy(dst, span.data + span_offset_, copied); + dst += copied; + span_offset_ += copied; + size -= copied; + } + return true; + } + + private: + void normalize() { + while (span_index_ < spans_.size() && + span_offset_ == spans_[span_index_].size) { + ++span_index_; + span_offset_ = 0; + } + } + + const std::vector &spans_; + size_t span_index_{0}; + size_t span_offset_{0}; +}; + +//! Calculate a row-major IVF block directly from its resident page spans. +//! Only a vector split by a page boundary is copied into scratch. +bool calculate_row_major_block(ScatterCursor *cursor, + IVFDistanceCalculator *calculator, + const void *query, size_t vector_count, + size_t element_size, size_t serialized_size, + uint8_t *scratch, float *distances, + bool calculate) { + if (element_size == 0 || + vector_count > std::numeric_limits::max() / element_size) { + return false; + } + const size_t vector_bytes = vector_count * element_size; + if (serialized_size < vector_bytes) { + return false; + } + if (!calculate) { + return cursor->skip(serialized_size); + } + + size_t processed = 0; + while (processed < vector_count) { + const size_t contiguous = cursor->contiguous_size(); + const size_t direct_count = + std::min(vector_count - processed, contiguous / element_size); + if (direct_count != 0) { + calculator->query_features_distance(query, cursor->data(), direct_count, + distances + processed); + if (!cursor->skip(direct_count * element_size)) { + return false; + } + processed += direct_count; + continue; + } + + if (!scratch || !cursor->copy(scratch, element_size)) { + return false; + } + calculator->query_features_distance(query, scratch, 1, + distances + processed); + ++processed; + } + return cursor->skip(serialized_size - vector_bytes); +} + +} // namespace + //! Initialize int IVFEntity::IVFReformerWrapper::init(const IndexMeta &imeta) { auto &name = imeta.reformer_name(); @@ -614,19 +730,14 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, std::vector distances(block_vecs); const size_t batch_size = kBatchBlocks; const size_t block_size = header_.block_size; + const size_t element_size = meta_.element_size(); + const bool row_major = meta_.major_order() == IndexMeta::MO_ROW; + if (row_major) { + scatter_vector_.resize(element_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 - 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, size) != size) { - LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); - return IndexError_ReadData; - } - //! Read keys size_t items = std::min(blocks * block_vecs, list_meta->vector_count - (i * block_vecs)); @@ -635,6 +746,24 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, return IndexError_ReadData; } + //! Read vectors after keys. A resident scatter read pins no pages while a + //! second segment is faulted into a small buffer pool. + const size_t off = list_meta->offset + i * block_size; + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + IndexStorage::Segment::ScatterBlock scatter_data; + if (row_major) { + if (inverted_->read_scatter(off, scatter_data, size) != size) { + LOG_ERROR("Failed to scatter-read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + } else if (inverted_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + ScatterCursor scatter_cursor(scatter_data); + //! Compute distances for each block for (size_t b = 0; b < blocks; ++b) { const size_t vecs_count = @@ -649,13 +778,39 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, ++(*context_stats->mutable_filtered_count()); } } + const size_t block_offset = b * block_size; + if (block_offset > size) { + LOG_ERROR("Invalid IVF block range, off=%zu, size=%zu", block_offset, + size); + return IndexError_ReadData; + } + const size_t serialized_size = std::min(block_size, size - block_offset); if (keeps == 0) { + if (row_major && !calculate_row_major_block( + &scatter_cursor, calculator_.get(), query, + vecs_count, element_size, serialized_size, + scatter_vector_.data(), distances.data(), false)) { + LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", + off + block_offset, serialized_size); + return IndexError_ReadData; + } continue; } - const void *block_data = static_cast(data) + b * block_size; - calculator_->query_features_distance(query, block_data, vecs_count, - distances.data()); + if (row_major) { + if (!calculate_row_major_block(&scatter_cursor, calculator_.get(), + query, vecs_count, element_size, + serialized_size, scatter_vector_.data(), + distances.data(), true)) { + LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", + off + block_offset, serialized_size); + return IndexError_ReadData; + } + } else { + const void *block_data = static_cast(data) + block_offset; + calculator_->query_features_distance(query, block_data, vecs_count, + distances.data()); + } *(context_stats->mutable_dist_calced_count()) += vecs_count; @@ -688,19 +843,14 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, std::vector distances(block_vecs); const size_t batch_size = kBatchBlocks; const size_t block_size = header_.block_size; + const size_t element_size = meta_.element_size(); + const bool row_major = meta_.major_order() == IndexMeta::MO_ROW; + if (row_major) { + scatter_vector_.resize(element_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 - 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, size) != size) { - LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); - return IndexError_ReadData; - } - //! Read keys size_t items = std::min(blocks * block_vecs, list_meta->vector_count - (i * block_vecs)); @@ -709,14 +859,49 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, return IndexError_ReadData; } + //! Read vectors after keys so scatter-page pins cover computation only. + const size_t off = list_meta->offset + i * block_size; + const size_t size = + std::min(blocks * block_size, + static_cast(header_.inverted_body_size - off)); + IndexStorage::Segment::ScatterBlock scatter_data; + if (row_major) { + if (inverted_->read_scatter(off, scatter_data, size) != size) { + LOG_ERROR("Failed to scatter-read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + } else if (inverted_->read(off, &data, size) != size) { + LOG_ERROR("Failed to read block, off=%zu, size=%zu", off, size); + return IndexError_ReadData; + } + ScatterCursor scatter_cursor(scatter_data); + //! Compute distances for each block 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; - calculator_->query_features_distance(query, block_data, vecs_count, - distances.data()); + const size_t block_offset = b * block_size; + if (block_offset > size) { + LOG_ERROR("Invalid IVF block range, off=%zu, size=%zu", block_offset, + size); + return IndexError_ReadData; + } + const size_t serialized_size = std::min(block_size, size - block_offset); + if (row_major) { + if (!calculate_row_major_block(&scatter_cursor, calculator_.get(), + query, vecs_count, element_size, + serialized_size, scatter_vector_.data(), + distances.data(), true)) { + LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", + off + block_offset, serialized_size); + return IndexError_ReadData; + } + } else { + const void *block_data = static_cast(data) + block_offset; + calculator_->query_features_distance(query, block_data, vecs_count, + distances.data()); + } for (size_t k = 0; k < vecs_count; ++k) { if (block_keys[k] != kInvalidKey) { uint32_t id = list_meta->id_offset + (i + b) * block_vecs + k; diff --git a/src/core/algorithm/ivf/ivf_entity.h b/src/core/algorithm/ivf/ivf_entity.h index d309e7367..3d2861ede 100644 --- a/src/core/algorithm/ivf/ivf_entity.h +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -171,17 +171,21 @@ class IVFEntity { return *static_cast(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 { - if (!mapping_) return nullptr; - const void *data = nullptr; - const size_t size = vector_count() * sizeof(uint32_t); - if (mapping_->read(0, &data, size) != size) { - return nullptr; + //! Test whether the sorted-rank to local-id mapping is available. + bool has_key_order_mapping() const { + return mapping_ != nullptr; + } + + //! Fetch a range from the key-order mapping (sorted rank -> local_id). + //! Explicit copying avoids retaining storage-owned transient pointers. + size_t get_key_order_mapping(size_t rank, uint32_t *out, size_t count) const { + if (!mapping_ || !out || count == 0 || rank >= vector_count()) { + return 0; } - return static_cast(data); + count = std::min(count, vector_count() - rank); + const size_t bytes = count * sizeof(uint32_t); + return mapping_->fetch(rank * sizeof(uint32_t), out, bytes) / + sizeof(uint32_t); } //! Retrieve vector by local id @@ -358,6 +362,8 @@ class IVFEntity { IndexStorage::Segment::Pointer features_{}; IndexStorage::Segment::Pointer integer_quantizer_params_{}; mutable std::string vector_{}; // temporary buffer for colomn major order + //! Temporary buffer for one row-major vector split across storage pages. + mutable std::vector scatter_vector_{}; float norm_value_{0.0f}; // normalize the inverted vector to orignal score bool norm_value_sqrt_{false}; // does the norm value need to sqrt InvertedIndexHeader header_; diff --git a/src/core/algorithm/ivf/ivf_index_provider.h b/src/core/algorithm/ivf/ivf_index_provider.h index 3caa98761..21afa122b 100644 --- a/src/core/algorithm/ivf/ivf_index_provider.h +++ b/src/core/algorithm/ivf/ivf_index_provider.h @@ -14,6 +14,7 @@ #pragma once #include +#include #include #include #include @@ -74,8 +75,8 @@ class IVFIndexProvider : public IndexProvider { public: SortedIterator(const IVFEntity::Pointer &entity) : entity_(entity) { count_ = entity_->vector_count(); - mapping_ = entity_->get_key_order_mapping(); - if (!mapping_) { + use_mapping_ = entity_->has_key_order_mapping(); + if (!use_mapping_) { // Fallback: compute sorting if mapping segment is unavailable fallback_.resize(count_); std::iota(fallback_.begin(), fallback_.end(), size_t(0)); @@ -89,33 +90,102 @@ class IVFIndexProvider : public IndexProvider { //! NOTICE: the vec feature will be changed after iterating to next, so //! the caller need to keep a copy of it before iterator to next vector const void *data() const override { - return entity_->get_vector(current_local_id()); + size_t local_id = current_local_id(); + if (local_id >= count_) { + return nullptr; + } + const void *result = entity_->get_vector(local_id); + if (!result) { + status_ = IndexError_ReadData; + } + return result; } //! Test if the iterator is valid bool is_valid() const override { - return pos_ < count_; + return status_ == 0 && pos_ < count_ && + (!use_mapping_ || ensure_mapping_chunk()); + } + + int status() const override { + return status_; } //! Retrieve primary key uint64_t key() const override { - return entity_->get_key(current_local_id()); + size_t local_id = current_local_id(); + if (local_id >= count_) { + return kInvalidKey; + } + const uint64_t result = entity_->get_key(local_id); + if (result == kInvalidKey) { + status_ = IndexError_ReadData; + } + return result; } //! Next iterator void next() override { - ++pos_; + if (status_ == 0 && pos_ < count_) { + ++pos_; + } } private: + bool ensure_mapping_chunk() const { + if (status_ != 0 || pos_ >= count_) { + return false; + } + if (pos_ >= mapping_chunk_begin_ && + pos_ - mapping_chunk_begin_ < mapping_chunk_.size()) { + return true; + } + mapping_chunk_begin_ = pos_; + const size_t chunk_count = + std::min(kMappingChunkEntries, count_ - mapping_chunk_begin_); + try { + mapping_chunk_.resize(chunk_count); + } catch (const std::bad_alloc &) { + status_ = IndexError_NoMemory; + return false; + } + if (entity_->get_key_order_mapping(mapping_chunk_begin_, + mapping_chunk_.data(), + chunk_count) != chunk_count) { + mapping_chunk_.clear(); + status_ = IndexError_ReadData; + return false; + } + if (std::any_of(mapping_chunk_.begin(), mapping_chunk_.end(), + [this](uint32_t id) { return id >= count_; })) { + mapping_chunk_.clear(); + status_ = IndexError_InvalidFormat; + return false; + } + return true; + } + size_t current_local_id() const { - return mapping_ ? static_cast(mapping_[pos_]) : fallback_[pos_]; + if (status_ != 0 || pos_ >= count_) { + return count_; + } + if (!use_mapping_) { + return fallback_[pos_]; + } + if (!ensure_mapping_chunk()) { + return count_; + } + return static_cast(mapping_chunk_[pos_ - mapping_chunk_begin_]); } //! Members + static constexpr size_t kMappingChunkEntries = 4096; IVFEntity::Pointer entity_; - const uint32_t *mapping_{nullptr}; // points into mapping_ segment data - std::vector fallback_; // used only if mapping_ unavailable + bool use_mapping_{false}; + mutable int status_{0}; + mutable std::vector mapping_chunk_; + mutable size_t mapping_chunk_begin_{0}; + std::vector fallback_; // used only if mapping_ unavailable size_t count_{0}; size_t pos_{0}; }; diff --git a/src/core/framework/index_helper.cc b/src/core/framework/index_helper.cc index 15534116a..22848143c 100644 --- a/src/core/framework/index_helper.cc +++ b/src/core/framework/index_helper.cc @@ -114,25 +114,59 @@ class TwoPassIndexHolder : public IndexHolder { //! Retrieve pointer of data const void *data() const override { - return front_iter_->data(); + if (status() != 0) { + return nullptr; + } + const void *result = front_iter_->data(); + holder_->status_ = front_iter_->status(); + if (!result && holder_->status_ == 0) { + holder_->status_ = IndexError_ReadData; + } + return holder_->status_ == 0 ? result : nullptr; } //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + if (status() != 0) { + return false; + } + const bool valid = front_iter_->is_valid(); + holder_->status_ = front_iter_->status(); + return holder_->status_ == 0 && valid; + } + + int status() const override { + if (holder_->status_ == 0) { + holder_->status_ = front_iter_->status(); + } + return holder_->status_; } //! Retrieve primary key uint64_t key() const override { - return front_iter_->key(); + const uint64_t result = front_iter_->key(); + (void)status(); + return result; } //! Next iterator void next() override { + if (!is_valid()) { + return; + } + const uint64_t record_key = key(); + if (status() != 0) { + return; + } + const void *record_data = data(); + if (status() != 0) { + return; + } holder_->features_.emplace_back( - front_iter_->key(), std::string((const char *)front_iter_->data(), - holder_->front_->element_size())); + record_key, std::string(static_cast(record_data), + holder_->front_->element_size())); front_iter_->next(); + (void)status(); } private: @@ -160,7 +194,12 @@ class TwoPassIndexHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return (features_iter_ != holder_->features_.end()); + return holder_->status_ == 0 && + features_iter_ != holder_->features_.end(); + } + + int status() const override { + return holder_->status_; } //! Retrieve primary key @@ -218,10 +257,12 @@ class TwoPassIndexHolder : public IndexHolder { ++pass_; if (pass_ == 1) { IndexHolder::Iterator::Pointer iter = front_->create_iterator(); - return iter ? IndexHolder::Iterator::Pointer( - new TwoPassIndexHolder::FirstPassIterator( - this, std::move(iter))) - : IndexHolder::Iterator::Pointer(); + if (!iter) { + status_ = IndexError_NoMemory; + return nullptr; + } + return IndexHolder::Iterator::Pointer( + new TwoPassIndexHolder::FirstPassIterator(this, std::move(iter))); } else if (pass_ == 2) { return IndexHolder::Iterator::Pointer( new TwoPassIndexHolder::SecondPassIterator(this)); @@ -236,6 +277,7 @@ class TwoPassIndexHolder : public IndexHolder { //! Members IndexHolder::Pointer front_{}; std::list> features_{}; + int status_{0}; size_t pass_{0}; IndexMeta::DataType data_type_{IndexMeta::DataType::DT_UNDEFINED}; size_t dimension_; diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 2ceb13bf9..58ec4acce 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -57,11 +57,24 @@ class MergeSourceIndexHolder final : public core::IndexHolder { } bool is_valid() const override { - return owner_->error_ == 0 && source_iter_ && source_iter_->is_valid(); + if (owner_->error_ != 0 || !source_iter_) { + return false; + } + const bool valid = source_iter_->is_valid(); + owner_->error_ = source_iter_->status(); + return owner_->error_ == 0 && valid; + } + + int status() const override { + return owner_->error_; } uint64_t key() const override { - return source_iter_->key(); + const uint64_t result = source_iter_->key(); + if (owner_->error_ == 0) { + owner_->error_ = source_iter_->status(); + } + return result; } void next() override { @@ -74,6 +87,10 @@ class MergeSourceIndexHolder final : public core::IndexHolder { data_ = nullptr; if (owner_->error_ != 0) return; while (!source_iter_ || !source_iter_->is_valid()) { + if (source_iter_ && source_iter_->status() != 0) { + owner_->error_ = source_iter_->status(); + return; + } source_iter_.reset(); if (source_index_ >= owner_->sources_.size()) return; source_ = &owner_->sources_[source_index_++]; @@ -83,7 +100,16 @@ class MergeSourceIndexHolder final : public core::IndexHolder { return; } } + if (source_iter_->status() != 0) { + owner_->error_ = source_iter_->status(); + return; + } data_ = source_iter_->data(); + if (source_iter_->status() != 0) { + owner_->error_ = source_iter_->status(); + data_ = nullptr; + return; + } if (!data_) { owner_->error_ = core::IndexError_ReadData; return; @@ -611,6 +637,14 @@ int Index::close() { LOG_ERROR("Failed to cleanup streamer"); return core::IndexError_Runtime; } + // Contexts are cached per index type in thread-local storage. IVF contexts + // own cloned storage segments, so leaving the current thread's context in + // the cache after Close would keep the buffer pool (and its metadata/pages) + // alive until another IVF search or thread exit. + if (context_index_ < context_list.size()) { + context_list[context_index_].reset(); + context_index_ = std::numeric_limits::max(); + } if (ailego_unlikely(storage_->close() != 0)) { LOG_ERROR("Failed to close storage"); return core::IndexError_Runtime; diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 00b3fe83f..410dc9662 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -18,6 +18,7 @@ #include #if DISKANN_SUPPORTED #include "algorithm/diskann/diskann_params.h" +#include "utility/utility_params.h" #include "holder_builder.h" #endif @@ -139,11 +140,7 @@ int DiskAnnIndex::open(const std::string &file_path, file_path_ = file_path; is_read_only_ = storage_options.read_only; switch (storage_options.type) { - case StorageOptions::StorageType::kMMAP: - case StorageOptions::StorageType::kBufferPool: { - // NOTE: DiskAnn index is dumped via FileDumper (plain binary file), which - // is not compatible with BufferStorage's IndexFormat layout. Fall back to - // FileReadStorage for both MMAP and BufferPool storage types. + case StorageOptions::StorageType::kMMAP: { storage_ = core::IndexFactory::CreateStorage("FileReadStorage"); if (storage_ == nullptr) { LOG_ERROR("Failed to create FileReadStorage"); @@ -157,6 +154,22 @@ int DiskAnnIndex::open(const std::string &file_path, } break; } + case StorageOptions::StorageType::kBufferPool: { + storage_params.set(core::BUFFER_READ_STORAGE_WARMUP_MODE, + core::BUFFER_READ_STORAGE_WARMUP_NONE); + storage_ = core::IndexFactory::CreateStorage("BufferReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create BufferReadStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init BufferReadStorage, path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + break; + } default: { LOG_ERROR("Unsupported storage type"); return core::IndexError_Unsupported; diff --git a/src/core/interface/indexes/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc index 17059940a..2b00fe0a2 100644 --- a/src/core/interface/indexes/ivf_index.cc +++ b/src/core/interface/indexes/ivf_index.cc @@ -18,6 +18,7 @@ #include #include "algorithm/cluster/cluster_params.h" #include "algorithm/ivf/ivf_params.h" +#include "utility/utility_params.h" #include "holder_builder.h" namespace zvec::core_interface { @@ -92,21 +93,23 @@ 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 is immutable after training and FileDumper already emits the + // IndexFormat consumed by BufferReadStorage. Keep construction on the + // FileDumper path and use the bounded page cache after dump/reopen. + // Opening an index must not prewarm the entire file or displace other + // collections' cached pages. Populate the cache on demand instead. + storage_params.set(core::BUFFER_READ_STORAGE_WARMUP_MODE, + core::BUFFER_READ_STORAGE_WARMUP_NONE); + storage_ = core::IndexFactory::CreateStorage("BufferReadStorage"); if (storage_ == nullptr) { - LOG_ERROR( - "Failed to create MMapFileReadStorage (IVF buffer-pool fallback)"); + LOG_ERROR("Failed to create BufferReadStorage for IVF"); return core::IndexError_Runtime; } int ret = storage_->init(storage_params); if (ret != 0) { LOG_ERROR( - "Failed to init MMapFileReadStorage (IVF buffer-pool fallback), " - "path: %s, err: %s", + "Failed to init BufferReadStorage for IVF, path: %s, " + "err: %s", file_path_.c_str(), core::IndexError::What(ret)); return ret; } diff --git a/src/core/mixed_reducer/merged_provider_index_holder.cc b/src/core/mixed_reducer/merged_provider_index_holder.cc index 1ac0e0cae..618f11569 100644 --- a/src/core/mixed_reducer/merged_provider_index_holder.cc +++ b/src/core/mixed_reducer/merged_provider_index_holder.cc @@ -60,6 +60,9 @@ class MergedProviderIndexHolder::Iterator final : public IndexHolder::Iterator { const auto &source = owner_->sources_[source_index_]; const void *source_data = source_iter_->data(); + if (source_iter_->status() != 0) { + return this->fail(source_iter_->status(), "Failed to read source vector"); + } if (source_data == nullptr) { return this->fail(IndexError_Runtime, "Source provider returned a null vector"); @@ -91,8 +94,19 @@ class MergedProviderIndexHolder::Iterator final : public IndexHolder::Iterator { } bool is_valid() const override { - return owner_->status() == 0 && source_index_ < owner_->sources_.size() && - source_iter_ && source_iter_->is_valid(); + if (this->status() != 0 || source_index_ >= owner_->sources_.size() || + !source_iter_) { + return false; + } + const bool valid = source_iter_->is_valid(); + return this->status() == 0 && valid; + } + + int status() const override { + if (source_iter_) { + owner_->set_status(source_iter_->status()); + } + return owner_->status(); } uint64_t key() const override { @@ -138,6 +152,9 @@ class MergedProviderIndexHolder::Iterator final : public IndexHolder::Iterator { } while (source_iter_->is_valid()) { + if (this->status() != 0) { + return; + } if (source_ordinal_ >= source.iterated_count) { LOG_ERROR( "Source provider iteration grew after filter planning, " @@ -153,6 +170,9 @@ class MergedProviderIndexHolder::Iterator final : public IndexHolder::Iterator { ++source_ordinal_; } + if (this->status() != 0) { + return; + } if (source_ordinal_ != source.iterated_count) { LOG_ERROR( "Source provider iteration changed after filter planning, " @@ -219,6 +239,9 @@ class MergedProviderIndexHolder::OrdinalReader final const auto &source = owner_->sources_[source_index]; size_t ordinal = 0; for (; iter->is_valid(); iter->next(), ++ordinal) { + if (iter->status() != 0) { + return fail(iter->status()); + } if (owner_->canceled()) { return fail(IndexError_Canceled); } @@ -226,9 +249,16 @@ class MergedProviderIndexHolder::OrdinalReader final return fail(IndexError_Mismatch); } if (owner_->keep(source_index, ordinal)) { - keys_.push_back(iter->key()); + const uint64_t key = iter->key(); + if (iter->status() != 0) { + return fail(iter->status()); + } + keys_.push_back(key); } } + if (iter->status() != 0) { + return fail(iter->status()); + } if (ordinal != source.iterated_count) { return fail(IndexError_Mismatch); } @@ -404,24 +434,35 @@ int MergedProviderIndexHolder::init(const IndexFilter &filter, size_t ordinal = 0; bool data_validated = false; for (; iter->is_valid(); iter->next(), ++ordinal) { + if (iter->status() != 0) { + this->set_status(iter->status()); + return this->status(); + } if (this->canceled()) { this->set_status(IndexError_Canceled); return this->status(); } - if (iter->key() > - std::numeric_limits::max() - source.logical_id_base) { + const uint64_t key = iter->key(); + if (iter->status() != 0) { + this->set_status(iter->status()); + return this->status(); + } + if (key > std::numeric_limits::max() - source.logical_id_base) { this->set_status(IndexError_Overflow); return this->status(); } - bool keep_item = - !has_filter_ || !filter(source.logical_id_base + iter->key()); + bool keep_item = !has_filter_ || !filter(source.logical_id_base + key); if (has_filter_) { AppendBit(&source.keep_bits, ordinal, keep_item); } if (keep_item) { if (!data_validated) { const void *source_data = iter->data(); + if (iter->status() != 0) { + this->set_status(iter->status()); + return this->status(); + } if (source_data == nullptr) { LOG_ERROR("Source provider returned a null vector, source=%zu", source_index); @@ -461,6 +502,10 @@ int MergedProviderIndexHolder::init(const IndexFilter &filter, ++filtered_count_; } } + if (iter->status() != 0) { + this->set_status(iter->status()); + return this->status(); + } source.iterated_count = ordinal; logical_id_base += source.provider_count; } diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.cc b/src/core/mixed_reducer/mixed_streamer_reducer.cc index 85b31f269..89eb51f56 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.cc +++ b/src/core/mixed_reducer/mixed_streamer_reducer.cc @@ -43,6 +43,9 @@ int MaterializeMergedInput(const MergedProviderIndexHolder::Pointer &source, } for (; iter->is_valid(); iter->next()) { const void *data = iter->data(); + if (iter->status() != 0) { + return iter->status(); + } if (source->status() != 0) { return source->status(); } @@ -51,10 +54,17 @@ int MaterializeMergedInput(const MergedProviderIndexHolder::Pointer &source, } ailego::NumericalVector vector(source->dimension()); std::memcpy(vector.data(), data, source->element_size()); - if (!snapshot->emplace(iter->key(), std::move(vector))) { + const uint64_t key = iter->key(); + if (iter->status() != 0) { + return iter->status(); + } + if (!snapshot->emplace(key, std::move(vector))) { return IndexError_Mismatch; } } + if (iter->status() != 0) { + return iter->status(); + } if (source->status() != 0) { return source->status(); } @@ -278,9 +288,11 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { [](int item) { return item == 0; }); }; - if (!check_results(read_results)) { + const auto read_error = std::find_if(read_results.begin(), read_results.end(), + [](int item) { return item != 0; }); + if (read_error != read_results.end()) { LOG_ERROR("Get vector from entities failed"); - return IndexError_Runtime; + return *read_error; } if (!check_results(add_results)) { @@ -376,17 +388,27 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index, } while (iterator->is_valid()) { + if (iterator->status() != 0) { + return iterator->status(); + } if (stop_flag_ != nullptr && stop_flag_->load(std::memory_order_relaxed)) { LOG_DEBUG("read_vec cancelled."); return 0; } - if (filter(iterator->key() + (uint64_t)id_offset)) { + const uint64_t key = iterator->key(); + if (iterator->status() != 0) { + return iterator->status(); + } + if (filter(key + (uint64_t)id_offset)) { (*stats_.mutable_filtered_count())++; iterator->next(); continue; } const void *vector_data = iterator->data(); + if (iterator->status() != 0) { + return iterator->status(); + } if (!vector_data) { LOG_ERROR("Failed to read source vector, index=%zu key=%zu", source_streamer_index, static_cast(iterator->key())); @@ -432,7 +454,7 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index, } iterator->next(); } - return 0; + return iterator->status(); } void MixedStreamerReducer::add_vec(int *result) { diff --git a/src/core/quantizer/binary_converter.cc b/src/core/quantizer/binary_converter.cc index d5fcc5412..a0eed5fb3 100644 --- a/src/core/quantizer/binary_converter.cc +++ b/src/core/quantizer/binary_converter.cc @@ -50,7 +50,11 @@ class BinaryConverterHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -67,8 +71,13 @@ class BinaryConverterHolder : public IndexHolder { private: //! Encode the data by quantizer inline void encode_record() { - if (front_iter_->is_valid()) { + if (this->is_valid()) { const float *vec = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (vec == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } quantizer_->encode(vec, dim_ / 2, buffer_.data()); } } @@ -78,6 +87,7 @@ class BinaryConverterHolder : public IndexHolder { IndexHolder::Iterator::Pointer front_iter_{}; std::shared_ptr quantizer_{}; size_t dim_{0u}; + int status_{0}; }; //! Constructor @@ -228,4 +238,4 @@ class BinaryConverter : public IndexConverter { INDEX_FACTORY_REGISTER_CONVERTER(BinaryConverter); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/cosine_converter.cc b/src/core/quantizer/cosine_converter.cc index 60e6a43df..1290990a3 100644 --- a/src/core/quantizer/cosine_converter.cc +++ b/src/core/quantizer/cosine_converter.cc @@ -92,7 +92,11 @@ class CosineConverterHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -109,7 +113,13 @@ class CosineConverterHolder : public IndexHolder { private: //! Encode the data by quantizer void convert_record() { - if (!front_iter_->is_valid()) { + if (!this->is_valid()) { + return; + } + const void *source = front_iter_->data(); + status_ = front_iter_->status(); + if (source == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; return; } @@ -119,8 +129,7 @@ class CosineConverterHolder : public IndexHolder { if (original_type_ == IndexMeta::DataType::DT_FP16) { ::memcpy(reinterpret_cast(&normalize_buffer_[0]), - reinterpret_cast(front_iter_->data()), - original_element_size); + reinterpret_cast(source), original_element_size); ailego::Float16 *buf = reinterpret_cast(&normalize_buffer_[0]); @@ -134,9 +143,8 @@ class CosineConverterHolder : public IndexHolder { &norm, NORM_SIZE); } else if (owner_->raw_fp16_storage_) { auto *buf = reinterpret_cast(buffer_.data()); - owner_->fp16_convert_func_( - static_cast(front_iter_->data()), - original_dimension_, buf); + owner_->fp16_convert_func_(static_cast(source), + original_dimension_, buf); float norm = 0.0F; ailego::Normalizer::L2(buf, original_dimension_, @@ -144,8 +152,7 @@ class CosineConverterHolder : public IndexHolder { ::memcpy(buffer_.data() + element_size - NORM_SIZE, &norm, NORM_SIZE); } else { // original_type_ == IndexMeta::DataType::DT_FP32 ::memcpy(reinterpret_cast(&normalize_buffer_[0]), - reinterpret_cast(front_iter_->data()), - original_element_size); + reinterpret_cast(source), original_element_size); float *buf = reinterpret_cast(&normalize_buffer_[0]); const float *vec = buf; @@ -191,6 +198,7 @@ class CosineConverterHolder : public IndexHolder { std::string normalize_buffer_{}; std::vector rotate_buffer_; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; size_t dimension_{0u}; size_t original_dimension_{0u}; IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED}; diff --git a/src/core/quantizer/half_float_converter.cc b/src/core/quantizer/half_float_converter.cc index 9e519445f..57de22bd1 100644 --- a/src/core/quantizer/half_float_converter.cc +++ b/src/core/quantizer/half_float_converter.cc @@ -123,7 +123,11 @@ class HalfFloatHolder : public IndexHolder, public OrdinalAccessHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -139,16 +143,21 @@ class HalfFloatHolder : public IndexHolder, public OrdinalAccessHolder { private: inline void transform_record() { - if (front_iter_->is_valid()) { - owner_->convert_func_( - reinterpret_cast(front_iter_->data()), - buffer_.size(), buffer_.data()); + if (this->is_valid()) { + const auto *source = static_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (source == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } + owner_->convert_func_(source, buffer_.size(), buffer_.data()); } } const HalfFloatHolder *owner_{nullptr}; std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; //! Constructor diff --git a/src/core/quantizer/integer_quantizer_converter.cc b/src/core/quantizer/integer_quantizer_converter.cc index a08ddb0a7..cad75bbb3 100644 --- a/src/core/quantizer/integer_quantizer_converter.cc +++ b/src/core/quantizer/integer_quantizer_converter.cc @@ -54,7 +54,11 @@ class IntegerQuantizerConverterHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -71,8 +75,13 @@ class IntegerQuantizerConverterHolder : public IndexHolder { private: //! Encode the data by quantizer inline void encode_record() { - if (front_iter_->is_valid()) { + if (this->is_valid()) { const float *vec = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (vec == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } quantizer_->encode( vec, dim_, reinterpret_cast(buffer_.data())); @@ -84,6 +93,7 @@ class IntegerQuantizerConverterHolder : public IndexHolder { IndexHolder::Iterator::Pointer front_iter_{}; std::shared_ptr quantizer_{}; size_t dim_{0u}; + int status_{0}; }; //! Constructor @@ -515,7 +525,11 @@ class IntegerStreamingConverter : public IndexConverter { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -532,9 +546,14 @@ class IntegerStreamingConverter : public IndexConverter { private: //! Encode the data by quantizer void encode_record() { - if (front_iter_->is_valid()) { + if (this->is_valid()) { const float *vec = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (vec == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } size_t dim = owner_->dimension_; if (owner_->rotator_) { float *rotate_buf = @@ -562,6 +581,7 @@ class IntegerStreamingConverter : public IndexConverter { std::string normalize_buffer_{}; std::string rotate_buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; //! Constructor diff --git a/src/core/quantizer/mips_converter.cc b/src/core/quantizer/mips_converter.cc index cc57744ca..43d2b2acf 100644 --- a/src/core/quantizer/mips_converter.cc +++ b/src/core/quantizer/mips_converter.cc @@ -95,7 +95,11 @@ class MipsConverterHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -112,11 +116,16 @@ class MipsConverterHolder : public IndexHolder { private: //! Transform the data void transform_data() { - if (!front_iter_->is_valid()) { + if (!this->is_valid()) { return; } const float *src = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (src == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } float *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -133,6 +142,7 @@ class MipsConverterHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; //! Constructor @@ -224,7 +234,11 @@ class MipsConverterForcedHalfHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -240,11 +254,16 @@ class MipsConverterForcedHalfHolder : public IndexHolder { private: void transform_record() { - if (!front_iter_->is_valid()) { + if (!this->is_valid()) { return; } const float *src = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (src == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } ailego::Float16 *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -261,6 +280,7 @@ class MipsConverterForcedHalfHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; //! Constructor @@ -354,7 +374,11 @@ class MipsConverterHalfHolder : public IndexHolder { //! Test if the iterator is valid bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } //! Retrieve primary key @@ -370,12 +394,17 @@ class MipsConverterHalfHolder : public IndexHolder { private: void transform_record() { - if (!front_iter_->is_valid()) { + if (!this->is_valid()) { return; } const ailego::Float16 *src = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (src == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } ailego::Float16 *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -392,6 +421,7 @@ class MipsConverterHalfHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; //! Constructor diff --git a/src/core/quantizer/raw_uint8_converter.cc b/src/core/quantizer/raw_uint8_converter.cc index 5b42a18ea..eff00bab3 100644 --- a/src/core/quantizer/raw_uint8_converter.cc +++ b/src/core/quantizer/raw_uint8_converter.cc @@ -60,7 +60,11 @@ class RawUint8Holder : public IndexHolder { } bool is_valid() const override { - return front_iterator_->is_valid(); + return this->status() == 0 && front_iterator_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iterator_->status(); } uint64_t key() const override { @@ -74,16 +78,22 @@ class RawUint8Holder : public IndexHolder { private: void transform_record() { - if (front_iterator_->is_valid()) { - owner_->convert_func_( - static_cast(front_iterator_->data()), buffer_.size(), - buffer_.data()); + if (this->is_valid()) { + const auto *source = + static_cast(front_iterator_->data()); + status_ = front_iterator_->status(); + if (source == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } + owner_->convert_func_(source, buffer_.size(), buffer_.data()); } } const RawUint8Holder *owner_{nullptr}; std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iterator_{}; + int status_{0}; }; RawUint8Holder(IndexHolder::Pointer holder, turbo::ConvertFunc convert_func) diff --git a/src/core/quantizer/uniform_uint4_converter.cc b/src/core/quantizer/uniform_uint4_converter.cc index 604c6e718..b7d0541e5 100644 --- a/src/core/quantizer/uniform_uint4_converter.cc +++ b/src/core/quantizer/uniform_uint4_converter.cc @@ -331,7 +331,10 @@ class UniformUint4Converter : public IndexConverter { return buffer_.data(); } bool is_valid() const override { - return front_->is_valid(); + return this->status() == 0 && front_->is_valid(); + } + int status() const override { + return status_ != 0 ? status_ : front_->status(); } uint64_t key() const override { return front_->key(); @@ -343,12 +346,18 @@ class UniformUint4Converter : public IndexConverter { private: void encode() { - if (!front_->is_valid()) return; + if (!this->is_valid()) return; + const void *source = front_->data(); + status_ = front_->status(); + if (source == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } const float *input = nullptr; if (owner_->source_type_ == IndexMeta::DataType::DT_FP32) { - input = static_cast(front_->data()); + input = static_cast(source); } else { - DecodeSource(front_->data(), owner_->source_type_, + DecodeSource(source, owner_->source_type_, owner_->original_dimension_, &decoded_); input = decoded_.data(); } @@ -367,6 +376,7 @@ class UniformUint4Converter : public IndexConverter { std::vector buffer_{}; std::vector decoded_{}; IndexHolder::Iterator::Pointer front_{}; + int status_{0}; }; UniformUint4Holder(IndexHolder::Pointer front, size_t original_dimension, diff --git a/src/core/quantizer/uniform_uint7_converter.cc b/src/core/quantizer/uniform_uint7_converter.cc index 09a7b478e..8fcfa13c7 100644 --- a/src/core/quantizer/uniform_uint7_converter.cc +++ b/src/core/quantizer/uniform_uint7_converter.cc @@ -258,7 +258,11 @@ class UniformUint7Converter : public IndexConverter { } bool is_valid() const override { - return front_iter_->is_valid(); + return this->status() == 0 && front_iter_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : front_iter_->status(); } uint64_t key() const override { @@ -272,10 +276,15 @@ class UniformUint7Converter : public IndexConverter { private: void encode_record() { - if (!front_iter_->is_valid()) { + if (!this->is_valid()) { return; } const float *vec = reinterpret_cast(front_iter_->data()); + status_ = front_iter_->status(); + if (vec == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } int8_t *out = buffer_.data(); const float scale = owner_->scale_; const float bias = owner_->bias_; @@ -297,6 +306,7 @@ class UniformUint7Converter : public IndexConverter { const UniformUint7Holder *owner_{nullptr}; std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + int status_{0}; }; UniformUint7Holder(IndexHolder::Pointer front, size_t original_dim, diff --git a/src/core/quantizer/uniform_uint8_converter.cc b/src/core/quantizer/uniform_uint8_converter.cc index 05682cae9..fc0960c54 100644 --- a/src/core/quantizer/uniform_uint8_converter.cc +++ b/src/core/quantizer/uniform_uint8_converter.cc @@ -242,7 +242,11 @@ class UniformUint8Converter : public IndexConverter { } bool is_valid() const override { - return iterator_ && iterator_->is_valid(); + return iterator_ && this->status() == 0 && iterator_->is_valid(); + } + + int status() const override { + return status_ != 0 ? status_ : iterator_->status(); } uint64_t key() const override { @@ -259,14 +263,20 @@ class UniformUint8Converter : public IndexConverter { if (!is_valid()) { return; } - EncodeRecord(static_cast(iterator_->data()), - owner_->original_dimension_, owner_->scale_, owner_->bias_, - buffer_.data()); + const auto *source = static_cast(iterator_->data()); + status_ = iterator_->status(); + if (source == nullptr || status_ != 0) { + if (status_ == 0) status_ = IndexError_Runtime; + return; + } + EncodeRecord(source, owner_->original_dimension_, owner_->scale_, + owner_->bias_, buffer_.data()); } const UniformUint8Holder *owner_; std::vector buffer_; IndexHolder::Iterator::Pointer iterator_; + int status_{0}; }; UniformUint8Holder(IndexHolder::Pointer holder, size_t original_dimension, diff --git a/src/core/utility/buffer_read_storage.cc b/src/core/utility/buffer_read_storage.cc index 2e8a02d9d..9dfca8af9 100644 --- a/src/core/utility/buffer_read_storage.cc +++ b/src/core/utility/buffer_read_storage.cc @@ -732,7 +732,7 @@ class BufferReadStorage : public IndexStorage { } std::shared_ptr vec_buffer_pool() const override { - return cache_enabled_ ? buffer_pool_ : nullptr; + return buffer_pool_; } //! Path of the opened index file (diagnostics / backend consistency). diff --git a/src/core/utility/buffer_storage.cc b/src/core/utility/buffer_storage.cc index 66ad06a59..a78ea5eb8 100644 --- a/src/core/utility/buffer_storage.cc +++ b/src/core/utility/buffer_storage.cc @@ -19,7 +19,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -1047,7 +1049,11 @@ class BufferStorage : public IndexStorage { } std::shared_ptr vec_buffer_pool() const override { - return cache_enabled_ ? buffer_pool_ : nullptr; + return buffer_pool_; + } + + std::string file_path() const override { + return file_name_; } //! Initialize storage @@ -1067,7 +1073,6 @@ class BufferStorage : public IndexStorage { //! Open storage int open(const std::string &path, bool create_if_missing) override { - file_name_ = path; if (!ailego::File::IsExist(path) && create_if_missing) { size_t last_slash = path.rfind('/'); if (last_slash != std::string::npos) { @@ -1081,11 +1086,25 @@ class BufferStorage : public IndexStorage { } } - // create_if_missing also indicates write intent, matching MMapFileStorage. - buffer_pool_ = std::make_shared( - path, /*writable=*/create_if_missing); - buffer_pool_handle_ = - std::make_shared(buffer_pool_); + try { + // Capture both owners before replacing the published state. In + // particular, a failed reopen must not change existing segments' pool. + // create_if_missing also indicates write intent, like MMapFileStorage. + auto candidate_pool = std::make_shared( + path, /*writable=*/create_if_missing); + auto candidate_handle = + std::make_shared(candidate_pool); + file_name_ = path; + buffer_pool_ = std::move(candidate_pool); + buffer_pool_handle_ = std::move(candidate_handle); + } catch (const std::bad_alloc &) { + LOG_ERROR("Out of memory opening BufferStorage: %s", path.c_str()); + return IndexError_NoMemory; + } catch (const std::runtime_error &error) { + LOG_ERROR("Failed to open BufferStorage file %s: %s", path.c_str(), + error.what()); + return IndexError_OpenFile; + } int ret = parse_to_mapping(); if (ret != 0) { this->close_index(); diff --git a/src/include/zvec/ailego/buffer/vector_page_table.h b/src/include/zvec/ailego/buffer/vector_page_table.h index a298c83ee..824e898fa 100644 --- a/src/include/zvec/ailego/buffer/vector_page_table.h +++ b/src/include/zvec/ailego/buffer/vector_page_table.h @@ -718,6 +718,18 @@ class ZVEC_AILEGO_API VecBufferPool { return file_size_; } + //! The backing file remains available when the budget cannot fit a cache. + bool cache_enabled() const { + return initialized_; + } + + //! Borrow the page-data descriptor (a CRT descriptor on Windows). The pool + //! owns it: callers must not close it or change its flags, and must retain + //! the pool until they have captured their own handle to the same file. + int file_descriptor() const { + return fd_; + } + //! Sequentially preload pages into the pool until pool is full. void warmup(); diff --git a/src/include/zvec/core/framework/index_holder.h b/src/include/zvec/core/framework/index_holder.h index e2ab2640e..dfd3866d2 100644 --- a/src/include/zvec/core/framework/index_holder.h +++ b/src/include/zvec/core/framework/index_holder.h @@ -47,6 +47,13 @@ struct IndexHolder { //! Test if the iterator is valid virtual bool is_valid() const = 0; + //! Sticky iteration error; zero also denotes normal end of iteration. + //! Consumers must check after iteration and before using a key or data + //! that may require I/O. Wrapping iterators must preserve source errors. + virtual int status() const { + return 0; + } + //! Retrieve primary key virtual uint64_t key() const = 0; diff --git a/src/include/zvec/core/framework/index_provider.h b/src/include/zvec/core/framework/index_provider.h index aaeba299c..bdcf5b0d7 100644 --- a/src/include/zvec/core/framework/index_provider.h +++ b/src/include/zvec/core/framework/index_provider.h @@ -357,108 +357,171 @@ inline IndexProvider::Pointer convert_holder_to_provider( auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const ailego::Float16 *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } ailego::NumericalVector vec(dimension); std::memcpy(vec.data(), data, dimension * sizeof(ailego::Float16)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_FP32: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const float *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } ailego::NumericalVector vec(dimension); std::memcpy(vec.data(), data, dimension * sizeof(float)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_FP64: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const double *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } ailego::NumericalVector vec(dimension); std::memcpy(vec.data(), data, dimension * sizeof(double)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_INT8: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const int8_t *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } ailego::NumericalVector vec(dimension); std::memcpy(vec.data(), data, dimension * sizeof(int8_t)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_INT16: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const int16_t *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } ailego::NumericalVector vec(dimension); std::memcpy(vec.data(), data, dimension * sizeof(int16_t)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_BINARY32: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const uint32_t *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } size_t binary_size = (dimension + 31) / 32; ailego::BinaryVector vec(dimension); std::memcpy(vec.data(), data, binary_size * sizeof(uint32_t)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } case IndexMeta::DataType::DT_BINARY64: { auto provider = std::make_shared< MultiPassIndexProvider>(dimension); auto iter = holder->create_iterator(); + if (!iter) { + return nullptr; + } while (iter->is_valid()) { uint64_t key = iter->key(); + if (iter->status() != 0) { + return nullptr; + } const uint64_t *data = static_cast(iter->data()); + if (!data || iter->status() != 0) { + return nullptr; + } size_t binary_size = (dimension + 63) / 64; ailego::BinaryVector vec(dimension); std::memcpy(vec.data(), data, binary_size * sizeof(uint64_t)); provider->emplace(key, std::move(vec)); iter->next(); } - return provider; + return iter->status() == 0 ? provider : nullptr; } default: diff --git a/src/include/zvec/core/framework/index_storage.h b/src/include/zvec/core/framework/index_storage.h index 1d6745b64..b64592236 100644 --- a/src/include/zvec/core/framework/index_storage.h +++ b/src/include/zvec/core/framework/index_storage.h @@ -690,7 +690,8 @@ class IndexStorage : public IndexModule { return MemoryBlock::MBT_MMAP; } - //! Return the shared page cache when this storage is backed by VecBufferPool. + //! Return the backing pool, including bypass-only mode. Check cache_enabled() + //! before using its page cache; its backing file is available in either mode. virtual std::shared_ptr vec_buffer_pool() const { return nullptr; } diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 14ff98a81..5f1dd6cc5 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include #include @@ -251,7 +252,7 @@ class ZVEC_CORE_API Index { // converter_/reformer_/metric_ stay null. std::shared_ptr turbo_quantizer_{}; - size_t context_index_; + size_t context_index_{std::numeric_limits::max()}; core::IndexStorage::Pointer storage_{}; bool is_open_{false}; diff --git a/tests/core/algorithm/diskann/CMakeLists.txt b/tests/core/algorithm/diskann/CMakeLists.txt index a6974a975..86ece479c 100644 --- a/tests/core/algorithm/diskann/CMakeLists.txt +++ b/tests/core/algorithm/diskann/CMakeLists.txt @@ -20,7 +20,7 @@ else() if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") list(FILTER ALL_TEST_SRCS EXCLUDE REGEX - "/diskann_file_reader_aio_test\\.cc$") + "/diskann_file_reader_aio_test\\.cc$|/diskann_libaio_lifetime_test\\.cc$") endif() endif() diff --git a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc index bec83670a..8887005db 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc @@ -12,6 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include "diskann_file_reader.h" + +#if defined(__linux) || defined(__linux__) + #include #include #include @@ -23,7 +28,6 @@ #include #include #include -#include "diskann_file_reader.h" namespace zvec { namespace core { @@ -34,6 +38,7 @@ int execute_io_libaio(io_context_t &ctx, int fd, } // namespace zvec using namespace zvec::core; +namespace ailego = zvec::ailego; namespace { @@ -144,6 +149,10 @@ class TemporaryFile { return fd_; } + const char *path() const { + return path_; + } + private: char path_[64] = "DiskAnnLinuxAioTest.XXXXXX"; int fd_; @@ -203,6 +212,178 @@ TEST(DiskAnnLinuxAioTest, AccumulatesPartialSubmissionsAndCompletions) { std::free(output); } +TEST(DiskAnnBufferPoolFileReaderTest, + ReadsScatteredRequestsThroughOnePinnedPageBatch) { + if (ailego::kVectorPageSize != DiskAnnUtil::kSectorSize) { + GTEST_SKIP() << "DiskAnn sectors require one native buffer-pool page"; + } + + constexpr size_t kPageCount = 4; + const size_t page_size = ailego::kVectorPageSize; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kPageCount * page_size); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * page_size, static_cast(page + 1), + page_size); + } + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + ASSERT_EQ(::fsync(file.fd()), 0); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + 0, memory_pool.init( + 8 * page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(kPageCount))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + EXPECT_FALSE(reader.requires_io_context()); + reader.open(file.path()); + + void *output = nullptr; + ASSERT_EQ(::posix_memalign(&output, page_size, 4 * page_size), 0); + ASSERT_NE(output, nullptr); + std::memset(output, 0, 4 * page_size); + std::vector requests; + requests.emplace_back(3 * page_size, page_size, output); + requests.emplace_back(page_size, 2 * page_size, + static_cast(output) + page_size); + requests.emplace_back(3 * page_size, page_size, + static_cast(output) + 3 * page_size); + + IOContext unused{}; + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output, source.data() + 3 * page_size, page_size), 0); + EXPECT_EQ(std::memcmp(static_cast(output) + page_size, + source.data() + page_size, 2 * page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 3 * page_size, + source.data() + 3 * page_size, page_size), + 0); + EXPECT_EQ(pool->stats().miss, 3u); + + pool->page_table_.force_evict_all_loaded(); + std::free(output); +} + +TEST(DiskAnnBufferPoolFileReaderTest, + BypassesColdMissUnderPressureAndFansOutDuplicates) { + if (ailego::kVectorPageSize != DiskAnnUtil::kSectorSize) { + GTEST_SKIP() << "DiskAnn sectors require one native buffer-pool page"; + } + + constexpr size_t kPageCount = 4; + const size_t page_size = ailego::kVectorPageSize; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kPageCount * page_size); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * page_size, static_cast(page + 1), + page_size); + } + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + ASSERT_EQ(::fsync(file.fd()), 0); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + 0, memory_pool.init( + page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(kPageCount))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + reader.open(file.path()); + + char *seed = pool->acquire_buffer(0, 10); + ASSERT_NE(seed, nullptr); + + void *output = nullptr; + ASSERT_EQ(::posix_memalign(&output, page_size, 4 * page_size), 0); + ASSERT_NE(output, nullptr); + std::vector requests; + requests.emplace_back(page_size, 3 * page_size, output); + requests.emplace_back(2 * page_size, page_size, + static_cast(output) + 3 * page_size); + IOContext unused{}; + + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output, source.data() + page_size, page_size), 0); + EXPECT_EQ(std::memcmp(static_cast(output) + page_size, + source.data() + 2 * page_size, page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 2 * page_size, + source.data() + 3 * page_size, page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 3 * page_size, + source.data() + 2 * page_size, page_size), + 0); + EXPECT_FALSE(pool->is_page_resident(1)); + EXPECT_FALSE(pool->is_page_resident(2)); + EXPECT_FALSE(pool->is_page_resident(3)); + EXPECT_EQ(pool->stats().admission_rejected, 3u); + EXPECT_EQ(pool->stats().bypass_reads, 1u); + EXPECT_EQ(pool->stats().bypass_bytes, 3 * page_size); + EXPECT_EQ(pool->stats().bypass_io_requests, 1u); + EXPECT_EQ(pool->stats().bypass_rechecks, 3u); + EXPECT_EQ(pool->stats().bypass_cache_joins, 0u); + + // The second observation promotes these ghost entries for admission, but + // the only cache page is still pinned. Admission failure must fall back to + // direct I/O instead of failing the query. + std::memset(output, 0, 4 * page_size); + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output, source.data() + page_size, page_size), 0); + EXPECT_EQ(std::memcmp(static_cast(output) + page_size, + source.data() + 2 * page_size, page_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output) + 2 * page_size, + source.data() + 3 * page_size, page_size), + 0); + + pool->page_table_.release_block(0); + pool->page_table_.force_evict_all_loaded(); + EXPECT_EQ(destroy_io_ctx(unused), 0); + std::free(output); +} + +TEST(DiskAnnBufferPoolFileReaderTest, RejectsNonPageAlignedRequests) { + if (ailego::kVectorPageSize <= 512) { + GTEST_SKIP() << "test requires a native page larger than 512 bytes"; + } + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(2 * ailego::kVectorPageSize, 0x5a); + ASSERT_EQ(::pwrite(file.fd(), source.data(), source.size(), 0), + static_cast(source.size())); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ(0, memory_pool.init( + 4 * ailego::kVectorPageSize + + ailego::VecBufferPool::metadata_bytes_for_page_count(2))); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + + ASSERT_EQ(reader.open_from_pool(file.path()), 0); + void *output = allocate_aligned(512); + ASSERT_NE(output, nullptr); + std::vector requests; + requests.emplace_back(512, 512, output); + IOContext unused{}; + EXPECT_EQ(reader.read(requests, unused), IndexError_InvalidArgument); + EXPECT_EQ(unused, nullptr); + std::free(output); +} + +#endif // __linux__ + TEST(DiskAnnLinuxAioTest, DrainsPartialSubmissionBeforePreadFallback) { TemporaryFile file; ASSERT_GE(file.fd(), 0); diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc index f1b2d0877..df41f9ffd 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include using namespace zvec::core; @@ -167,6 +168,223 @@ TEST(DiskAnnFileReaderTest, ReadBeforeOpenReturnsError) { EXPECT_NE(reader.read(requests, ctx, false), 0); } +TEST(DiskAnnFileReaderTest, BufferPoolReadsSectorSlicesFromNativePages) { + namespace ailego = zvec::ailego; + const size_t native_page_size = ailego::kVectorPageSize; + const size_t sector_size = DiskAnnUtil::kSectorSize; + ASSERT_GE(native_page_size, sector_size); + ASSERT_EQ(native_page_size % sector_size, 0U); + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(2 * native_page_size); + for (size_t sector = 0; sector < source.size() / sector_size; ++sector) { + std::memset(source.data() + sector * sector_size, + static_cast(sector + 1), sector_size); + } + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + memory_pool.init(native_page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(2)), + 0); + auto pool = + std::make_shared(file.path(), /*writable=*/false); + ASSERT_EQ(pool->init(), 0); + ailego::block_id_t seed_page = 0; + ASSERT_NE(pool->acquire_buffer(seed_page, 10), nullptr); + + AlignedBuffer output = make_aligned_buffer(4 * sector_size); + ASSERT_NE(output, nullptr); + std::vector requests{ + {sector_size, sector_size, output.get()}, + {native_page_size - sector_size, 2 * sector_size, + static_cast(output.get()) + sector_size}, + {sector_size, sector_size, + static_cast(output.get()) + 3 * sector_size}, + }; + + BufferPoolAlignedFileReader reader(pool); + reader.open(file.path()); + IOContext unused{}; + ASSERT_EQ(reader.read(requests, unused), 0); + EXPECT_EQ(std::memcmp(output.get(), source.data() + sector_size, sector_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output.get()) + sector_size, + source.data() + native_page_size - sector_size, + 2 * sector_size), + 0); + EXPECT_EQ(std::memcmp(static_cast(output.get()) + 3 * sector_size, + source.data() + sector_size, sector_size), + 0); + EXPECT_GE(pool->stats().bypass_bytes, sector_size); + pool->release_pages(&seed_page, 1); + reader.release_io_ctx(unused); + EXPECT_EQ(destroy_io_ctx(unused), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, + BufferPoolCacheAndBypassFollowOriginalFileAfterPathReplacement) { + namespace ailego = zvec::ailego; + const size_t native_page_size = ailego::kVectorPageSize; + const size_t sector_size = DiskAnnUtil::kSectorSize; + TemporaryFile original; + TemporaryFile replacement; + ASSERT_GE(original.fd(), 0); + ASSERT_GE(replacement.fd(), 0); + std::vector original_data(2 * native_page_size, 0x3a); + std::vector replacement_data(original_data.size(), 0xc7); + ASSERT_TRUE(original.write_all(original_data.data(), original_data.size())); + ASSERT_TRUE( + replacement.write_all(replacement_data.data(), replacement_data.size())); + original.close(); + replacement.close(); + + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + memory_pool.init(native_page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count(2)), + 0); + auto pool = std::make_shared(original.path(), false); + ASSERT_EQ(pool->init(), 0); + ailego::block_id_t pinned_page = 0; + ASSERT_NE(pool->acquire_buffer(pinned_page, 10), nullptr); + + // Pin the only page the budget can hold. The cross-page request must use + // both the old cached page and direct I/O for the uncached second page. + ASSERT_EQ(::rename(replacement.path(), original.path()), 0); + BufferPoolAlignedFileReader reader(pool); + ASSERT_EQ(reader.open_from_pool(original.path()), 0); + AlignedBuffer output = make_aligned_buffer(2 * sector_size); + ASSERT_NE(output, nullptr); + const size_t offset = native_page_size - sector_size; + std::vector requests{{offset, 2 * sector_size, output.get()}}; + const auto before = pool->stats(); + IOContext ctx{}; + EXPECT_EQ(reader.read(requests, ctx), 0); + EXPECT_EQ( + std::memcmp(output.get(), original_data.data() + offset, 2 * sector_size), + 0); + const auto after = pool->stats(); + EXPECT_EQ(after.bypass_bytes - before.bypass_bytes, sector_size); + pool->release_pages(&pinned_page, 1); + reader.release_io_ctx(ctx); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); + pool.reset(); + EXPECT_EQ(memory_pool.used(), 0U); + EXPECT_EQ(memory_pool.metadata_used(), 0U); +} + +TEST(DiskAnnFileReaderTest, BufferPoolWithoutPageTableUsesBypassOnly) { + namespace ailego = zvec::ailego; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(2 * ailego::kVectorPageSize, 0x5a); + std::memset(source.data() + ailego::kVectorPageSize, 0xa5, + ailego::kVectorPageSize); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ(memory_pool.init(1), 0); + auto pool = std::make_shared(file.path(), false); + + // BufferReadStorage deliberately does not initialize a page table when + // metadata and one page cannot fit within the shared memory budget. + ASSERT_FALSE(pool->cache_enabled()); + BufferPoolAlignedFileReader reader(pool); + ASSERT_EQ(reader.open_from_pool(file.path()), 0); + AlignedBuffer output = make_aligned_buffer(2 * kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{ + {ailego::kVectorPageSize, kPageSize, output.get()}, + {0, kPageSize, static_cast(output.get()) + kPageSize}}; + IOContext ctx{}; + EXPECT_EQ(reader.read(requests, ctx), 0); + EXPECT_EQ(std::memcmp(output.get(), source.data() + ailego::kVectorPageSize, + kPageSize), + 0); + EXPECT_EQ(std::memcmp(static_cast(output.get()) + kPageSize, + source.data(), kPageSize), + 0); + EXPECT_EQ(pool->stats().bypass_bytes, 2 * kPageSize); + EXPECT_EQ(pool->stats().page_table_metadata_bytes, 0U); + EXPECT_EQ(memory_pool.used(), 0U); + EXPECT_EQ(memory_pool.metadata_used(), 0U); + reader.release_io_ctx(ctx); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderTest, BufferPoolReadBeforeOpenReturnsError) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(kPageSize, 0x5a); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + auto pool = std::make_shared(file.path(), false); + BufferPoolAlignedFileReader reader(pool); + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get()}}; + IOContext ctx{}; + EXPECT_NE(reader.read(requests, ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderTest, BufferPoolRejectsMissingPool) { + BufferPoolAlignedFileReader reader(nullptr); + EXPECT_NE(reader.open_from_pool("missing"), 0); + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get()}}; + IOContext ctx{}; + EXPECT_NE(reader.read(requests, ctx), 0); + PendingBatch batch; + EXPECT_NE(reader.submit(batch, requests, ctx), 0); + EXPECT_EQ(batch.n_submitted, 0U); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderTest, BufferPoolRejectsUnalignedRequestsAfterOpen) { + namespace ailego = zvec::ailego; + const size_t sector_size = DiskAnnUtil::kSectorSize; + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector source(2 * ailego::kVectorPageSize, 0x5a); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + file.close(); + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + memory_pool.init(16 * ailego::kVectorPageSize + + ailego::VecBufferPool::metadata_bytes_for_page_count(2)), + 0); + auto pool = std::make_shared(file.path(), false); + ASSERT_EQ(pool->init(), 0); + BufferPoolAlignedFileReader reader(pool); + ASSERT_EQ(reader.open_from_pool(file.path()), 0); + AlignedBuffer output = make_aligned_buffer(sector_size); + ASSERT_NE(output, nullptr); + + IOContext ctx{}; + std::vector bad_offset{{1, sector_size, output.get()}}; + EXPECT_EQ(reader.read(bad_offset, ctx), IndexError_InvalidArgument); + EXPECT_EQ(ctx, nullptr); + std::vector bad_length{{0, sector_size - 1, output.get()}}; + EXPECT_EQ(reader.read(bad_length, ctx), IndexError_InvalidArgument); + EXPECT_EQ(ctx, nullptr); + + std::vector valid{{0, sector_size, output.get()}}; + EXPECT_EQ(reader.read(valid, ctx), 0); + EXPECT_EQ(std::memcmp(output.get(), source.data(), sector_size), 0); + reader.release_io_ctx(ctx); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + TEST(DiskAnnFileReaderTest, OpenFromHandleSurvivesPathReplacementBeforeHandoff) { TemporaryFile original; diff --git a/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc index 5a515f92a..dbc7e44b5 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include "diskann_file_reader.h" @@ -258,6 +259,106 @@ class ScopedCurrentDirectory { } // namespace +TEST(DiskAnnFileReaderWindowsTest, BufferPoolReadsCachedAndBypassedPages) { + namespace ailego = zvec::ailego; + const size_t native_page_size = ailego::kVectorPageSize; + ASSERT_GE(native_page_size, kPageSize); + ASSERT_EQ(native_page_size % kPageSize, 0U); + const size_t native_page_count = kPageSize * kPageCount / native_page_size; + ASSERT_GE(native_page_count, 2U); + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ( + memory_pool.init(native_page_size + + ailego::VecBufferPool::metadata_bytes_for_page_count( + native_page_count)), + 0); + auto pool = std::make_shared(file.path(), false); + ASSERT_EQ(pool->init(), 0); + ailego::block_id_t pinned_page = 0; + ASSERT_NE(pool->acquire_buffer(pinned_page, 10), nullptr); + + BufferPoolAlignedFileReader reader(pool); + ASSERT_EQ(reader.open_from_pool(file.path()), 0); + AlignedBuffer output = make_aligned_buffer(2 * kPageSize); + ASSERT_NE(output, nullptr); + const size_t offset = native_page_size - kPageSize; + std::vector requests{{offset, 2 * kPageSize, output.get()}}; + const auto before = pool->stats(); + IOContext ctx = nullptr; + EXPECT_EQ(reader.read(requests, ctx), 0); + EXPECT_TRUE(verify_page(output.get(), offset / kPageSize)); + EXPECT_TRUE( + verify_page(output.get() + kPageSize, native_page_size / kPageSize)); + const auto after = pool->stats(); + EXPECT_EQ(after.bypass_bytes - before.bypass_bytes, kPageSize); + pool->release_pages(&pinned_page, 1); + reader.release_io_ctx(ctx); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); + pool.reset(); + EXPECT_EQ(memory_pool.used(), 0U); + EXPECT_EQ(memory_pool.metadata_used(), 0U); +} + +TEST(DiskAnnFileReaderWindowsTest, BufferPoolWithoutPageTableUsesBypassOnly) { + namespace ailego = zvec::ailego; + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + auto &memory_pool = ailego::MemoryLimitPool::get_instance(); + ASSERT_EQ(memory_pool.init(1), 0); + auto pool = std::make_shared(file.path(), false); + ASSERT_FALSE(pool->cache_enabled()); + BufferPoolAlignedFileReader reader(pool); + ASSERT_EQ(reader.open_from_pool(file.path()), 0); + AlignedBuffer output = make_aligned_buffer(2 * kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{7 * kPageSize, kPageSize, output.get()}, + {0, kPageSize, output.get() + kPageSize}}; + IOContext ctx = nullptr; + EXPECT_EQ(reader.read(requests, ctx), 0); + EXPECT_TRUE(verify_page(output.get(), 7)); + EXPECT_TRUE(verify_page(output.get() + kPageSize, 0)); + EXPECT_EQ(pool->stats().bypass_bytes, 2 * kPageSize); + EXPECT_EQ(pool->stats().page_table_metadata_bytes, 0U); + EXPECT_EQ(memory_pool.used(), 0U); + EXPECT_EQ(memory_pool.metadata_used(), 0U); + reader.release_io_ctx(ctx); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +TEST(DiskAnnFileReaderWindowsTest, BufferPoolReadBeforeOpenReturnsError) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + auto pool = std::make_shared(file.path(), false); + BufferPoolAlignedFileReader reader(pool); + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get()}}; + IOContext ctx = nullptr; + EXPECT_NE(reader.read(requests, ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, BufferPoolRejectsMissingPool) { + BufferPoolAlignedFileReader reader(nullptr); + EXPECT_NE(reader.open_from_pool("missing"), 0); + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get()}}; + IOContext ctx = nullptr; + EXPECT_NE(reader.read(requests, ctx), 0); + PendingBatch batch; + EXPECT_NE(reader.submit(batch, requests, ctx), 0); + EXPECT_EQ(batch.n_submitted, 0U); + EXPECT_EQ(ctx, nullptr); +} + TEST(DiskAnnFileReaderWindowsTest, OpenKeepsStableHandleUnbuffered) { TemporaryFile file; ASSERT_TRUE(file.valid()); diff --git a/tests/core/algorithm/diskann/diskann_libaio_lifetime_test.cc b/tests/core/algorithm/diskann/diskann_libaio_lifetime_test.cc new file mode 100644 index 000000000..5debd8341 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_libaio_lifetime_test.cc @@ -0,0 +1,95 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +class WorkerAioContext { + public: + ~WorkerAioContext() { + if (ctx_ == nullptr || LibAioLoader::Instance().io_destroy(ctx_) != 0) { + std::_Exit(EXIT_FAILURE); + } + std::fputs("libaio context destroyed at shutdown\n", stderr); + } + + io_context_t ctx_{nullptr}; +}; + +// Mirror GlobalResource: its thread pool is constructed before libaio is +// loaded, but its workers retain AIO contexts until shutdown. +class ShutdownWorker { + public: + void start() { + auto ready = ready_.get_future(); + worker_ = std::thread([this] { + // Keep cleanup on the worker stack: on musl, libstdc++ can delete its + // TLS destructor key before ShutdownWorker joins this thread. A + // thread_local destructor would then be skipped by the runtime. + WorkerAioContext context; + ready_.set_value(LibAioLoader::Instance().io_setup(1, &context.ctx_)); + stop_.get_future().wait(); + }); + if (ready.get() != 0) { + std::_Exit(EXIT_FAILURE); + } + } + + ~ShutdownWorker() { + stop_.set_value(); + worker_.join(); + } + + private: + std::promise ready_; + std::promise stop_; + std::thread worker_; +}; + +[[noreturn]] void exit_with_worker_context() { + // Registration order matters: the old loader destructor runs before the + // worker destructor and unloads the code needed by its io_destroy(). + static ShutdownWorker worker; + if (!LibAioLoader::Instance().load()) { + std::_Exit(EXIT_FAILURE); + } + worker.start(); + std::exit(EXIT_SUCCESS); +} + +} // namespace + +TEST(DiskAnnLibAioDeathTest, WorkerContextOutlivesStaticDestructors) { + // Check the optional dependency without constructing the loader singleton. + void *handle = dlopen("libaio.so.1", RTLD_LAZY); + if (handle == nullptr) { + handle = dlopen("libaio.so.1t64", RTLD_LAZY); + } + if (handle == nullptr) { + GTEST_SKIP() << "libaio is not installed"; + } + dlclose(handle); + + // Re-exec to keep singleton initialization independent of other tests. + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(exit_with_worker_context(), ::testing::ExitedWithCode(0), + "libaio context destroyed at shutdown"); +} diff --git a/tests/core/algorithm/ivf/ivf_builder_test.cc b/tests/core/algorithm/ivf/ivf_builder_test.cc index 6977c855d..dcb6f0c88 100644 --- a/tests/core/algorithm/ivf/ivf_builder_test.cc +++ b/tests/core/algorithm/ivf/ivf_builder_test.cc @@ -16,9 +16,11 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -72,6 +74,225 @@ void IVFBuilderTest::prepare_index_holder(uint32_t base_key, uint32_t num) { holder_.reset(holder); } +enum class IteratorErrorOperation { kValidity, kKey, kData }; + +class IteratorErrorHolder : public IndexHolder { + public: + IteratorErrorHolder(IndexHolder::Pointer delegate, size_t failed_ordinal, + IteratorErrorOperation operation) + : delegate_(std::move(delegate)), + failed_ordinal_(failed_ordinal), + operation_(operation) {} + + size_t count() const override { + return delegate_->count(); + } + size_t dimension() const override { + return delegate_->dimension(); + } + IndexMeta::DataType data_type() const override { + return delegate_->data_type(); + } + size_t element_size() const override { + return delegate_->element_size(); + } + bool multipass() const override { + return true; + } + IndexHolder::Iterator::Pointer create_iterator() override { + return std::make_unique(delegate_->create_iterator(), + failed_ordinal_, operation_); + } + + private: + class ErrorIterator : public IndexHolder::Iterator { + public: + ErrorIterator(IndexHolder::Iterator::Pointer delegate, + size_t failed_ordinal, IteratorErrorOperation operation) + : delegate_(std::move(delegate)), + failed_ordinal_(failed_ordinal), + operation_(operation) {} + + bool is_valid() const override { + return !failed(IteratorErrorOperation::kValidity) && + delegate_->is_valid(); + } + const void *data() const override { + return failed(IteratorErrorOperation::kData) ? nullptr + : delegate_->data(); + } + uint64_t key() const override { + return failed(IteratorErrorOperation::kKey) + ? std::numeric_limits::max() + : delegate_->key(); + } + void next() override { + ++ordinal_; + delegate_->next(); + } + int status() const override { + return status_; + } + + private: + bool failed(IteratorErrorOperation operation) const { + if (operation == operation_ && ordinal_ == failed_ordinal_) { + status_ = IndexError_ReadData; + } + return status_ != 0; + } + IndexHolder::Iterator::Pointer delegate_; + size_t failed_ordinal_; + IteratorErrorOperation operation_; + size_t ordinal_{0}; + mutable int status_{0}; + }; + + IndexHolder::Pointer delegate_; + size_t failed_ordinal_; + IteratorErrorOperation operation_; +}; + +class SinglePassIteratorErrorHolder : public IteratorErrorHolder { + public: + using IteratorErrorHolder::IteratorErrorHolder; + + bool multipass() const override { + return false; + } +}; + +TEST_F(IVFBuilderTest, TwoPassHolderDoesNotExposePartialCacheAfterReadError) { + prepare_index_holder(0, 8); + for (auto operation : + {IteratorErrorOperation::kValidity, IteratorErrorOperation::kKey, + IteratorErrorOperation::kData}) { + for (size_t failed_ordinal : {0U, 3U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto source = std::make_shared( + holder_, failed_ordinal, operation); + auto two_pass = IndexHelper::MakeTwoPassHolder(source); + ASSERT_NE(nullptr, two_pass); + ASSERT_NE(source.get(), two_pass.get()); + auto first = two_pass->create_iterator(); + ASSERT_NE(nullptr, first); + size_t read_count = 0; + for (; first->is_valid(); first->next()) { + (void)first->key(); + if (first->status() != 0) { + break; + } + (void)first->data(); + if (first->status() != 0) { + break; + } + ++read_count; + } + EXPECT_EQ(failed_ordinal, read_count); + EXPECT_EQ(IndexError_ReadData, first->status()); + first.reset(); + + auto second = two_pass->create_iterator(); + ASSERT_NE(nullptr, second); + // The first pass may have cached a prefix, but it is not a valid input. + EXPECT_FALSE(second->is_valid()); + EXPECT_EQ(IndexError_ReadData, second->status()); + } + } +} + +TEST_F(IVFBuilderTest, TwoPassHolderKeepsSuccessfulSecondPass) { + prepare_index_holder(0, 8); + auto source = std::make_shared( + holder_, holder_->count() + 1, IteratorErrorOperation::kValidity); + auto two_pass = IndexHelper::MakeTwoPassHolder(source); + ASSERT_NE(nullptr, two_pass); + ASSERT_NE(source.get(), two_pass.get()); + for (size_t pass = 0; pass < 2; ++pass) { + SCOPED_TRACE(pass); + auto iter = two_pass->create_iterator(); + ASSERT_NE(nullptr, iter); + size_t read_count = 0; + for (; iter->is_valid(); iter->next()) { + EXPECT_EQ(read_count, iter->key()); + const auto *data = static_cast(iter->data()); + ASSERT_NE(nullptr, data); + EXPECT_FLOAT_EQ(static_cast(read_count), data[0]); + EXPECT_EQ(0, iter->status()); + ++read_count; + } + EXPECT_EQ(holder_->count(), read_count); + EXPECT_EQ(0, iter->status()); + } +} + +TEST_F(IVFBuilderTest, MaterializationPreservesIteratorReadError) { + dimension_ = 16; + index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_); + params_.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4"); + prepare_index_holder(0, 128); + threads_ = std::make_shared(1, false); + for (auto operation : + {IteratorErrorOperation::kValidity, IteratorErrorOperation::kKey, + IteratorErrorOperation::kData}) { + for (size_t failed_ordinal : {0U, 65U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + IVFBuilder builder; + ASSERT_EQ(0, builder.init(index_meta_, params_)); + ASSERT_EQ(0, builder.train(threads_, holder_)); + auto failing = std::make_shared( + holder_, failed_ordinal, operation); + ASSERT_EQ(IndexError_ReadData, builder.build(threads_, failing)); + // The partial copy must not become a successfully built index. + EXPECT_EQ(IndexError_Runtime, builder.dump(nullptr)); + } + } +} + +TEST_F(IVFBuilderTest, ConvertedIteratorsPreserveSourceReadError) { + prepare_index_holder(0, 8); + for (const char *name : {"HalfFloatConverter", "CosineFp32Converter"}) { + for (auto operation : + {IteratorErrorOperation::kValidity, IteratorErrorOperation::kKey, + IteratorErrorOperation::kData}) { + for (size_t failed_ordinal : {0U, 3U}) { + SCOPED_TRACE(name); + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto converter = IndexFactory::CreateConverter(name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(index_meta_, Params{})); + auto failing = std::make_shared( + holder_, failed_ordinal, operation); + ASSERT_EQ(0, IndexConverter::TrainAndTransform(converter, failing)); + auto converted = converter->result(); + ASSERT_NE(nullptr, converted); + auto iter = converted->create_iterator(); + ASSERT_NE(nullptr, iter); + size_t read_count = 0; + for (; iter->is_valid(); iter->next()) { + (void)iter->key(); + if (iter->status() != 0) { + break; + } + ASSERT_NE(nullptr, iter->data()); + if (iter->status() != 0) { + break; + } + ++read_count; + } + EXPECT_EQ(failed_ordinal, read_count); + EXPECT_EQ(IndexError_ReadData, iter->status()); + iter->next(); + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(IndexError_ReadData, iter->status()); + } + } + } +} + // Defer execution until wait_finish() to exercise the worst case: producers // outrun all consumers. Track how many copied-vector batches remain live. class DeferredLabelThreads : public IndexThreads { diff --git a/tests/core/algorithm/ivf/ivf_index_provider_test.cc b/tests/core/algorithm/ivf/ivf_index_provider_test.cc new file mode 100644 index 000000000..599773098 --- /dev/null +++ b/tests/core/algorithm/ivf/ivf_index_provider_test.cc @@ -0,0 +1,326 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ivf_index_provider.h" +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core; + +namespace { + +constexpr size_t kMappingChunkEntries = 4096; +constexpr size_t kVectorCount = kMappingChunkEntries + 17; +constexpr size_t kDimension = 2; + +template +class TestSegment : public IndexStorage::Segment { + public: + explicit TestSegment(size_t count) : values_(count) {} + + size_t data_size() const override { + return values_.size() * sizeof(Value); + } + uint32_t data_crc() const override { + return 0; + } + size_t padding_size() const override { + return 0; + } + size_t capacity() const override { + return data_size(); + } + size_t fetch(size_t offset, void *buffer, size_t length) const override { + ++fetch_calls_; + if (offset > data_size() || length > data_size() - offset) { + return 0; + } + if (offset == failed_offset_) { + length = short_read_ && length != 0 ? length - 1 : 0; + } + if (length != 0) { + std::memcpy(buffer, + reinterpret_cast(values_.data()) + offset, + length); + } + return length; + } + size_t read(size_t offset, const void **data, size_t length) override { + if (offset > data_size() || length > data_size() - offset) { + *data = nullptr; + return 0; + } + *data = reinterpret_cast(values_.data()) + offset; + return length; + } + size_t read(size_t offset, IndexStorage::MemoryBlock &block, + size_t length) override { + const void *data = nullptr; + const size_t result = read(offset, &data, length); + block.reset(const_cast(data)); + return result; + } + size_t write(size_t, const void *, size_t) override { + return 0; + } + size_t resize(size_t) override { + return 0; + } + void update_data_crc(uint32_t) override {} + Pointer clone() override { + return std::make_shared>(*this); + } + + void fail_at(size_t offset, bool short_read) { + failed_offset_ = offset; + short_read_ = short_read; + } + void clear_failure() { + failed_offset_ = std::numeric_limits::max(); + } + size_t fetch_calls() const { + return fetch_calls_; + } + std::vector &values() { + return values_; + } + + private: + std::vector values_; + size_t failed_offset_{std::numeric_limits::max()}; + bool short_read_{false}; + mutable size_t fetch_calls_{0}; +}; + +class MappingStorage : public IndexStorage { + public: + int init(const zvec::ailego::Params &) override { + return 0; + } + int cleanup() override { + return 0; + } + int open(const std::string &, bool) override { + return 0; + } + int flush() override { + return 0; + } + int close() override { + return 0; + } + int append(const std::string &, size_t) override { + return IndexError_NotImplemented; + } + void refresh(uint64_t) override {} + uint64_t check_point() const override { + return 0; + } + Segment::Pointer get(const std::string &id, int = -1) override { + const auto it = segments.find(id); + return it == segments.end() ? nullptr : it->second; + } + bool has(const std::string &id) const override { + return segments.find(id) != segments.end(); + } + uint32_t magic() const override { + return 0; + } + + std::map segments; +}; + +struct MappingProvider { + MappingProvider() + : mapping(std::make_shared>(kVectorCount)), + entity(std::make_shared()) { + // IVFEntity contains a header with a flexible array member, so MSVC + // cannot use it as a base class. Load a real entity through its public + // storage interface instead of subclassing it to populate its fields. + auto storage = std::make_shared(); + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + std::string serialized_meta; + meta.serialize(&serialized_meta); + + InvertedIndexHeader header{}; + header.index_meta_size = static_cast(serialized_meta.size()); + header.header_size = sizeof(header) + header.index_meta_size; + header.total_vector_count = kVectorCount; + header.inverted_list_count = 1; + header.block_vector_count = kVectorCount; + header.block_size = kVectorCount * meta.element_size(); + header.block_count = 1; + header.inverted_body_size = header.block_size; + auto header_segment = + std::make_shared>(header.header_size); + std::memcpy(header_segment->values().data(), &header, sizeof(header)); + std::memcpy(header_segment->values().data() + sizeof(header), + serialized_meta.data(), serialized_meta.size()); + + auto keys = std::make_shared>(kVectorCount); + auto features = + std::make_shared>(kVectorCount * kDimension); + auto offsets = std::make_shared>( + kVectorCount * sizeof(InvertedVecLocation)); + for (size_t id = 0; id < kVectorCount; ++id) { + keys->values()[id] = kVectorCount - id; + for (size_t column = 0; column < kDimension; ++column) { + features->values()[id * kDimension + column] = + static_cast(kVectorCount - id); + } + mapping->values()[id] = static_cast(kVectorCount - id - 1); + const InvertedVecLocation location(id * meta.element_size(), false); + std::memcpy(offsets->values().data() + id * sizeof(location), &location, + sizeof(location)); + } + auto inverted_meta = std::make_shared>(1); + inverted_meta->values()[0].block_count = 1; + inverted_meta->values()[0].vector_count = kVectorCount; + storage->segments = {{IVF_INVERTED_HEADER_SEG_ID, header_segment}, + {IVF_INVERTED_BODY_SEG_ID, features}, + {IVF_INVERTED_META_SEG_ID, inverted_meta}, + {IVF_KEYS_SEG_ID, keys}, + {IVF_OFFSETS_SEG_ID, offsets}, + {IVF_MAPPING_SEG_ID, mapping}, + {IVF_FEATURES_SEG_ID, features}}; + + load_status = entity->load(storage); + if (load_status == 0) { + provider = std::make_shared(entity->meta(), entity, + "MappingProviderTest"); + } + } + + int load_status{IndexError_Runtime}; + std::shared_ptr> mapping; + IVFEntity::Pointer entity; + IndexProvider::Pointer provider; +}; + +} // namespace + +TEST(IVFIndexProviderTest, MappingFetchFailuresAreStickyAndNotNormalEof) { + for (size_t failed_rank : {size_t{0}, kMappingChunkEntries}) { + for (bool short_read : {false, true}) { + SCOPED_TRACE(::testing::Message() << "failed_rank=" << failed_rank + << " short_read=" << short_read); + MappingProvider fixture; + ASSERT_EQ(fixture.load_status, 0); + fixture.mapping->fail_at(failed_rank * sizeof(uint32_t), short_read); + auto iter = fixture.provider->create_iterator(); + ASSERT_NE(iter, nullptr); + for (size_t rank = 0; rank < failed_rank; ++rank) { + ASSERT_TRUE(iter->is_valid()); + EXPECT_EQ(iter->key(), rank + 1); + ASSERT_NE(iter->data(), nullptr); + iter->next(); + } + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(iter->status(), IndexError_ReadData); + const size_t failed_fetch_calls = fixture.mapping->fetch_calls(); + + // A consumer must reject even a complete first chunk instead of + // accepting a truncated provider when the next chunk cannot be read. + EXPECT_EQ(convert_holder_to_provider(fixture.provider), nullptr); + fixture.mapping->clear_failure(); + const size_t fetch_calls_after_conversion = + fixture.mapping->fetch_calls(); + EXPECT_GE(fetch_calls_after_conversion, failed_fetch_calls); + for (size_t retry = 0; retry < 3; ++retry) { + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(iter->status(), IndexError_ReadData); + EXPECT_EQ(iter->key(), kInvalidKey); + EXPECT_EQ(iter->data(), nullptr); + iter->next(); + } + EXPECT_EQ(fixture.mapping->fetch_calls(), fetch_calls_after_conversion); + + // The error belongs to the failed iterator, not to later traversals. + auto fresh = fixture.provider->create_iterator(); + ASSERT_NE(fresh, nullptr); + EXPECT_TRUE(fresh->is_valid()); + EXPECT_EQ(fresh->status(), 0); + } + } +} + +TEST(IVFIndexProviderTest, + SortedIterationOwnsMappingChunksAndEndsWithoutError) { + MappingProvider fixture; + ASSERT_EQ(fixture.load_status, 0); + auto iter = fixture.provider->create_iterator(); + ASSERT_NE(iter, nullptr); + ASSERT_TRUE(iter->is_valid()); + ASSERT_EQ(fixture.mapping->fetch_calls(), 1U); + + // Reusing storage-owned memory after fetch must not overwrite the chunk + // already copied into the iterator. Leave the next chunk valid to exercise + // the transition across the 4096-entry boundary as well. + std::fill_n(fixture.mapping->values().begin(), kMappingChunkEntries, + std::numeric_limits::max()); + for (size_t rank = 0; rank < kVectorCount; ++rank) { + ASSERT_TRUE(iter->is_valid()) << rank; + EXPECT_EQ(iter->status(), 0); + EXPECT_EQ(iter->key(), rank + 1); + const auto *data = static_cast(iter->data()); + ASSERT_NE(data, nullptr); + for (size_t column = 0; column < kDimension; ++column) { + EXPECT_FLOAT_EQ(data[column], static_cast(rank + 1)); + } + iter->next(); + } + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(iter->status(), 0); + EXPECT_EQ(fixture.mapping->fetch_calls(), 2U); +} + +TEST(IVFIndexProviderTest, InvalidMappingIdsHaveStickyFormatError) { + for (size_t corrupt_rank : {size_t{0}, kMappingChunkEntries}) { + SCOPED_TRACE(corrupt_rank); + MappingProvider fixture; + ASSERT_EQ(fixture.load_status, 0); + fixture.mapping->values()[corrupt_rank] = kVectorCount; + auto iter = fixture.provider->create_iterator(); + ASSERT_NE(iter, nullptr); + for (size_t rank = 0; rank < corrupt_rank; ++rank) { + ASSERT_TRUE(iter->is_valid()); + EXPECT_EQ(iter->key(), rank + 1); + iter->next(); + } + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(iter->status(), IndexError_InvalidFormat); + EXPECT_EQ(convert_holder_to_provider(fixture.provider), nullptr); + + fixture.mapping->values()[corrupt_rank] = + static_cast(kVectorCount - corrupt_rank - 1); + const size_t fetch_calls = fixture.mapping->fetch_calls(); + iter->next(); + EXPECT_FALSE(iter->is_valid()); + EXPECT_EQ(iter->status(), IndexError_InvalidFormat); + EXPECT_EQ(iter->key(), kInvalidKey); + EXPECT_EQ(iter->data(), nullptr); + EXPECT_EQ(fixture.mapping->fetch_calls(), fetch_calls); + + auto fresh = fixture.provider->create_iterator(); + ASSERT_NE(fresh, nullptr); + EXPECT_TRUE(fresh->is_valid()); + EXPECT_EQ(fresh->status(), 0); + } +} diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 8f204e73b..10f6072da 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include #include @@ -22,8 +23,10 @@ #include #include #include +#include #include #include +#include #include #include #include "tests/test_util.h" @@ -32,6 +35,7 @@ #include "zvec/core/framework/index_provider.h" #endif #include +#include #include #include #include @@ -783,6 +787,11 @@ TEST(IndexInterface, BufferGeneral) { auto func = [&](const BaseIndexParam::Pointer ¶m, const BaseIndexQueryParam::Pointer &query_param) { + const float value_tolerance = + param->quantizer_param && + param->quantizer_param->type == QuantizerType::kInt4 + ? 0.1f + : 1e-6f; std::string real_index_name = index_name; zvec::test_util::RemoveTestFiles(index_name + "*"); auto write_index = IndexFactory::CreateAndInitIndex(*param); @@ -797,6 +806,7 @@ TEST(IndexInterface, BufferGeneral) { VectorData vector_data; vector_data.vector = DenseVector{vector.data()}; ASSERT_TRUE(0 == write_index->add(vector_data, 233)); + ASSERT_TRUE(0 == write_index->train()); write_index->close(); auto read_index = IndexFactory::CreateAndInitIndex(*param); @@ -810,7 +820,7 @@ TEST(IndexInterface, BufferGeneral) { read_index->search(query, query_param, &result); ASSERT_EQ(1, result.doc_list_.size()); ASSERT_EQ(233, result.doc_list_[0].key()); - ASSERT_FLOAT_EQ(5.0f, result.doc_list_[0].score()); + ASSERT_NEAR(5.0f, result.doc_list_[0].score(), value_tolerance); if (query_param->fetch_vector) { auto &doc = result.doc_list_[0]; if (result.reverted_vector_list_.size() != 0) { @@ -818,12 +828,12 @@ TEST(IndexInterface, BufferGeneral) { ASSERT_EQ(1, result.reverted_vector_list_.size()); auto reverted_vector = reinterpret_cast( result.reverted_vector_list_[0].data()); - ASSERT_FLOAT_EQ(1.0f, reverted_vector[1]); - ASSERT_FLOAT_EQ(2.0f, reverted_vector[2]); + ASSERT_NEAR(1.0f, reverted_vector[1], value_tolerance); + ASSERT_NEAR(2.0f, reverted_vector[2], value_tolerance); } else { auto vector = reinterpret_cast(doc.vector()); - ASSERT_FLOAT_EQ(1.0f, vector[1]); - ASSERT_FLOAT_EQ(2.0f, vector[2]); + ASSERT_NEAR(1.0f, vector[1], value_tolerance); + ASSERT_NEAR(2.0f, vector[2], value_tolerance); } } @@ -834,8 +844,8 @@ TEST(IndexInterface, BufferGeneral) { float *fetched_vector = reinterpret_cast( std::get(fetched_vector_data.vector_buffer) .data.data()); - ASSERT_FLOAT_EQ(1.0f, fetched_vector[1]); - ASSERT_FLOAT_EQ(2.0f, fetched_vector[2]); + ASSERT_NEAR(1.0f, fetched_vector[1], value_tolerance); + ASSERT_NEAR(2.0f, fetched_vector[2], value_tolerance); result.doc_list_.clear(); read_index->close(); zvec::test_util::RemoveTestFiles(index_name + "*"); @@ -884,9 +894,337 @@ TEST(IndexInterface, BufferGeneral) { .with_fetch_vector(true) .with_ef_search(20) .build()); + func(IVFIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_n_list(10) + .build(), + IVFQueryParamBuilder().with_topk(10).with_fetch_vector(true).build()); + func(IVFIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_n_list(10) + .with_quantizer_param(QuantizerParam(QuantizerType::kFP16)) + .build(), + IVFQueryParamBuilder().with_topk(10).with_fetch_vector(true).build()); + func(IVFIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_n_list(10) + .with_quantizer_param(QuantizerParam(QuantizerType::kInt4)) + .build(), + IVFQueryParamBuilder().with_topk(10).with_fetch_vector(true).build()); +} + +TEST(IndexInterface, IvfBufferPoolDefersWarmupUntilReads) { + constexpr uint32_t kDimension = 256; + constexpr uint32_t kDocCount = 1024; + constexpr size_t kVectorBytes = kDimension * kDocCount * sizeof(float); + constexpr size_t kBufferBudget = 32UL * 1024UL * 1024UL; + const std::string index_name{"test_ivf_buffer_lazy_open.index"}; + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + const size_t previous_capacity = memory_pool.capacity(); + ASSERT_EQ(0u, memory_pool.used()); + auto cleanup = zvec::ailego::ScopeGuard::Make([&]() { + zvec::test_util::RemoveTestFiles(index_name + "*"); + EXPECT_EQ(0u, memory_pool.used()); + EXPECT_EQ(0, memory_pool.init(previous_capacity)); + }); + zvec::test_util::RemoveTestFiles(index_name + "*"); + ASSERT_EQ(0, memory_pool.init(kBufferBudget)); + + auto param = IVFIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_n_list(8) + .with_n_iters(4) + .build(); + auto verify_lazy_reads = [&](const Index::Pointer &index) { + const size_t open_page_bytes = memory_pool.stats().page_used; + ASSERT_GT(memory_pool.metadata_used(), 0u); + // Opening may read headers and centroids, but not all posting-list pages. + // The pool fits the entire index, so eviction cannot hide eager warmup. + ASSERT_LT(open_page_bytes, kVectorBytes / 2); + + for (uint32_t id : {0U, 127U, 511U, 895U, 1023U}) { + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(id, &fetched)); + const auto &buffer = std::get(fetched.vector_buffer); + ASSERT_EQ(kDimension * sizeof(float), buffer.data.size()); + const auto *values = reinterpret_cast(buffer.data.data()); + for (uint32_t dim = 0; dim < kDimension; ++dim) { + ASSERT_FLOAT_EQ(static_cast(id), values[dim]); + } + } + EXPECT_GT(memory_pool.stats().page_used, open_page_bytes); + + std::vector vector(kDimension, 511.0f); + VectorData query{DenseVector{vector.data()}}; + auto query_param = + IVFQueryParamBuilder().with_topk(1).with_nprobe(8).build(); + SearchResult result; + ASSERT_EQ(0, index->search(query, query_param, &result)); + ASSERT_EQ(1u, result.doc_list_.size()); + EXPECT_EQ(511u, result.doc_list_[0].key()); + }; + + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, + index->open(index_name, {StorageOptions::StorageType::kBufferPool, + /*create_new=*/true})); + std::vector vector(kDimension); + for (uint32_t id = 0; id < kDocCount; ++id) { + std::fill(vector.begin(), vector.end(), static_cast(id)); + ASSERT_EQ(0, index->add(VectorData{DenseVector{vector.data()}}, id)); + } + // train() dumps the file and opens its read storage before returning. + ASSERT_EQ(0, index->train()); + verify_lazy_reads(index); + ASSERT_EQ(0, index->close()); + } + ASSERT_EQ(0u, memory_pool.used()); + + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ( + 0, index->open(index_name, {StorageOptions::StorageType::kBufferPool, + /*create_new=*/false, /*read_only=*/true})); + verify_lazy_reads(index); + ASSERT_EQ(0, index->close()); + } +} + +TEST(IndexInterface, IvfBufferPoolSearchAfterOpenThreadExits) { + constexpr uint32_t kDimension = 768; + constexpr uint32_t kDocCount = 256; + // Enough for VecBufferPool metadata plus only a fraction of this index's + // data pages, so searches exercise real cache pressure rather than the + // bypass-only fallback. + constexpr size_t kBufferBudget = 1024 * 1024; + const std::string index_name{"test_ivf_buffer_eviction.index"}; + zvec::test_util::RemoveTestFiles(index_name + "*"); + + auto param = IVFIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_n_list(16) + .with_n_iters(4) + .build(); + + { + auto write_index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, write_index); + ASSERT_EQ(0, write_index->open(index_name, + {StorageOptions::StorageType::kMMAP, true})); + + std::vector vector(kDimension); + for (uint32_t id = 0; id < kDocCount; ++id) { + std::fill(vector.begin(), vector.end(), static_cast(id)); + VectorData vector_data{DenseVector{vector.data()}}; + ASSERT_EQ(0, write_index->add(vector_data, id)); + } + ASSERT_EQ(0, write_index->train()); + ASSERT_EQ(0, write_index->close()); + } + + ASSERT_EQ(0, + zvec::ailego::MemoryLimitPool::get_instance().init(kBufferBudget)); + auto read_index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, read_index); + + // Load on a short-lived thread. Buffer-backed sub-indexes must own every + // pointer they retain after Open rather than referencing that thread's TLS. + int open_rc = -1; + std::thread opener([&]() { + open_rc = read_index->open( + index_name, {StorageOptions::StorageType::kBufferPool, false, true}); + }); + opener.join(); + ASSERT_EQ(0, open_rc); + + auto query_param = IVFQueryParamBuilder() + .with_topk(1) + .with_nprobe(16) + .with_fetch_vector(true) + .build(); + auto search_one = [&](uint32_t id) { + std::vector query_vector(kDimension, static_cast(id)); + VectorData query{DenseVector{query_vector.data()}}; + SearchResult result; + if (read_index->search(query, query_param, &result) != 0 || + result.doc_list_.size() != 1 || result.doc_list_[0].key() != id) { + return false; + } + auto vector = static_cast(result.doc_list_[0].vector()); + return vector != nullptr && vector[0] == static_cast(id); + }; + + for (uint32_t id : {0U, 63U, 127U, 191U, 255U}) { + ASSERT_TRUE(search_one(id)); + } + + std::array concurrent_ok{}; + std::array workers; + for (size_t i = 0; i < workers.size(); ++i) { + workers[i] = std::thread([&, i]() { + concurrent_ok[i] = search_one(static_cast(i * 63)); + }); + } + for (auto &worker : workers) { + worker.join(); + } + for (bool ok : concurrent_ok) { + ASSERT_TRUE(ok); + } + + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + ASSERT_GT(memory_pool.metadata_used(), 0u); + ASSERT_GT(memory_pool.used(), memory_pool.metadata_used()); + + ASSERT_EQ(0, read_index->close()); + read_index.reset(); + ASSERT_EQ(0u, memory_pool.metadata_used()); + ASSERT_EQ(0u, memory_pool.used()); + zvec::test_util::RemoveTestFiles(index_name + "*"); + ASSERT_EQ(0, memory_pool.init(100 * 1024 * 1024)); } +#if DISKANN_SUPPORTED +TEST(IndexInterface, DiskAnnBufferPoolSearchAcrossBudgets) { + constexpr uint32_t kDimension = 64; + constexpr uint32_t kDocCount = 512; + constexpr size_t kFullBudget = 8UL * 1024UL * 1024UL; + const std::string index_name{"test_diskann_buffer_budgets.index"}; + auto &memory_pool = zvec::ailego::MemoryLimitPool::get_instance(); + const size_t previous_capacity = memory_pool.capacity(); + ASSERT_EQ(0u, memory_pool.used()); + auto cleanup = zvec::ailego::ScopeGuard::Make([&]() { + zvec::test_util::RemoveTestFiles(index_name + "*"); + EXPECT_EQ(0u, memory_pool.used()); + EXPECT_EQ(0, memory_pool.init(previous_capacity)); + }); + zvec::test_util::RemoveTestFiles(index_name + "*"); + ASSERT_EQ(0, memory_pool.init(kFullBudget)); + + std::mt19937 random(317); + std::uniform_real_distribution uniform(-2.0f, 2.0f); + std::vector values(kDocCount * kDimension); + for (float &value : values) { + value = uniform(random); + } + auto param = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_max_degree(24) + .with_list_size(64) + .with_pq_chunk_num(8) + .build(); + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ( + 0, index->open(index_name, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t id = 0; id < kDocCount; ++id) { + ASSERT_EQ( + 0, index->add( + VectorData{DenseVector{values.data() + id * kDimension}}, id)); + } + ASSERT_EQ(0, index->train()); + ASSERT_EQ(0, index->close()); + } + + const size_t file_bytes = ReadIndexBytesForTest(index_name).size(); + const size_t page_size = zvec::ailego::kVectorPageSize; + const size_t page_count = (file_bytes + page_size - 1) / page_size; + const size_t metadata_bytes = + zvec::ailego::VecBufferPool::metadata_bytes_for_page_count(page_count); + ASSERT_GT(file_bytes, 4 * page_size); + ASSERT_GT(metadata_bytes, 0u); + + auto query_param = std::make_shared(); + query_param->topk = 10; + query_param->list_size = 64; + query_param->fetch_vector = true; + using Result = std::vector>; + std::array baseline; + auto query = [&](const Index::Pointer &index, size_t number, Result *out) { + const auto id = static_cast(number * 61); + SearchResult result; + ASSERT_EQ(0, index->search( + VectorData{DenseVector{values.data() + id * kDimension}}, + query_param, &result)); + ASSERT_EQ(10u, result.doc_list_.size()); + out->clear(); + for (const auto &doc : result.doc_list_) { + out->emplace_back(doc.key(), doc.score()); + ASSERT_LT(doc.key(), kDocCount); + ASSERT_EQ(kDimension * sizeof(float), doc.vector_string().size()); + EXPECT_EQ(0, std::memcmp(doc.vector_string().data(), + values.data() + doc.key() * kDimension, + kDimension * sizeof(float))); + } + }; + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(index_name, {StorageOptions::StorageType::kMMAP, + false, true})); + for (size_t i = 0; i < baseline.size(); ++i) { + ASSERT_NO_FATAL_FAILURE(query(index, i, &baseline[i])); + } + ASSERT_EQ(0, index->close()); + } + + // Cover both bypass-only thresholds, then actual cache pressure and a pool + // large enough for the entire index. Every mode reads the same built file. + for (size_t budget : {page_size - 1, metadata_bytes + page_size - 1, + metadata_bytes + 4 * page_size, kFullBudget}) { + SCOPED_TRACE(budget); + ASSERT_EQ(0u, memory_pool.used()); + ASSERT_EQ(0, memory_pool.init(budget)); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, + index->open(index_name, {StorageOptions::StorageType::kBufferPool, + false, true})); + const bool cache_enabled = budget >= metadata_bytes + page_size; + EXPECT_EQ(cache_enabled ? metadata_bytes : 0u, memory_pool.metadata_used()); + for (size_t repeat = 0; repeat < 2; ++repeat) { + for (size_t i = 0; i < baseline.size(); ++i) { + Result actual; + ASSERT_NO_FATAL_FAILURE(query(index, i, &actual)); + EXPECT_EQ(baseline[i], actual); + if (cache_enabled) { + EXPECT_LE(memory_pool.stats().page_used, budget - metadata_bytes); + } else { + EXPECT_EQ(0u, memory_pool.metadata_used()); + EXPECT_EQ(0u, memory_pool.stats().page_used); + } + } + } + ASSERT_EQ(0, index->close()); + index.reset(); + EXPECT_EQ(0u, memory_pool.metadata_used()); + EXPECT_EQ(0u, memory_pool.stats().page_used); + EXPECT_EQ(0u, memory_pool.used()); + } +} +#endif // DISKANN_SUPPORTED + TEST(IndexInterface, SparseGeneral) { constexpr uint32_t kSparseCount = 3; const std::string index_name{"test.index"}; diff --git a/tests/core/mixed_reducer/merged_provider_index_holder_test.cc b/tests/core/mixed_reducer/merged_provider_index_holder_test.cc index c35c84908..f5944825e 100644 --- a/tests/core/mixed_reducer/merged_provider_index_holder_test.cc +++ b/tests/core/mixed_reducer/merged_provider_index_holder_test.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -285,6 +286,88 @@ IndexStreamer::Pointer MakeStreamer( return MakeStreamer(docs, std::make_shared()); } +enum class IteratorFailureOperation { kValidity, kKey, kData }; + +struct IteratorReadFailure { + IteratorFailureOperation operation{IteratorFailureOperation::kValidity}; + size_t ordinal{0}; + bool enabled{true}; +}; + +class FailingIteratorProvider final + : public MultiPassIndexProvider { + public: + FailingIteratorProvider(std::shared_ptr failure, + size_t count) + : MultiPassIndexProvider(kDimension), failure_(std::move(failure)) { + for (size_t i = 0; i < count; ++i) { + ailego::NumericalVector vector(kDimension); + vector[0] = static_cast(i); + vector[1] = static_cast(i) + 0.5F; + EXPECT_TRUE(this->emplace(i, std::move(vector))); + } + } + + IndexHolder::Iterator::Pointer create_iterator() override { + return std::make_unique( + MultiPassIndexProvider::create_iterator(), failure_); + } + + private: + class FailingIterator final : public IndexHolder::Iterator { + public: + FailingIterator(IndexHolder::Iterator::Pointer delegate, + std::shared_ptr failure) + : delegate_(std::move(delegate)), failure_(std::move(failure)) {} + + bool is_valid() const override { + return !failed(IteratorFailureOperation::kValidity) && + delegate_->is_valid(); + } + const void *data() const override { + return failed(IteratorFailureOperation::kData) ? nullptr + : delegate_->data(); + } + uint64_t key() const override { + return failed(IteratorFailureOperation::kKey) + ? std::numeric_limits::max() + : delegate_->key(); + } + void next() override { + delegate_->next(); + ++ordinal_; + } + int status() const override { + return status_; + } + + private: + bool failed(IteratorFailureOperation operation) const { + if (failure_->enabled && failure_->operation == operation && + failure_->ordinal == ordinal_) { + status_ = IndexError_ReadData; + } + return status_ != 0; + } + + IndexHolder::Iterator::Pointer delegate_; + std::shared_ptr failure_; + size_t ordinal_{0}; + mutable int status_{0}; + }; + + std::shared_ptr failure_; +}; + +IndexStreamer::Pointer MakeFailingIteratorStreamer( + const std::shared_ptr &failure, size_t count) { + return std::make_shared( + [failure, count](size_t) { + return std::make_shared(failure, count); + }, + std::make_shared()); +} + IndexStreamer::Pointer MakeBlockStreamer( const std::vector> &docs, const std::shared_ptr &stats, @@ -324,6 +407,138 @@ std::vector> ReadAll( return docs; } +TEST(MergedProviderIndexHolderTest, PlanningRejectsMappingErrorsBeforeFilter) { + for (auto operation : + {IteratorFailureOperation::kValidity, IteratorFailureOperation::kKey}) { + for (size_t failed_ordinal : {0U, 4096U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto failure = std::make_shared(); + failure->operation = operation; + failure->ordinal = failed_ordinal; + MergedProviderIndexHolder holder( + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + {MakeSource(MakeFailingIteratorStreamer(failure, 4098))}); + size_t filter_calls = 0; + IndexFilter filter; + filter.set([&](uint64_t) { + ++filter_calls; + return true; + }); + EXPECT_EQ(IndexError_ReadData, holder.init(filter)); + EXPECT_EQ(IndexError_ReadData, holder.status()); + EXPECT_EQ(failed_ordinal, filter_calls); + EXPECT_EQ(nullptr, holder.create_iterator()); + } + } +} + +TEST(MergedProviderIndexHolderTest, PlanningPreservesVectorReadError) { + auto failure = std::make_shared(); + failure->operation = IteratorFailureOperation::kData; + MergedProviderIndexHolder holder( + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + {MakeSource(MakeFailingIteratorStreamer(failure, 3))}); + EXPECT_EQ(IndexError_ReadData, holder.init({})); + EXPECT_EQ(IndexError_ReadData, holder.status()); +} + +TEST(MergedProviderIndexHolderTest, SequentialPassPreservesIteratorReadError) { + for (auto operation : + {IteratorFailureOperation::kValidity, IteratorFailureOperation::kData}) { + for (size_t failed_ordinal : {0U, 2U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto failure = std::make_shared(); + failure->operation = operation; + failure->ordinal = failed_ordinal; + failure->enabled = false; + MergedProviderIndexHolder holder( + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + {MakeSource(MakeFailingIteratorStreamer(failure, 3))}); + ASSERT_EQ(0, holder.init({})); + failure->enabled = true; + auto iter = holder.create_iterator(); + ASSERT_NE(nullptr, iter); + size_t read_count = 0; + for (; iter->is_valid(); iter->next()) { + (void)iter->data(); + if (iter->status() != 0) { + break; + } + ++read_count; + } + EXPECT_EQ(failed_ordinal, read_count); + EXPECT_EQ(IndexError_ReadData, iter->status()); + EXPECT_EQ(IndexError_ReadData, holder.status()); + EXPECT_FALSE(iter->is_valid()); + } + } +} + +TEST(MergedProviderIndexHolderTest, OrdinalKeyPassPreservesIteratorReadError) { + for (auto operation : + {IteratorFailureOperation::kValidity, IteratorFailureOperation::kKey}) { + for (size_t failed_ordinal : {0U, 2U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto failure = std::make_shared(); + failure->operation = operation; + failure->ordinal = failed_ordinal; + failure->enabled = false; + MergedProviderIndexHolder holder( + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension), + {MakeSource(MakeFailingIteratorStreamer(failure, 3))}); + ASSERT_EQ(0, holder.init({})); + failure->enabled = true; + OrdinalAccessHolder::Reader::Pointer reader; + EXPECT_EQ(IndexError_ReadData, holder.create_ordinal_reader(&reader)); + EXPECT_EQ(nullptr, reader); + EXPECT_EQ(IndexError_ReadData, holder.status()); + } + } +} + +TEST(MergedProviderIndexHolderTest, StreamerMergePreservesIteratorReadError) { + for (auto operation : + {IteratorFailureOperation::kValidity, IteratorFailureOperation::kKey, + IteratorFailureOperation::kData}) { + for (size_t failed_ordinal : {0U, 2U}) { + SCOPED_TRACE(static_cast(operation)); + SCOPED_TRACE(failed_ordinal); + auto failure = std::make_shared(); + failure->operation = operation; + failure->ordinal = failed_ordinal; + auto source = MakeFailingIteratorStreamer(failure, 3); + auto target = MakeStreamer({}); + ailego::ThreadPool pool(1, false); + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1); + ASSERT_EQ(0, reducer.init(params)); + reducer.set_thread_pool(&pool); + ASSERT_EQ(0, + reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, kDimension))); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(source, nullptr)); + size_t filter_calls = 0; + IndexFilter filter; + filter.set([&](uint64_t key) { + ++filter_calls; + return operation != IteratorFailureOperation::kData || + key < failed_ordinal; + }); + EXPECT_EQ(IndexError_ReadData, reducer.reduce(filter)); + EXPECT_EQ(failed_ordinal + + (operation == IteratorFailureOperation::kData ? 1U : 0U), + filter_calls); + // A failed read must not publish the target as ready for dumping. + EXPECT_EQ(IndexError_NoReady, reducer.dump(nullptr)); + } + } +} + class ReadFailureReformer : public IndexReformer { public: int init(const ailego::Params &) override { @@ -361,12 +576,12 @@ class RetainingTestBuilder : public IndexBuilder { int train(IndexThreads::Pointer threads, IndexHolder::Pointer) override { ++train_calls; train_thread_count = threads ? threads->count() : 0; - train_used_expected_pool = UsesExpectedPool(threads); + train_used_expected_pool = uses_expected_pool(threads); return 0; } int build(IndexThreads::Pointer threads, IndexHolder::Pointer input) override { build_thread_count = threads ? threads->count() : 0; - build_used_expected_pool = UsesExpectedPool(threads); + build_used_expected_pool = uses_expected_pool(threads); holder = std::move(input); return 0; } @@ -386,7 +601,7 @@ class RetainingTestBuilder : public IndexBuilder { IndexHolder::Pointer holder; private: - bool UsesExpectedPool(const IndexThreads::Pointer &threads) const { + bool uses_expected_pool(const IndexThreads::Pointer &threads) const { if (!threads || !expected_pool_) { return false; } diff --git a/tests/core/utility/buffer_read_storage_test.cc b/tests/core/utility/buffer_read_storage_test.cc index 5d1a741ee..81a5e506a 100644 --- a/tests/core/utility/buffer_read_storage_test.cc +++ b/tests/core/utility/buffer_read_storage_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include "utility/utility_params.h" @@ -56,7 +57,7 @@ class BufferReadStorageTest : public testing::Test { std::remove(file_path_.c_str()); } - IndexStorage::Pointer CreateStorage(const std::string &warmup_mode) { + IndexStorage::Pointer create_storage(const std::string &warmup_mode) { auto storage = IndexFactory::CreateStorage("BufferReadStorage"); EXPECT_NE(storage, nullptr); if (!storage) { @@ -74,13 +75,15 @@ class BufferReadStorageTest : public testing::Test { }; TEST_F(BufferReadStorageTest, NoneDefersPagePopulationUntilFirstRead) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto &pool = ailego::MemoryLimitPool::get_instance(); EXPECT_EQ(0u, pool.stats().page_used); - EXPECT_NE(nullptr, storage->vec_buffer_pool()); + ASSERT_NE(nullptr, storage->vec_buffer_pool()); + EXPECT_TRUE(storage->vec_buffer_pool()->cache_enabled()); + EXPECT_GE(storage->vec_buffer_pool()->file_descriptor(), 0); auto segment = storage->get("payload"); ASSERT_NE(segment, nullptr); @@ -91,7 +94,7 @@ TEST_F(BufferReadStorageTest, NoneDefersPagePopulationUntilFirstRead) { } TEST_F(BufferReadStorageTest, SequentialPreservesExistingWarmupBehavior) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); @@ -105,7 +108,7 @@ TEST_F(BufferReadStorageTest, SequentialPreservesExistingWarmupBehavior) { } TEST_F(BufferReadStorageTest, ResidentScatterReadReturnsPinnedPageSpans) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_SEQUENTIAL); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -130,7 +133,7 @@ TEST_F(BufferReadStorageTest, ResidentScatterReadReturnsPinnedPageSpans) { } TEST_F(BufferReadStorageTest, ColdScatterReadKeepsContiguousFallback) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -159,27 +162,33 @@ TEST_F(BufferReadStorageTest, PoolSmallerThanOnePageFallsBackToBypass) { const size_t kTooSmall = ailego::kVectorPageSize - 1; ASSERT_EQ(0, ailego::MemoryLimitPool::get_instance().init(kTooSmall)); - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().metadata_used); - EXPECT_EQ(nullptr, storage->vec_buffer_pool()); + ASSERT_NE(nullptr, storage->vec_buffer_pool()); + EXPECT_FALSE(storage->vec_buffer_pool()->cache_enabled()); + EXPECT_GE(storage->vec_buffer_pool()->file_descriptor(), 0); auto segment = storage->get("payload"); ASSERT_NE(segment, nullptr); std::string actual(payload_.size(), '\0'); ASSERT_EQ(actual.size(), segment->fetch(0, actual.data(), actual.size())); EXPECT_EQ(payload_, actual); + EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().page_used); } TEST_F(BufferReadStorageTest, PoolWithoutRoomForMetadataFallsBackToBypass) { ASSERT_EQ( 0, ailego::MemoryLimitPool::get_instance().init(ailego::kVectorPageSize)); - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().metadata_used); + ASSERT_NE(nullptr, storage->vec_buffer_pool()); + EXPECT_FALSE(storage->vec_buffer_pool()->cache_enabled()); + EXPECT_GE(storage->vec_buffer_pool()->file_descriptor(), 0); auto segment = storage->get("payload"); ASSERT_NE(segment, nullptr); @@ -187,10 +196,11 @@ TEST_F(BufferReadStorageTest, PoolWithoutRoomForMetadataFallsBackToBypass) { ASSERT_EQ(64u, segment->read(0, block, 64)); EXPECT_EQ(0, std::memcmp(payload_.data(), block.data(), 64)); EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + EXPECT_EQ(0u, ailego::MemoryLimitPool::get_instance().stats().page_used); } TEST_F(BufferReadStorageTest, CachePressureFallsBackToOwnedRead) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -212,7 +222,7 @@ TEST_F(BufferReadStorageTest, MemoryBlockKeepsPoolAliveAfterStorageClose) { IndexStorage::MemoryBlock block; IndexStorage::MemoryBlock copy; { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -234,7 +244,7 @@ TEST_F(BufferReadStorageTest, MemoryBlockKeepsPoolAliveAfterStorageClose) { } TEST_F(BufferReadStorageTest, BorrowedReadAvoidsOwningHandleOnResidentPage) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -253,7 +263,7 @@ TEST_F(BufferReadStorageTest, BorrowedReadAvoidsOwningHandleOnResidentPage) { } TEST_F(BufferReadStorageTest, PointerReadPinsUntilNextPointerRead) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -287,7 +297,7 @@ TEST_F(BufferReadStorageTest, PointerReadPinsUntilNextPointerRead) { } TEST_F(BufferReadStorageTest, PointerReadsUsePerThreadScratchBuffers) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); @@ -323,7 +333,7 @@ TEST_F(BufferReadStorageTest, PointerReadsUsePerThreadScratchBuffers) { } TEST_F(BufferReadStorageTest, MissingFileReturnsErrorWithoutThrowing) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); const std::string missing_path = file_path_ + ".missing"; @@ -334,7 +344,7 @@ TEST_F(BufferReadStorageTest, MissingFileReturnsErrorWithoutThrowing) { } TEST_F(BufferReadStorageTest, FailedReopenPreservesPublishedState) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); @@ -365,7 +375,7 @@ TEST_F(BufferReadStorageTest, RejectsOutOfRangeContainerOffset) { } TEST_F(BufferReadStorageTest, RangeChecksDoNotOverflow) { - auto storage = CreateStorage(BUFFER_READ_STORAGE_WARMUP_NONE); + auto storage = create_storage(BUFFER_READ_STORAGE_WARMUP_NONE); ASSERT_NE(storage, nullptr); ASSERT_EQ(0, storage->open(file_path_, false)); auto segment = storage->get("payload"); diff --git a/tests/core/utility/buffer_storage_write_test.cc b/tests/core/utility/buffer_storage_write_test.cc index f8fd5660d..c698fee49 100644 --- a/tests/core/utility/buffer_storage_write_test.cc +++ b/tests/core/utility/buffer_storage_write_test.cc @@ -22,7 +22,9 @@ #include #include #include +#include #include +#include #include #include @@ -70,6 +72,67 @@ class BufferStorageWriteTest : public ::testing::Test { std::string file_path_; }; +TEST_F(BufferStorageWriteTest, MissingFileReturnsErrorAndAllowsRetry) { + auto storage = IndexFactory::CreateStorage("BufferStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params{})); + int result = 0; + ASSERT_NO_THROW(result = storage->open(file_path_, false)); + EXPECT_EQ(IndexError_OpenFile, result); + EXPECT_EQ(nullptr, storage->vec_buffer_pool()); + EXPECT_TRUE(storage->file_path().empty()); + + { + auto writer = open_writable(); + ASSERT_NE(nullptr, writer); + ASSERT_EQ(0, writer->close()); + } + ASSERT_EQ(0, storage->open(file_path_, false)); + EXPECT_NE(nullptr, storage->vec_buffer_pool()); + EXPECT_EQ(file_path_, storage->file_path()); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(BufferStorageWriteTest, FailedFileCapturePreservesPublishedState) { + const std::string expected = "original storage survives failed reopen"; + { + auto writer = open_writable(); + ASSERT_NE(nullptr, writer); + ASSERT_EQ(0, writer->append("payload", 4096)); + auto segment = writer->get("payload"); + ASSERT_NE(nullptr, segment); + ASSERT_EQ(expected.size(), + segment->write(0, expected.data(), expected.size())); + ASSERT_EQ(0, writer->flush()); + ASSERT_EQ(0, writer->close()); + } + auto storage = open_read_only(); + ASSERT_NE(nullptr, storage); + auto published_pool = storage->vec_buffer_pool(); + auto published_segment = storage->get("payload"); + ASSERT_NE(nullptr, published_pool); + ASSERT_NE(nullptr, published_segment); + const std::string missing_path = file_path_ + ".missing"; + ailego::File::Delete(missing_path); + int result = 0; + ASSERT_NO_THROW(result = storage->open(missing_path, false)); + EXPECT_EQ(IndexError_OpenFile, result); + EXPECT_EQ(file_path_, storage->file_path()); + EXPECT_EQ(published_pool, storage->vec_buffer_pool()); + EXPECT_TRUE(storage->has("payload")); + + std::string actual(expected.size(), '\0'); + ASSERT_EQ(actual.size(), + published_segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(expected, actual); + auto segment = storage->get("payload"); + ASSERT_NE(nullptr, segment); + std::fill(actual.begin(), actual.end(), '\0'); + ASSERT_EQ(actual.size(), segment->fetch(0, actual.data(), actual.size())); + EXPECT_EQ(expected, actual); + EXPECT_EQ(0, storage->close()); +} + // ===== Basic Write Tests ===== // Test: Create new index via BufferStorage, append segment, write data, read @@ -1801,12 +1864,17 @@ TEST_F(BufferStorageWriteTest, CR_ReadOnlyMetadataPressureFallsBackToBypass) { { auto storage = open_read_only(); ASSERT_TRUE(storage); + ASSERT_NE(nullptr, storage->vec_buffer_pool()); + EXPECT_FALSE(storage->vec_buffer_pool()->cache_enabled()); + EXPECT_GE(storage->vec_buffer_pool()->file_descriptor(), 0); + EXPECT_EQ(0u, pool.stats().metadata_used); auto seg = storage->get("seg1"); ASSERT_TRUE(seg); IndexStorage::MemoryBlock block; ASSERT_EQ(expected.size(), seg->read(0, block, expected.size())); EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); EXPECT_EQ(0, std::memcmp(expected.data(), block.data(), expected.size())); + EXPECT_EQ(0u, pool.stats().page_used); ASSERT_EQ(0, storage->close()); } ASSERT_EQ(0, pool.init(64UL * 1024UL * 1024UL));