From 10ebe391d03c7e3324c90ebfdc5a3df05f0357f6 Mon Sep 17 00:00:00 2001 From: LHT129 Date: Thu, 13 Aug 2026 15:53:56 +0800 Subject: [PATCH] perf(io): batch cache misses in multiread Co-authored-by: opencode Signed-off-by: LHT129 --- src/io/common/basic_io.h | 398 +++++++++++++++++++++++--- src/io/read_cache/page_cache.cpp | 101 ++++++- src/io/read_cache/page_cache.h | 47 +++ src/io/read_cache/page_cache_test.cpp | 34 +++ src/io/read_cache/read_cache_test.cpp | 394 +++++++++++++++++++++++++ tests/test_ivf.cpp | 4 +- 6 files changed, 935 insertions(+), 43 deletions(-) diff --git a/src/io/common/basic_io.h b/src/io/common/basic_io.h index cbd2eabc11..a9fc66071b 100644 --- a/src/io/common/basic_io.h +++ b/src/io/common/basic_io.h @@ -23,7 +23,9 @@ #include #include #include +#include +#include "hash_types.h" #include "io/common/io_parameter.h" #include "io/read_cache/lru_page_cache.h" #include "io/read_cache/page.h" @@ -52,6 +54,14 @@ struct SupportsZeroSizeResize class BasicIO { +private: + using PageIdList = std::vector>; + + struct ReadCacheSnapshot { + std::shared_ptr cache; + uint64_t page_id_base{0}; + }; + public: /// Checks if the IO object is in-memory. static constexpr bool InMemory = IOTmpl::InMemory; @@ -65,7 +75,8 @@ class BasicIO { * * @param allocator A pointer to the Allocator object. */ - explicit BasicIO(Allocator* allocator) : allocator_(allocator){}; + explicit BasicIO(Allocator* allocator) + : allocator_(allocator), cached_direct_reads_(allocator){}; /** * @brief Writes data to the IO object at a specified offset. @@ -101,8 +112,9 @@ class BasicIO { Read(uint64_t size, uint64_t offset, uint8_t* data) const { static_assert(has_ReadImpl::value); if constexpr (not InMemory) { - if (cache_ != nullptr) { - return ReadCached(size, offset, data); + const auto cache = GetReadCacheSnapshot(); + if (cache.cache != nullptr) { + return ReadCached(size, offset, data, cache); } } return cast().ReadImpl(size, offset, data); @@ -123,7 +135,8 @@ class BasicIO { Read(uint64_t size, uint64_t offset, bool& need_release) const { static_assert(has_DirectReadImpl::value); if constexpr (not InMemory) { - if (cache_ != nullptr) { + const auto cache = GetReadCacheSnapshot(); + if (cache.cache != nullptr) { need_release = false; if (size == 0 or not IsValidRange(size, offset)) { return nullptr; @@ -132,11 +145,19 @@ class BasicIO { if (data == nullptr) { return nullptr; } - if (not ReadCached(size, offset, data)) { + if (not ReadCached(size, offset, data, cache)) { allocator_->Deallocate(data); return nullptr; } need_release = true; + try { + std::scoped_lock lock(cached_direct_reads_mutex_); + cached_direct_reads_.emplace(data); + } catch (...) { + need_release = false; + allocator_->Deallocate(data); + throw; + } return data; } } @@ -159,14 +180,102 @@ class BasicIO { MultiRead(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { static_assert(has_MultiReadImpl::value); if constexpr (not InMemory) { - if (cache_ != nullptr) { + const auto cache = GetReadCacheSnapshot(); + if (cache.cache != nullptr) { + if (count == 0) { + return true; + } + if (sizes == nullptr or offsets == nullptr) { + return false; + } + bool has_data = false; for (uint64_t i = 0; i < count; ++i) { - if (not ReadCached(sizes[i], offsets[i], datas)) { + if (not IsValidRange(sizes[i], offsets[i])) { return false; } - datas += sizes[i]; + has_data = has_data or sizes[i] > 0; } - return true; + if (not has_data) { + return true; + } + if (datas == nullptr) { + return false; + } + struct PageCopy { + uint64_t page_id; + uint64_t page_offset; + uint64_t output_offset; + uint64_t size; + }; + using PageCopyList = std::vector>; + PageIdList page_ids(allocator_); + PageCopyList page_copies(allocator_); + UnorderedSet seen_page_ids(allocator_); + uint64_t output_offset = 0; + for (uint64_t i = 0; i < count; ++i) { + if (sizes[i] > UINT64_MAX - output_offset) { + return false; + } + if (sizes[i] == 0) { + continue; + } + uint64_t copied = 0; + while (copied < sizes[i]) { + const uint64_t current_offset = offsets[i] + copied; + const uint64_t page_id = current_offset / Page::DEFAULT_PAGE_SIZE; + const uint64_t page_offset = current_offset % Page::DEFAULT_PAGE_SIZE; + const uint64_t copy_size = + std::min(sizes[i] - copied, Page::DEFAULT_PAGE_SIZE - page_offset); + if (seen_page_ids.emplace(page_id).second) { + page_ids.emplace_back(page_id); + } + page_copies.emplace_back( + PageCopy{page_id, page_offset, output_offset + copied, copy_size}); + copied += copy_size; + } + output_offset += sizes[i]; + } + + std::sort(page_ids.begin(), page_ids.end()); + std::sort( + page_copies.begin(), page_copies.end(), [](const auto& lhs, const auto& rhs) { + return lhs.page_id < rhs.page_id; + }); + // Each miss temporarily needs both a Page and contiguous MultiRead storage. + // Keep that temporary working set near 8 MiB while preserving batched backend IO. + constexpr uint64_t MAX_BATCH_READ_TEMP_BYTES = 8ULL * 1024ULL * 1024ULL; + const uint64_t max_batch_pages = std::max( + 1, MAX_BATCH_READ_TEMP_BYTES / (2 * Page::DEFAULT_PAGE_SIZE)); + UnorderedMap pages(allocator_); + uint64_t copy_index = 0; + for (uint64_t batch_begin = 0; batch_begin < page_ids.size(); + batch_begin += max_batch_pages) { + const uint64_t batch_end = + std::min(page_ids.size(), batch_begin + max_batch_pages); + PageIdList batch_page_ids(allocator_); + batch_page_ids.reserve(batch_end - batch_begin); + for (uint64_t i = batch_begin; i < batch_end; ++i) { + batch_page_ids.emplace_back(page_ids[i]); + } + pages.clear(); + if (not LoadCachedPages(batch_page_ids, pages, true, cache)) { + return false; + } + const uint64_t last_page_id = page_ids[batch_end - 1]; + while (copy_index < page_copies.size() and + page_copies[copy_index].page_id <= last_page_id) { + const auto& copy = page_copies[copy_index]; + const auto iter = pages.find(copy.page_id); + if (iter == pages.end() or iter->second == nullptr) { + return false; + } + std::memcpy(datas + copy.output_offset, + iter->second->Data() + copy.page_offset, + copy.size); + ++copy_index; + } + } + return copy_index == page_copies.size(); } } return cast().MultiReadImpl(datas, sizes, offsets, count); @@ -258,8 +367,9 @@ class BasicIO { inline void Release(const uint8_t* data) const { - if constexpr (not InMemory) { - if (cache_ != nullptr) { + { + std::scoped_lock lock(cached_direct_reads_mutex_); + if (cached_direct_reads_.erase(data) != 0) { allocator_->Deallocate(const_cast(data)); return; } @@ -281,7 +391,8 @@ class BasicIO { inline void InitIO(const IOParamPtr& io_param) { if constexpr (not InMemory) { - if (cache_ == nullptr) { + const auto cache = GetReadCacheSnapshot(); + if (cache.cache == nullptr) { EnableReadCache(io_param); } else if (io_param != nullptr and io_param->enable_read_cache_) { EnableReadCache(io_param); @@ -346,6 +457,7 @@ class BasicIO { void EnableReadCache(const IOParamPtr& io_param) { if constexpr (not InMemory) { + std::scoped_lock lock(cache_mutex_); if (io_param == nullptr or not io_param->enable_read_cache_) { cache_.reset(); cache_page_id_base_ = 0; @@ -450,8 +562,17 @@ class BasicIO { return offset <= size_ and size <= size_ - offset; } + ReadCacheSnapshot + GetReadCacheSnapshot() const { + std::scoped_lock lock(cache_mutex_); + return {cache_, cache_page_id_base_}; + } + bool - ReadCached(uint64_t size, uint64_t offset, uint8_t* data) const { + ReadCached(uint64_t size, + uint64_t offset, + uint8_t* data, + const ReadCacheSnapshot& cache) const { if (not IsValidRange(size, offset)) { return false; } @@ -461,7 +582,7 @@ class BasicIO { uint64_t page_id = current_offset / Page::DEFAULT_PAGE_SIZE; uint64_t page_offset = current_offset % Page::DEFAULT_PAGE_SIZE; uint64_t copy_size = std::min(size - copied, Page::DEFAULT_PAGE_SIZE - page_offset); - auto page = GetOrLoadPage(page_id); + auto page = GetOrLoadPage(page_id, cache); if (page == nullptr) { return false; } @@ -472,52 +593,249 @@ class BasicIO { } PagePtr - GetOrLoadPage(uint64_t page_id) const { + GetOrLoadPage(uint64_t page_id, const ReadCacheSnapshot& cache) const { if (page_id > UINT64_MAX / Page::DEFAULT_PAGE_SIZE) { return nullptr; } - uint64_t offset = page_id * Page::DEFAULT_PAGE_SIZE; + const uint64_t offset = page_id * Page::DEFAULT_PAGE_SIZE; if (offset >= size_) { return nullptr; } - std::scoped_lock lock(cache_mutex_); - if (page_id > UINT64_MAX - cache_page_id_base_) { + if (cache.cache == nullptr or page_id > UINT64_MAX - cache.page_id_base) { return nullptr; } - uint64_t cache_page_id = cache_page_id_base_ + page_id; - auto page = cache_->Get(cache_page_id); - if (page != nullptr) { - return page; - } - auto new_page = std::make_shared(allocator_); - if (new_page->Data() == nullptr) { + const uint64_t cache_page_id = cache.page_id_base + page_id; + while (true) { + auto result = cache.cache->Acquire(cache_page_id); + if (result.page != nullptr) { + return result.page; + } + if (not result.should_load) { + auto page = cache.cache->Wait(result.handle); + if (page != nullptr) { + return page; + } + if (cache.cache->IsStale(result.handle)) { + continue; + } + return nullptr; + } + + PagePtr page = nullptr; + bool success = false; + try { + page = std::make_shared(allocator_); + success = page->Data() != nullptr; + if (success) { + const uint64_t read_size = std::min(Page::DEFAULT_PAGE_SIZE, size_ - offset); + success = cast().ReadImpl(read_size, offset, page->Data()); + } + } catch (...) { + cache.cache->Complete(cache_page_id, result.handle, nullptr, false); + throw; + } + page = cache.cache->Complete( + cache_page_id, result.handle, success ? std::move(page) : nullptr, success); + if (page != nullptr) { + return page; + } + if (cache.cache->IsStale(result.handle)) { + continue; + } return nullptr; } - uint64_t read_size = std::min(Page::DEFAULT_PAGE_SIZE, size_ - offset); - if (not cast().ReadImpl(read_size, offset, new_page->Data())) { - return nullptr; + } + + bool + LoadCachedPages(const PageIdList& page_ids, + UnorderedMap& pages, + bool batch_read, + const ReadCacheSnapshot& cache) const { + PageIdList ordered_page_ids(page_ids.begin(), page_ids.end(), allocator_); + std::sort(ordered_page_ids.begin(), ordered_page_ids.end()); + + while (true) { + pages.clear(); + bool stale = false; + const bool loaded = + LoadCachedPagesOnce(ordered_page_ids, pages, batch_read, stale, cache); + if (loaded and not stale) { + return true; + } + if (not stale) { + return false; + } } - return cache_->Insert(cache_page_id, std::move(new_page)); + } + + bool + LoadCachedPagesOnce(const PageIdList& page_ids, + UnorderedMap& pages, + bool batch_read, + bool& stale, + const ReadCacheSnapshot& cache) const { + stale = false; + if (cache.cache == nullptr or + std::any_of(page_ids.begin(), page_ids.end(), [&cache](uint64_t page_id) { + return page_id > UINT64_MAX - cache.page_id_base; + })) { + return false; + } + using PageLoadList = + std::vector, + AllocatorWrapper>>; + PageLoadList owned_loads(allocator_); + pages.reserve(page_ids.size()); + owned_loads.reserve(page_ids.size()); + const auto abandon_owned_loads = [&cache, &owned_loads]() { + for (const auto& [page_id, handle] : owned_loads) { + cache.cache->Complete(cache.page_id_base + page_id, handle, nullptr, false); + } + owned_loads.clear(); + }; + try { + for (const uint64_t page_id : page_ids) { + auto result = cache.cache->Acquire(cache.page_id_base + page_id); + if (result.page != nullptr) { + pages.emplace(page_id, std::move(result.page)); + } else if (result.should_load) { + owned_loads.emplace_back(page_id, std::move(result.handle)); + } else { + auto page = cache.cache->Wait(result.handle); + stale = stale or cache.cache->IsStale(result.handle); + if (page == nullptr or stale) { + abandon_owned_loads(); + return false; + } + pages.emplace(page_id, std::move(page)); + } + } + } catch (...) { + abandon_owned_loads(); + throw; + } + + std::vector> loaded_pages(allocator_); + std::vector> read_sizes(allocator_); + std::vector> read_offsets(allocator_); + std::vector> read_data(allocator_); + bool success = true; + try { + uint64_t total_size = 0; + loaded_pages.reserve(owned_loads.size()); + read_sizes.reserve(owned_loads.size()); + read_offsets.reserve(owned_loads.size()); + for (const auto& [page_id, _] : owned_loads) { + if (page_id > UINT64_MAX / Page::DEFAULT_PAGE_SIZE) { + success = false; + break; + } + const uint64_t offset = page_id * Page::DEFAULT_PAGE_SIZE; + if (offset >= size_) { + success = false; + break; + } + const uint64_t read_size = std::min(Page::DEFAULT_PAGE_SIZE, size_ - offset); + if (read_size > UINT64_MAX - total_size) { + success = false; + break; + } + auto page = std::make_shared(allocator_); + if (page->Data() == nullptr) { + success = false; + break; + } + loaded_pages.emplace_back(std::move(page)); + read_sizes.emplace_back(read_size); + read_offsets.emplace_back(offset); + total_size += read_size; + } + if (success and not owned_loads.empty()) { + if (batch_read) { + read_data.resize(total_size); + success = cast().MultiReadImpl(read_data.data(), + read_sizes.data(), + read_offsets.data(), + owned_loads.size()); + uint64_t copied = 0; + for (uint64_t i = 0; success and i < loaded_pages.size(); ++i) { + std::memcpy( + loaded_pages[i]->Data(), read_data.data() + copied, read_sizes[i]); + copied += read_sizes[i]; + } + } else { + for (uint64_t i = 0; success and i < loaded_pages.size(); ++i) { + success = cast().ReadImpl( + read_sizes[i], read_offsets[i], loaded_pages[i]->Data()); + } + } + } + } catch (...) { + for (const auto& [page_id, handle] : owned_loads) { + cache.cache->Complete(cache.page_id_base + page_id, handle, nullptr, false); + } + throw; + } + + for (uint64_t i = 0; i < owned_loads.size(); ++i) { + const auto& [page_id, handle] = owned_loads[i]; + PagePtr loaded_page = nullptr; + if (success) { + loaded_page = std::move(loaded_pages[i]); + } + try { + auto page = cache.cache->Complete( + cache.page_id_base + page_id, handle, std::move(loaded_page), success); + stale = stale or cache.cache->IsStale(handle); + if (page != nullptr) { + pages.emplace(page_id, std::move(page)); + } + } catch (...) { + for (uint64_t j = i + 1; j < owned_loads.size(); ++j) { + const auto& [remaining_page_id, remaining_handle] = owned_loads[j]; + cache.cache->Complete( + cache.page_id_base + remaining_page_id, remaining_handle, nullptr, false); + } + throw; + } + } + if (not success) { + return false; + } + + return std::all_of(page_ids.begin(), page_ids.end(), [&pages](uint64_t page_id) { + const auto iter = pages.find(page_id); + return iter != pages.end() and iter->second != nullptr; + }); } void InvalidateCacheRange(uint64_t size, uint64_t offset) { - if (cache_ == nullptr or size == 0) { + if (size == 0) { + return; + } + std::shared_ptr cache; + uint64_t cache_page_id_base; + { + std::scoped_lock lock(cache_mutex_); + cache = cache_; + cache_page_id_base = cache_page_id_base_; + } + if (cache == nullptr) { return; } - std::scoped_lock lock(cache_mutex_); if (offset > UINT64_MAX - (size - 1)) { - cache_->Clear(); + cache->Clear(); return; } uint64_t first_page = offset / Page::DEFAULT_PAGE_SIZE; uint64_t last_page = (offset + size - 1) / Page::DEFAULT_PAGE_SIZE; - if (last_page > UINT64_MAX - cache_page_id_base_) { - cache_->Clear(); + if (last_page > UINT64_MAX - cache_page_id_base) { + cache->Clear(); return; } for (uint64_t page_id = first_page;; ++page_id) { - cache_->Remove(cache_page_id_base_ + page_id); + cache->Remove(cache_page_id_base + page_id); if (page_id == last_page) { break; } @@ -526,9 +844,13 @@ class BasicIO { void ClearCache() { - if (cache_ != nullptr) { + std::shared_ptr cache; + { std::scoped_lock lock(cache_mutex_); - cache_->Clear(); + cache = cache_; + } + if (cache != nullptr) { + cache->Clear(); } } @@ -540,6 +862,8 @@ class BasicIO { mutable std::mutex cache_mutex_; mutable std::shared_ptr cache_; uint64_t cache_page_id_base_{0}; + mutable std::mutex cached_direct_reads_mutex_; + mutable UnorderedSet cached_direct_reads_; bool has_deserialized_{false}; private: diff --git a/src/io/read_cache/page_cache.cpp b/src/io/read_cache/page_cache.cpp index 7ae0dc69b8..c8e0a594c8 100644 --- a/src/io/read_cache/page_cache.cpp +++ b/src/io/read_cache/page_cache.cpp @@ -15,9 +15,19 @@ #include "io/read_cache/page_cache.h" #include +#include +#include +#include namespace vsag { +struct PageCache::LoadingPage { + std::condition_variable cv; + PagePtr page; + bool done{false}; + bool stale{false}; +}; + PageCache::PageCache(uint64_t max_pages) : max_pages_(max_pages) { } @@ -32,9 +42,77 @@ PageCache::Get(uint64_t page_id) { return it->second; } +PageCache::LoadResult +PageCache::Acquire(uint64_t page_id) { + std::scoped_lock lock(mutex_); + if (const auto iter = pages_.find(page_id); iter != pages_.end()) { + OnAccess(page_id); + LoadResult result; + result.page = iter->second; + return result; + } + if (const auto iter = loading_pages_.find(page_id); iter != loading_pages_.end()) { + LoadResult result; + result.handle.state_ = iter->second; + return result; + } + auto state = std::make_shared(); + loading_pages_.emplace(page_id, state); + LoadResult result; + result.handle.state_ = std::move(state); + result.should_load = true; + return result; +} + +PagePtr +PageCache::Wait(const LoadHandle& handle) { + std::unique_lock lock(mutex_); + handle.state_->cv.wait(lock, [&handle] { return handle.state_->done; }); + return handle.state_->page; +} + +bool +PageCache::IsStale(const LoadHandle& handle) const { + std::scoped_lock lock(mutex_); + return handle.state_->stale; +} + +PagePtr +PageCache::Complete(uint64_t page_id, const LoadHandle& handle, PagePtr page, bool success) { + PagePtr result; + std::exception_ptr error; + { + std::scoped_lock lock(mutex_); + const auto state = handle.state_; + try { + if (success and not state->stale) { + result = InsertLocked(page_id, std::move(page)); + } + } catch (...) { + error = std::current_exception(); + } + state->page = result; + state->done = true; + if (const auto iter = loading_pages_.find(page_id); + iter != loading_pages_.end() and iter->second == state) { + loading_pages_.erase(iter); + } + } + handle.state_->cv.notify_all(); + if (error != nullptr) { + std::rethrow_exception(error); + } + return result; +} + PagePtr PageCache::Insert(uint64_t page_id, PagePtr page) { std::scoped_lock lock(mutex_); + return InsertLocked(page_id, std::move(page)); +} + +PagePtr +PageCache::InsertLocked(uint64_t page_id, PagePtr page) { auto existing = pages_.find(page_id); if (existing != pages_.end()) { OnAccess(page_id); @@ -64,15 +142,30 @@ PageCache::Remove(uint64_t page_id) { OnRemove(page_id); pages_.erase(it); } + if (const auto loading = loading_pages_.find(page_id); loading != loading_pages_.end()) { + loading->second->stale = true; + } } void PageCache::Clear() { - std::scoped_lock lock(mutex_); - for (const auto& page_pair : pages_) { - OnRemove(page_pair.first); + std::vector> states_to_notify; + { + std::scoped_lock lock(mutex_); + for (const auto& page_pair : pages_) { + OnRemove(page_pair.first); + } + pages_.clear(); + for (const auto& [page_id, state] : loading_pages_) { + state->stale = true; + state->done = true; + states_to_notify.push_back(state); + } + loading_pages_.clear(); + } + for (const auto& state : states_to_notify) { + state->cv.notify_all(); } - pages_.clear(); } uint64_t diff --git a/src/io/read_cache/page_cache.h b/src/io/read_cache/page_cache.h index 71fce90315..eca598ee10 100644 --- a/src/io/read_cache/page_cache.h +++ b/src/io/read_cache/page_cache.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include @@ -31,7 +32,23 @@ namespace vsag { * from the cache concurrently. */ class PageCache { +private: + struct LoadingPage; + public: + class LoadHandle { + friend class PageCache; + + private: + std::shared_ptr state_; + }; + + struct LoadResult { + PagePtr page; + LoadHandle handle; + bool should_load{false}; + }; + explicit PageCache(uint64_t max_pages); virtual ~PageCache() = default; @@ -45,6 +62,32 @@ class PageCache { virtual PagePtr Get(uint64_t page_id); + /** + * @brief Get a cached page or register a single owner to load a cache miss. + * + * Other callers for the same miss receive a handle that can be waited on. + */ + LoadResult + Acquire(uint64_t page_id); + + /** + * @brief Wait for an in-flight page load to complete. + */ + PagePtr + Wait(const LoadHandle& handle); + + /** + * @brief Returns whether an in-flight load was invalidated before completion. + */ + bool + IsStale(const LoadHandle& handle) const; + + /** + * @brief Publish a page load result and wake callers waiting on the page. + */ + PagePtr + Complete(uint64_t page_id, const LoadHandle& handle, PagePtr page, bool success); + /** * @brief Insert a page, evicting victims first if the cache is full. * @@ -90,9 +133,13 @@ class PageCache { virtual uint64_t PickVictim() = 0; + PagePtr + InsertLocked(uint64_t page_id, PagePtr page); + protected: mutable std::mutex mutex_; std::unordered_map pages_; + std::unordered_map> loading_pages_; uint64_t max_pages_{0}; }; diff --git a/src/io/read_cache/page_cache_test.cpp b/src/io/read_cache/page_cache_test.cpp index dc229b9825..ea3115ffc5 100644 --- a/src/io/read_cache/page_cache_test.cpp +++ b/src/io/read_cache/page_cache_test.cpp @@ -15,7 +15,10 @@ #include "io/read_cache/page_cache.h" #include +#include #include +#include +#include #include "impl/allocator/safe_allocator.h" #include "unittest.h" @@ -50,6 +53,9 @@ class FifoPageCache : public PageCache { uint64_t PickVictim() override { + if (throw_on_pick_) { + throw std::runtime_error("pick failed"); + } if (order_.empty()) { return UINT64_MAX; } @@ -59,6 +65,8 @@ class FifoPageCache : public PageCache { public: uint64_t last_access_{UINT64_MAX}; + bool throw_on_pick_{false}; + private: std::deque order_; }; @@ -112,3 +120,29 @@ TEST_CASE("PageCache Eviction Test", "[PageCache][ut]") { REQUIRE(cache.Get(2) != nullptr); REQUIRE(cache.Get(3) != nullptr); } + +TEST_CASE("PageCache Complete wakes waiters when insertion throws", "[PageCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + FifoPageCache cache(1); + cache.Insert(1, MakePage(allocator.get(), 1)); + + auto owner = cache.Acquire(2); + REQUIRE(owner.should_load); + auto waiter = cache.Acquire(2); + REQUIRE_FALSE(waiter.should_load); + + auto waited_page = std::async(std::launch::async, + [&cache, handle = waiter.handle] { return cache.Wait(handle); }); + + cache.throw_on_pick_ = true; + REQUIRE_THROWS_AS(cache.Complete(2, owner.handle, MakePage(allocator.get(), 2), true), + std::runtime_error); + REQUIRE(waited_page.wait_for(std::chrono::seconds(1)) == std::future_status::ready); + REQUIRE(waited_page.get() == nullptr); + + cache.throw_on_pick_ = false; + auto retry = cache.Acquire(2); + REQUIRE(retry.should_load); + auto retried_page = MakePage(allocator.get(), 2); + REQUIRE(cache.Complete(2, retry.handle, retried_page, true) == retried_page); +} diff --git a/src/io/read_cache/read_cache_test.cpp b/src/io/read_cache/read_cache_test.cpp index 36e5faea68..f62e86bad9 100644 --- a/src/io/read_cache/read_cache_test.cpp +++ b/src/io/read_cache/read_cache_test.cpp @@ -12,7 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include +#include #include +#include +#include #include #include "impl/allocator/safe_allocator.h" @@ -63,6 +69,143 @@ class TestReader : public Reader { const std::vector& data_; }; +class CountingIO : public BasicIO { +public: + static constexpr bool InMemory = false; + static constexpr bool SkipDeserialize = false; + + explicit CountingIO(Allocator* allocator) : BasicIO(allocator) { + } + + void + WriteImpl(const uint8_t* data, uint64_t size, uint64_t offset) { + std::scoped_lock lock(multi_read_mutex_); + if (data_.size() < offset + size) { + data_.resize(offset + size); + size_ = offset + size; + } + std::memcpy(data_.data() + offset, data, size); + } + + bool + ReadImpl(uint64_t size, uint64_t offset, uint8_t* data) const { + ++read_count; + { + std::unique_lock lock(read_mutex_); + if (block_read_) { + read_blocked_ = true; + read_blocked_cv_.notify_all(); + read_resume_cv_.wait(lock, [this] { return not block_read_; }); + } + } + std::scoped_lock lock(multi_read_mutex_); + std::memcpy(data, data_.data() + offset, size); + return true; + } + + const uint8_t* + DirectReadImpl(uint64_t size, uint64_t offset, bool& need_release) const { + std::scoped_lock lock(multi_read_mutex_); + need_release = true; + if (offset > data_.size() or size > data_.size() - offset) { + return nullptr; + } + return data_.data() + offset; + } + + void + ReleaseImpl(const uint8_t*) const { + ++release_count; + } + + bool + MultiReadImpl(uint8_t* datas, uint64_t* sizes, uint64_t* offsets, uint64_t count) const { + ++multi_read_count; + { + std::unique_lock lock(multi_read_mutex_); + if (block_multi_read_) { + multi_read_blocked_ = true; + multi_read_blocked_cv_.notify_all(); + multi_read_resume_cv_.wait(lock, [this] { return not block_multi_read_; }); + } + } + if (fail_next_multi_read.exchange(false)) { + return false; + } + std::scoped_lock lock(multi_read_mutex_); + for (uint64_t i = 0; i < count; ++i) { + std::memcpy(datas, data_.data() + offsets[i], sizes[i]); + datas += sizes[i]; + } + return true; + } + + void + BlockMultiRead() { + std::scoped_lock lock(multi_read_mutex_); + block_multi_read_ = true; + multi_read_blocked_ = false; + } + + bool + WaitForMultiReadBlock() const { + std::unique_lock lock(multi_read_mutex_); + return multi_read_blocked_cv_.wait_for( + lock, std::chrono::seconds(5), [this] { return multi_read_blocked_; }); + } + + void + UnblockMultiRead() { + { + std::scoped_lock lock(multi_read_mutex_); + block_multi_read_ = false; + } + multi_read_resume_cv_.notify_all(); + } + + void + BlockRead() { + std::scoped_lock lock(read_mutex_); + block_read_ = true; + read_blocked_ = false; + } + + bool + WaitForReadBlock() const { + std::unique_lock lock(read_mutex_); + return read_blocked_cv_.wait_for( + lock, std::chrono::seconds(5), [this] { return read_blocked_; }); + } + + void + UnblockRead() { + { + std::scoped_lock lock(read_mutex_); + block_read_ = false; + } + read_resume_cv_.notify_all(); + } + + mutable std::atomic read_count{0}; + mutable std::atomic multi_read_count{0}; + mutable std::atomic fail_next_multi_read{false}; + + mutable std::atomic release_count{0}; + +private: + std::vector data_; + mutable std::mutex multi_read_mutex_; + mutable std::condition_variable multi_read_blocked_cv_; + mutable std::condition_variable multi_read_resume_cv_; + mutable bool block_multi_read_{false}; + mutable bool multi_read_blocked_{false}; + mutable std::mutex read_mutex_; + mutable std::condition_variable read_blocked_cv_; + mutable std::condition_variable read_resume_cv_; + mutable bool block_read_{false}; + mutable bool read_blocked_{false}; +}; + } // namespace TEST_CASE("BasicIO cache component basic test", "[ReadCache][ut]") { @@ -122,6 +265,257 @@ TEST_CASE("BasicIO cache component direct and multi read", "[ReadCache][ut]") { REQUIRE(std::memcmp(result.data(), data.data(), result.size()) == 0); } +TEST_CASE("BasicIO cache releases direct-read buffers by their allocation source", + "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(128, 0x5A); + io.Write(source.data(), source.size(), 0); + + bool need_release = false; + const auto* cached = io.Read(source.size(), 0, need_release); + REQUIRE(cached != nullptr); + REQUIRE(need_release); + REQUIRE(std::memcmp(cached, source.data(), source.size()) == 0); + + io.SetReadCache(nullptr); + io.Release(cached); + REQUIRE(io.release_count == 0); + + const auto* direct = io.Read(source.size(), 0, need_release); + REQUIRE(direct != nullptr); + REQUIRE(need_release); + io.Release(direct); + REQUIRE(io.release_count == 1); +} + +TEST_CASE("BasicIO cache validates multi-read pointers", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0x5A); + io.Write(source.data(), source.size(), 0); + + uint64_t zero = 0; + uint64_t one = 1; + REQUIRE(io.MultiRead(nullptr, nullptr, nullptr, 0)); + REQUIRE_FALSE(io.MultiRead(nullptr, nullptr, &zero, 1)); + REQUIRE_FALSE(io.MultiRead(nullptr, &zero, nullptr, 1)); + REQUIRE(io.MultiRead(nullptr, &zero, &zero, 1)); + REQUIRE_FALSE(io.MultiRead(nullptr, &one, &zero, 1)); +} + +TEST_CASE("BasicIO cache batches missing pages", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam(4)); + + std::vector source(Page::DEFAULT_PAGE_SIZE * 3); + for (uint64_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i); + } + io.Write(source.data(), source.size(), 0); + + uint64_t sizes[] = {32, 64, 48}; + uint64_t offsets[] = {Page::DEFAULT_PAGE_SIZE - 16, + Page::DEFAULT_PAGE_SIZE + 20, + Page::DEFAULT_PAGE_SIZE * 2 + 40}; + std::vector result(144); + REQUIRE(io.MultiRead(result.data(), sizes, offsets, 3)); + REQUIRE(io.multi_read_count == 1); + REQUIRE(io.read_count == 0); + REQUIRE(std::memcmp(result.data(), source.data() + offsets[0], sizes[0]) == 0); + REQUIRE(std::memcmp(result.data() + sizes[0], source.data() + offsets[1], sizes[1]) == 0); + REQUIRE(std::memcmp( + result.data() + sizes[0] + sizes[1], source.data() + offsets[2], sizes[2]) == 0); +} + +TEST_CASE("BasicIO cache bounds sparse multi-read batches", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam(4)); + + constexpr uint64_t temp_budget = 8ULL * 1024ULL * 1024ULL; + const uint64_t page_count = temp_budget / (2 * Page::DEFAULT_PAGE_SIZE) + 1; + std::vector source(page_count * Page::DEFAULT_PAGE_SIZE); + std::vector sizes(page_count, 1); + std::vector offsets(page_count); + std::vector result(page_count); + for (uint64_t i = 0; i < page_count; ++i) { + offsets[i] = i * Page::DEFAULT_PAGE_SIZE; + source[offsets[i]] = static_cast(i); + } + io.Write(source.data(), source.size(), 0); + + REQUIRE(io.MultiRead(result.data(), sizes.data(), offsets.data(), page_count)); + REQUIRE(io.multi_read_count == 2); + for (uint64_t i = 0; i < page_count; ++i) { + REQUIRE(result[i] == source[offsets[i]]); + } +} + +TEST_CASE("BasicIO cache uses single-read backend for direct reads", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0xA5); + std::vector result(source.size()); + io.Write(source.data(), source.size(), 0); + + REQUIRE(io.Read(result.size(), 0, result.data())); + REQUIRE(io.read_count == 1); + REQUIRE(io.multi_read_count == 0); + REQUIRE(result == source); +} + +TEST_CASE("BasicIO cache shares in-flight page loads and retries failures", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0xA5); + io.Write(source.data(), source.size(), 0); + io.BlockMultiRead(); + std::mutex start_mutex; + std::condition_variable start_cv; + uint64_t readers_ready = 0; + bool start = false; + std::vector first(source.size()); + std::vector second(source.size()); + std::atomic first_succeeded{false}; + std::atomic second_succeeded{false}; + auto read_page = [&](std::vector& result, std::atomic& succeeded) { + { + std::unique_lock lock(start_mutex); + ++readers_ready; + start_cv.notify_all(); + if (not start_cv.wait_for(lock, std::chrono::seconds(5), [&] { return start; })) { + return; + } + } + uint64_t size = result.size(); + uint64_t offset = 0; + succeeded = io.MultiRead(result.data(), &size, &offset, 1); + }; + std::thread first_reader(read_page, std::ref(first), std::ref(first_succeeded)); + std::thread second_reader(read_page, std::ref(second), std::ref(second_succeeded)); + bool readers_ready_in_time = false; + { + std::unique_lock lock(start_mutex); + readers_ready_in_time = + start_cv.wait_for(lock, std::chrono::seconds(5), [&] { return readers_ready == 2; }); + start = true; + } + start_cv.notify_all(); + const bool multi_read_blocked = io.WaitForMultiReadBlock(); + io.UnblockMultiRead(); + first_reader.join(); + second_reader.join(); + + REQUIRE(readers_ready_in_time); + REQUIRE(multi_read_blocked); + REQUIRE(first_succeeded); + REQUIRE(second_succeeded); + REQUIRE(io.multi_read_count == 1); + REQUIRE(first == source); + REQUIRE(second == source); + + CountingIO failing_io(allocator.get()); + failing_io.EnableReadCache(MakeReadCacheParam()); + failing_io.Write(source.data(), source.size(), 0); + failing_io.fail_next_multi_read = true; + uint64_t size = source.size(); + uint64_t offset = 0; + std::vector result(source.size()); + REQUIRE_FALSE(failing_io.MultiRead(result.data(), &size, &offset, 1)); + REQUIRE(failing_io.MultiRead(result.data(), &size, &offset, 1)); + REQUIRE(failing_io.multi_read_count == 2); + REQUIRE(result == source); +} + +TEST_CASE("BasicIO cache keeps an in-flight multi-read snapshot", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0xA5); + std::vector result(source.size()); + io.Write(source.data(), source.size(), 0); + io.BlockMultiRead(); + + std::atomic read_succeeded{false}; + std::thread reader([&] { + uint64_t size = result.size(); + uint64_t offset = 0; + read_succeeded = io.MultiRead(result.data(), &size, &offset, 1); + }); + const bool multi_read_blocked = io.WaitForMultiReadBlock(); + io.SetReadCache(nullptr); + io.UnblockMultiRead(); + reader.join(); + + REQUIRE(multi_read_blocked); + REQUIRE(read_succeeded); + REQUIRE(result == source); + REQUIRE(io.multi_read_count == 1); +} + +TEST_CASE("BasicIO cache reloads invalidated in-flight pages", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0xA5); + std::vector result(source.size()); + io.Write(source.data(), source.size(), 0); + io.BlockMultiRead(); + + std::atomic read_succeeded{false}; + std::thread reader([&] { + uint64_t size = result.size(); + uint64_t offset = 0; + read_succeeded = io.MultiRead(result.data(), &size, &offset, 1); + }); + const bool multi_read_blocked = io.WaitForMultiReadBlock(); + source[0] = 0x5A; + io.Write(source.data(), 1, 0); + io.UnblockMultiRead(); + reader.join(); + + REQUIRE(multi_read_blocked); + REQUIRE(read_succeeded); + REQUIRE(result == source); + REQUIRE(io.multi_read_count == 2); +} + +TEST_CASE("BasicIO cache reloads invalidated direct reads", "[ReadCache][ut]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + CountingIO io(allocator.get()); + io.EnableReadCache(MakeReadCacheParam()); + + std::vector source(Page::DEFAULT_PAGE_SIZE, 0xA5); + std::vector result(source.size()); + io.Write(source.data(), source.size(), 0); + io.BlockRead(); + + std::atomic read_succeeded{false}; + std::thread reader([&] { read_succeeded = io.Read(result.size(), 0, result.data()); }); + const bool read_blocked = io.WaitForReadBlock(); + source[0] = 0x5A; + io.Write(source.data(), 1, 0); + io.UnblockRead(); + reader.join(); + + REQUIRE(read_blocked); + REQUIRE(read_succeeded); + REQUIRE(result == source); + REQUIRE(io.read_count == 2); +} + TEST_CASE("BasicIO cache component initializes ReaderIO", "[ReadCache][ut]") { std::vector data(1024); for (uint64_t i = 0; i < data.size(); ++i) { diff --git a/tests/test_ivf.cpp b/tests/test_ivf.cpp index 8740fd257b..62420bc118 100644 --- a/tests/test_ivf.cpp +++ b/tests/test_ivf.cpp @@ -691,7 +691,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, expected_search.value()->GetDistances()[0]) < 2e-6F); REQUIRE(precise_reader->ReadBytes() > 0); if (enable_read_cache) { - REQUIRE(precise_reader->ReadCalls() > 0); + REQUIRE(precise_reader->ReadCalls() + precise_reader->MultiReadCalls() > 0); precise_reader->ResetCounters(); auto cached_search = loaded.value()->KnnSearch(query, 1, search_param); REQUIRE(cached_search.has_value()); @@ -723,7 +723,7 @@ TEST_CASE_PERSISTENT_FIXTURE(IVFTestIndex, auto first_distances = fresh_loaded.value()->CalcDistancesById(query_vector, ids, batch_count); REQUIRE(first_distances.has_value()); - REQUIRE(fresh_reader->ReadCalls() > 0); + REQUIRE(fresh_reader->ReadCalls() + fresh_reader->MultiReadCalls() > 0); REQUIRE(fresh_reader->ReadBytes() > 0); for (int64_t i = 0; i < batch_count; ++i) { REQUIRE(std::abs(first_distances.value()->GetDistances()[i] -